src/shuf.c

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