src/shuf.c

   1 /* Shuffle lines of text.
   2
   3    Copyright (C) 2006-2012 Free Software Foundation, Inc.
   4
   5    This program is free software: you can redistribute it and/or modify
   6    it under the terms of the GNU General Public License as published by
   7    the Free Software Foundation, either version 3 of the License, or
   8    (at your option) any later version.
   9
  10    This program is distributed in the hope that it will be useful,
  11    but WITHOUT ANY WARRANTY; without even the implied warranty of
  12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  13    GNU General Public License for more details.
  14
  15    You should have received a copy of the GNU General Public License
  16    along with this program.  If not, see <http://www.gnu.org/licenses/>.
  17
  18    Written by Paul Eggert.  */
  19
  20 #include <config.h>
  21
  22 #include <sys/types.h>
  23 #include "system.h"
  24
  25 #include "error.h"
  26 #include "fadvise.h"
  27 #include "getopt.h"
  28 #include "quote.h"
  29 #include "quotearg.h"
  30 #include "randint.h"
  31 #include "randperm.h"
  32 #include "read-file.h"
  33 #include "stdio--.h"
  34 #include "xstrtol.h"
  35
  36 /* The official name of this program (e.g., no 'g' prefix).  */
  37 #define PROGRAM_NAME "shuf"
  38
  39 #define AUTHORS proper_name ("Paul Eggert")
  40
  41 void
  42 usage (int status)
  43 {
  44   if (status != EXIT_SUCCESS)
  45     emit_try_help ();
  46   else
  47     {
  48       printf (_("\
  49 Usage: %s [OPTION]... [FILE]\n\
  50   or:  %s -e [OPTION]... [ARG]...\n\
  51   or:  %s -i LO-HI [OPTION]...\n\
  52 "),
  53               program_name, program_name, program_name);
  54       fputs (_("\
  55 Write a random permutation of the input lines to standard output.\n\
  56 \n\
  57 "), stdout);
  58       fputs (_("\
  59 Mandatory arguments to long options are mandatory for short options too.\n\
  60 "), stdout);
  61       fputs (_("\
  62   -e, --echo                treat each ARG as an input line\n\
  63   -i, --input-range=LO-HI   treat each number LO through HI as an input line\n\
  64   -n, --head-count=COUNT    output at most COUNT lines\n\
  65   -o, --output=FILE         write result to FILE instead of standard output\n\
  66       --random-source=FILE  get random bytes from FILE\n\
  67   -z, --zero-terminated     end lines with 0 byte, not newline\n\
  68 "), stdout);
  69       fputs (HELP_OPTION_DESCRIPTION, stdout);
  70       fputs (VERSION_OPTION_DESCRIPTION, stdout);
  71       fputs (_("\
  72 \n\
  73 With no FILE, or when FILE is -, read standard input.\n\
  74 "), stdout);
  75       emit_ancillary_info ();
  76     }
  77
  78   exit (status);
  79 }
  80
  81 /* For long options that have no equivalent short option, use a
  82    non-character as a pseudo short option, starting with CHAR_MAX + 1.  */
  83 enum
  84 {
  85   RANDOM_SOURCE_OPTION = CHAR_MAX + 1
  86 };
  87
  88 static struct option const long_opts[] =
  89 {
  90   {"echo", no_argument, NULL, 'e'},
  91   {"input-range", required_argument, NULL, 'i'},
  92   {"head-count", required_argument, NULL, 'n'},
  93   {"output", required_argument, NULL, 'o'},
  94   {"random-source", required_argument, NULL, RANDOM_SOURCE_OPTION},
  95   {"zero-terminated", no_argument, NULL, 'z'},
  96   {GETOPT_HELP_OPTION_DECL},
  97   {GETOPT_VERSION_OPTION_DECL},
  98   {0, 0, 0, 0},
  99 };
 100
 101 static bool
 102 input_numbers_option_used (size_t lo_input, size_t hi_input)
 103 {
 104   return ! (lo_input == SIZE_MAX && hi_input == 0);
 105 }
 106
 107 static void
 108 input_from_argv (char **operand, int n_operands, char eolbyte)
 109 {
 110   char *p;
 111   size_t size = n_operands;
 112   int i;
 113
 114   for (i = 0; i < n_operands; i++)
 115     size += strlen (operand[i]);
 116   p = xmalloc (size);
 117
 118   for (i = 0; i < n_operands; i++)
 119     {
 120       char *p1 = stpcpy (p, operand[i]);
 121       operand[i] = p;
 122       p = p1;
 123       *p++ = eolbyte;
 124     }
 125
 126   operand[n_operands] = p;
 127 }
 128
 129 /* Return the start of the next line after LINE.  The current line
 130    ends in EOLBYTE, and is guaranteed to end before LINE + N.  */
 131
 132 static char *
 133 next_line (char *line, char eolbyte, size_t n)
 134 {
 135   char *p = memchr (line, eolbyte, n);
 136   return p + 1;
 137 }
 138
 139 /* Read data from file IN.  Input lines are delimited by EOLBYTE;
 140    silently append a trailing EOLBYTE if the file ends in some other
 141    byte.  Store a pointer to the resulting array of lines into *PLINE.
 142    Return the number of lines read.  Report an error and exit on
 143    failure.  */
 144
 145 static size_t
 146 read_input (FILE *in, char eolbyte, char ***pline)
 147 {
 148   char *p;
 149   char *buf = NULL;
 150   size_t used;
 151   char *lim;
 152   char **line;
 153   size_t i;
 154   size_t n_lines;
 155
 156   if (!(buf = fread_file (in, &used)))
 157     error (EXIT_FAILURE, errno, _("read error"));
 158
 159   if (used && buf[used - 1] != eolbyte)
 160     buf[used++] = eolbyte;
 161
 162   lim = buf + used;
 163
 164   n_lines = 0;
 165   for (p = buf; p < lim; p = next_line (p, eolbyte, lim - p))
 166     n_lines++;
 167
 168   *pline = line = xnmalloc (n_lines + 1, sizeof *line);
 169
 170   line[0] = p = buf;
 171   for (i = 1; i <= n_lines; i++)
 172     line[i] = p = next_line (p, eolbyte, lim - p);
 173
 174   return n_lines;
 175 }
 176
 177 static int
 178 write_permuted_output (size_t n_lines, char * const *line, size_t lo_input,
 179                        size_t const *permutation, char eolbyte)
 180 {
 181   size_t i;
 182
 183   if (line)
 184     for (i = 0; i < n_lines; i++)
 185       {
 186         char * const *p = line + permutation[i];
 187         size_t len = p[1] - p[0];
 188         if (fwrite (p[0], sizeof *p[0], len, stdout) != len)
 189           return -1;
 190       }
 191   else
 192     for (i = 0; i < n_lines; i++)
 193       {
 194         unsigned long int n = lo_input + permutation[i];
 195         if (printf ("%lu%c", n, eolbyte) < 0)
 196           return -1;
 197       }
 198
 199   return 0;
 200 }
 201
 202 int
 203 main (int argc, char **argv)
 204 {
 205   bool echo = false;
 206   size_t lo_input = SIZE_MAX;
 207   size_t hi_input = 0;
 208   size_t head_lines = SIZE_MAX;
 209   char const *outfile = NULL;
 210   char *random_source = NULL;
 211   char eolbyte = '\n';
 212   char **input_lines = NULL;
 213
 214   int optc;
 215   int n_operands;
 216   char **operand;
 217   size_t n_lines;
 218   char **line;
 219   struct randint_source *randint_source;
 220   size_t *permutation;
 221
 222   initialize_main (&argc, &argv);
 223   set_program_name (argv[0]);
 224   setlocale (LC_ALL, "");
 225   bindtextdomain (PACKAGE, LOCALEDIR);
 226   textdomain (PACKAGE);
 227
 228   atexit (close_stdout);
 229
 230   while ((optc = getopt_long (argc, argv, "ei:n:o:z", long_opts, NULL)) != -1)
 231     switch (optc)
 232       {
 233       case 'e':
 234         echo = true;
 235         break;
 236
 237       case 'i':
 238         {
 239           unsigned long int argval = 0;
 240           char *p = strchr (optarg, '-');
 241           char const *hi_optarg = optarg;
 242           bool invalid = !p;
 243
 244           if (input_numbers_option_used (lo_input, hi_input))
 245             error (EXIT_FAILURE, 0, _("multiple -i options specified"));
 246
 247           if (p)
 248             {
 249               *p = '\0';
 250               invalid = ((xstrtoul (optarg, NULL, 10, &argval, NULL)
 251                           != LONGINT_OK)
 252                          || SIZE_MAX < argval);
 253               *p = '-';
 254               lo_input = argval;
 255               hi_optarg = p + 1;
 256             }
 257
 258           invalid |= ((xstrtoul (hi_optarg, NULL, 10, &argval, NULL)
 259                        != LONGINT_OK)
 260                       || SIZE_MAX < argval);
 261           hi_input = argval;
 262           n_lines = hi_input - lo_input + 1;
 263           invalid |= ((lo_input <= hi_input) == (n_lines == 0));
 264           if (invalid)
 265             error (EXIT_FAILURE, 0, _("invalid input range %s"),
 266                    quote (optarg));
 267         }
 268         break;
 269
 270       case 'n':
 271         {
 272           unsigned long int argval;
 273           strtol_error e = xstrtoul (optarg, NULL, 10, &argval, NULL);
 274
 275           if (e == LONGINT_OK)
 276             head_lines = MIN (head_lines, argval);
 277           else if (e != LONGINT_OVERFLOW)
 278             error (EXIT_FAILURE, 0, _("invalid line count %s"),
 279                    quote (optarg));
 280         }
 281         break;
 282
 283       case 'o':
 284         if (outfile && !STREQ (outfile, optarg))
 285           error (EXIT_FAILURE, 0, _("multiple output files specified"));
 286         outfile = optarg;
 287         break;
 288
 289       case RANDOM_SOURCE_OPTION:
 290         if (random_source && !STREQ (random_source, optarg))
 291           error (EXIT_FAILURE, 0, _("multiple random sources specified"));
 292         random_source = optarg;
 293         break;
 294
 295       case 'z':
 296         eolbyte = '\0';
 297         break;
 298
 299       case_GETOPT_HELP_CHAR;
 300       case_GETOPT_VERSION_CHAR (PROGRAM_NAME, AUTHORS);
 301       default:
 302         usage (EXIT_FAILURE);
 303       }
 304
 305   n_operands = argc - optind;
 306   operand = argv + optind;
 307
 308   if (echo)
 309     {
 310       if (input_numbers_option_used (lo_input, hi_input))
 311         error (EXIT_FAILURE, 0, _("cannot combine -e and -i options"));
 312       input_from_argv (operand, n_operands, eolbyte);
 313       n_lines = n_operands;
 314       line = operand;
 315     }
 316   else if (input_numbers_option_used (lo_input, hi_input))
 317     {
 318       if (n_operands)
 319         {
 320           error (0, 0, _("extra operand %s"), quote (operand[0]));
 321           usage (EXIT_FAILURE);
 322         }
 323       n_lines = hi_input - lo_input + 1;
 324       line = NULL;
 325     }
 326   else
 327     {
 328       switch (n_operands)
 329         {
 330         case 0:
 331           break;
 332
 333         case 1:
 334           if (! (STREQ (operand[0], "-") || freopen (operand[0], "r", stdin)))
 335             error (EXIT_FAILURE, errno, "%s", operand[0]);
 336           break;
 337
 338         default:
 339           error (0, 0, _("extra operand %s"), quote (operand[1]));
 340           usage (EXIT_FAILURE);
 341         }
 342
 343       fadvise (stdin, FADVISE_SEQUENTIAL);
 344
 345       n_lines = read_input (stdin, eolbyte, &input_lines);
 346       line = input_lines;
 347     }
 348
 349   head_lines = MIN (head_lines, n_lines);
 350
 351   randint_source = randint_all_new (random_source,
 352                                     randperm_bound (head_lines, n_lines));
 353   if (! randint_source)
 354     error (EXIT_FAILURE, errno, "%s", quotearg_colon (random_source));
 355
 356   /* Close stdin now, rather than earlier, so that randint_all_new
 357      doesn't have to worry about opening something other than
 358      stdin.  */
 359   if (! (echo || input_numbers_option_used (lo_input, hi_input))
 360       && (fclose (stdin) != 0))
 361     error (EXIT_FAILURE, errno, _("read error"));
 362
 363   permutation = randperm_new (randint_source, head_lines, n_lines);
 364
 365   if (outfile && ! freopen (outfile, "w", stdout))
 366     error (EXIT_FAILURE, errno, "%s", quotearg_colon (outfile));
 367   if (write_permuted_output (head_lines, line, lo_input, permutation, eolbyte)
 368       != 0)
 369     error (EXIT_FAILURE, errno, _("write error"));
 370
 371 #ifdef lint
 372   free (permutation);
 373   randint_all_free (randint_source);
 374   if (input_lines)
 375     {
 376       free (input_lines[0]);
 377       free (input_lines);
 378     }
 379 #endif
 380
 381   return EXIT_SUCCESS;
 382 }