src/shuf.c

   1 /* Shuffle lines of text.
   2
   3    Copyright (C) 2006-2016 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 "linebuffer.h"
  29 #include "quote.h"
  30 #include "randint.h"
  31 #include "randperm.h"
  32 #include "read-file.h"
  33 #include "stdio--.h"
  34 #include "xdectoint.h"
  35 #include "xstrtol.h"
  36
  37 /* The official name of this program (e.g., no 'g' prefix).  */
  38 #define PROGRAM_NAME "shuf"
  39
  40 #define AUTHORS proper_name ("Paul Eggert")
  41
  42 /* For reservoir-sampling, allocate the reservoir lines in batches.  */
  43 enum { RESERVOIR_LINES_INCREMENT = 1024 };
  44
  45 /* reservoir-sampling introduces CPU overhead for small inputs.
  46    So only enable it for inputs >= this limit.
  47    This limit was determined using these commands:
  48      $ for p in $(seq 7); do src/seq $((10**$p)) > 10p$p.in; done
  49      $ for p in $(seq 7); do time shuf-nores -n10 10p$p.in >/dev/null; done
  50      $ for p in $(seq 7); do time shuf -n10 10p$p.in >/dev/null; done  .*/
  51 enum { RESERVOIR_MIN_INPUT = 8192 * 1024 };
  52
  53
  54 void
  55 usage (int status)
  56 {
  57   if (status != EXIT_SUCCESS)
  58     emit_try_help ();
  59   else
  60     {
  61       printf (_("\
  62 Usage: %s [OPTION]... [FILE]\n\
  63   or:  %s -e [OPTION]... [ARG]...\n\
  64   or:  %s -i LO-HI [OPTION]...\n\
  65 "),
  66               program_name, program_name, program_name);
  67       fputs (_("\
  68 Write a random permutation of the input lines to standard output.\n\
  69 "), stdout);
  70
  71       emit_stdin_note ();
  72       emit_mandatory_arg_note ();
  73
  74       fputs (_("\
  75   -e, --echo                treat each ARG as an input line\n\
  76   -i, --input-range=LO-HI   treat each number LO through HI as an input line\n\
  77   -n, --head-count=COUNT    output at most COUNT lines\n\
  78   -o, --output=FILE         write result to FILE instead of standard output\n\
  79       --random-source=FILE  get random bytes from FILE\n\
  80   -r, --repeat              output lines can be repeated\n\
  81 "), stdout);
  82       fputs (_("\
  83   -z, --zero-terminated     line delimiter is NUL, not newline\n\
  84 "), stdout);
  85       fputs (HELP_OPTION_DESCRIPTION, stdout);
  86       fputs (VERSION_OPTION_DESCRIPTION, stdout);
  87       emit_ancillary_info (PROGRAM_NAME);
  88     }
  89
  90   exit (status);
  91 }
  92
  93 /* For long options that have no equivalent short option, use a
  94    non-character as a pseudo short option, starting with CHAR_MAX + 1.  */
  95 enum
  96 {
  97   RANDOM_SOURCE_OPTION = CHAR_MAX + 1
  98 };
  99
 100 static struct option const long_opts[] =
 101 {
 102   {"echo", no_argument, NULL, 'e'},
 103   {"input-range", required_argument, NULL, 'i'},
 104   {"head-count", required_argument, NULL, 'n'},
 105   {"output", required_argument, NULL, 'o'},
 106   {"random-source", required_argument, NULL, RANDOM_SOURCE_OPTION},
 107   {"repeat", no_argument, NULL, 'r'},
 108   {"zero-terminated", no_argument, NULL, 'z'},
 109   {GETOPT_HELP_OPTION_DECL},
 110   {GETOPT_VERSION_OPTION_DECL},
 111   {0, 0, 0, 0},
 112 };
 113
 114 static void
 115 input_from_argv (char **operand, int n_operands, char eolbyte)
 116 {
 117   char *p;
 118   size_t size = n_operands;
 119   int i;
 120
 121   for (i = 0; i < n_operands; i++)
 122     size += strlen (operand[i]);
 123   p = xmalloc (size);
 124
 125   for (i = 0; i < n_operands; i++)
 126     {
 127       char *p1 = stpcpy (p, operand[i]);
 128       operand[i] = p;
 129       p = p1;
 130       *p++ = eolbyte;
 131     }
 132
 133   operand[n_operands] = p;
 134 }
 135
 136 /* Return the start of the next line after LINE.  The current line
 137    ends in EOLBYTE, and is guaranteed to end before LINE + N.  */
 138
 139 static char *
 140 next_line (char *line, char eolbyte, size_t n)
 141 {
 142   char *p = memchr (line, eolbyte, n);
 143   return p + 1;
 144 }
 145
 146 /* Return the size of the input if possible or OFF_T_MAX if not.  */
 147
 148 static off_t
 149 input_size (void)
 150 {
 151   off_t file_size;
 152
 153   struct stat stat_buf;
 154   if (fstat (STDIN_FILENO, &stat_buf) != 0)
 155     return OFF_T_MAX;
 156   if (usable_st_size (&stat_buf))
 157     file_size = stat_buf.st_size;
 158   else
 159     return OFF_T_MAX;
 160
 161   off_t input_offset = lseek (STDIN_FILENO, 0, SEEK_CUR);
 162   if (input_offset < 0)
 163     return OFF_T_MAX;
 164
 165   file_size -= input_offset;
 166
 167   return file_size;
 168 }
 169
 170 /* Read all lines and store up to K permuted lines in *OUT_RSRV.
 171    Return the number of lines read, up to a maximum of K.  */
 172
 173 static size_t
 174 read_input_reservoir_sampling (FILE *in, char eolbyte, size_t k,
 175                                struct randint_source *s,
 176                                struct linebuffer **out_rsrv)
 177 {
 178   randint n_lines = 0;
 179   size_t n_alloc_lines = MIN (k, RESERVOIR_LINES_INCREMENT);
 180   struct linebuffer *line = NULL;
 181   struct linebuffer *rsrv;
 182
 183   rsrv = xcalloc (n_alloc_lines, sizeof (struct linebuffer));
 184
 185   /* Fill the first K lines, directly into the reservoir.  */
 186   while (n_lines < k
 187          && (line =
 188              readlinebuffer_delim (&rsrv[n_lines], in, eolbyte)) != NULL)
 189     {
 190       n_lines++;
 191
 192       /* Enlarge reservoir.  */
 193       if (n_lines >= n_alloc_lines)
 194         {
 195           n_alloc_lines += RESERVOIR_LINES_INCREMENT;
 196           rsrv = xnrealloc (rsrv, n_alloc_lines, sizeof (struct linebuffer));
 197           memset (&rsrv[n_lines], 0,
 198                   RESERVOIR_LINES_INCREMENT * sizeof (struct linebuffer));
 199         }
 200     }
 201
 202   /* last line wasn't NULL - so there may be more lines to read.  */
 203   if (line != NULL)
 204     {
 205       struct linebuffer dummy;
 206       initbuffer (&dummy);  /* space for lines not put in reservoir.  */
 207
 208       /* Choose the fate of the next line, with decreasing probability (as
 209          n_lines increases in size).
 210
 211          If the line will be used, store it directly in the reservoir.
 212          Otherwise, store it in dummy space.
 213
 214          With 'struct linebuffer', storing into existing buffer will reduce
 215          re-allocations (will only re-allocate if the new line is longer than
 216          the currently allocated space).  */
 217       do
 218         {
 219           randint j = randint_choose (s, n_lines + 1);  /* 0 .. n_lines.  */
 220           line = (j < k) ? (&rsrv[j]) : (&dummy);
 221         }
 222       while (readlinebuffer_delim (line, in, eolbyte) != NULL && n_lines++);
 223
 224       if (! n_lines)
 225         error (EXIT_FAILURE, EOVERFLOW, _("too many input lines"));
 226
 227       freebuffer (&dummy);
 228     }
 229
 230   /* no more input lines, or an input error.  */
 231   if (ferror (in))
 232     error (EXIT_FAILURE, errno, _("read error"));
 233
 234   *out_rsrv = rsrv;
 235   return MIN (k, n_lines);
 236 }
 237
 238 static int
 239 write_permuted_output_reservoir (size_t n_lines, struct linebuffer *lines,
 240                                  size_t const *permutation)
 241 {
 242   size_t i;
 243
 244   for (i = 0; i < n_lines; i++)
 245     {
 246       const struct linebuffer *p = &lines[permutation[i]];
 247       if (fwrite (p->buffer, sizeof (char), p->length, stdout) != p->length)
 248         return -1;
 249     }
 250
 251   return 0;
 252 }
 253
 254 /* Read data from file IN.  Input lines are delimited by EOLBYTE;
 255    silently append a trailing EOLBYTE if the file ends in some other
 256    byte.  Store a pointer to the resulting array of lines into *PLINE.
 257    Return the number of lines read.  Report an error and exit on
 258    failure.  */
 259
 260 static size_t
 261 read_input (FILE *in, char eolbyte, char ***pline)
 262 {
 263   char *p;
 264   char *buf = NULL;
 265   size_t used;
 266   char *lim;
 267   char **line;
 268   size_t i;
 269   size_t n_lines;
 270
 271   /* TODO: We should limit the amount of data read here,
 272      to less than RESERVOIR_MIN_INPUT.  I.e., adjust fread_file() to support
 273      taking a byte limit.  We'd then need to ensure we handle a line spanning
 274      this boundary.  With that in place we could set use_reservoir_sampling
 275      when used==RESERVOIR_MIN_INPUT, and have read_input_reservoir_sampling()
 276      call a wrapper function to populate a linebuffer from the internal pline
 277      or if none left, stdin.  Doing that would give better performance by
 278      avoiding the reservoir CPU overhead when reading < RESERVOIR_MIN_INPUT
 279      from a pipe, and allow us to dispense with the input_size() function.  */
 280   if (!(buf = fread_file (in, &used)))
 281     error (EXIT_FAILURE, errno, _("read error"));
 282
 283   if (used && buf[used - 1] != eolbyte)
 284     buf[used++] = eolbyte;
 285
 286   lim = buf + used;
 287
 288   n_lines = 0;
 289   for (p = buf; p < lim; p = next_line (p, eolbyte, lim - p))
 290     n_lines++;
 291
 292   *pline = line = xnmalloc (n_lines + 1, sizeof *line);
 293
 294   line[0] = p = buf;
 295   for (i = 1; i <= n_lines; i++)
 296     line[i] = p = next_line (p, eolbyte, lim - p);
 297
 298   return n_lines;
 299 }
 300
 301 /* Output N_LINES lines to stdout from LINE array,
 302    chosen by the indices in PERMUTATION.
 303    PERMUTATION and LINE must have at least N_LINES elements.
 304    Strings in LINE must include the line-terminator character.  */
 305 static int
 306 write_permuted_lines (size_t n_lines, char *const *line,
 307                       size_t const *permutation)
 308 {
 309   size_t i;
 310
 311   for (i = 0; i < n_lines; i++)
 312     {
 313       char *const *p = line + permutation[i];
 314       size_t len = p[1] - p[0];
 315       if (fwrite (p[0], sizeof *p[0], len, stdout) != len)
 316         return -1;
 317     }
 318
 319   return 0;
 320 }
 321
 322 /* Output N_LINES of numbers to stdout, from PERMUTATION array.
 323    PERMUTATION must have at least N_LINES elements.  */
 324 static int
 325 write_permuted_numbers (size_t n_lines, size_t lo_input,
 326                         size_t const *permutation, char eolbyte)
 327 {
 328   size_t i;
 329
 330   for (i = 0; i < n_lines; i++)
 331     {
 332       unsigned long int n = lo_input + permutation[i];
 333       if (printf ("%lu%c", n, eolbyte) < 0)
 334         return -1;
 335     }
 336
 337   return 0;
 338 }
 339
 340 /* Output COUNT numbers to stdout, chosen randomly from range
 341    LO_INPUT through HI_INPUT.  */
 342 static int
 343 write_random_numbers (struct randint_source *s, size_t count,
 344                       size_t lo_input, size_t hi_input, char eolbyte)
 345 {
 346   size_t i;
 347   const randint range = hi_input - lo_input + 1;
 348
 349   for (i = 0; i < count; i++)
 350     {
 351       unsigned long int j = lo_input + randint_choose (s, range);
 352       if (printf ("%lu%c", j, eolbyte) < 0)
 353         return -1;
 354     }
 355
 356   return 0;
 357 }
 358
 359 /* Output COUNT lines to stdout from LINES array.
 360    LINES must have at least N_LINES elements in it.
 361    Strings in LINES_ must include the line-terminator character.  */
 362 static int
 363 write_random_lines (struct randint_source *s, size_t count,
 364                     char *const *lines, size_t n_lines)
 365 {
 366   size_t i;
 367
 368   for (i = 0; i < count; i++)
 369     {
 370       const randint j = randint_choose (s, n_lines);
 371       char *const *p = lines + j;
 372       size_t len = p[1] - p[0];
 373       if (fwrite (p[0], sizeof *p[0], len, stdout) != len)
 374         return -1;
 375     }
 376
 377   return 0;
 378 }
 379
 380 int
 381 main (int argc, char **argv)
 382 {
 383   bool echo = false;
 384   bool input_range = false;
 385   size_t lo_input = SIZE_MAX;
 386   size_t hi_input = 0;
 387   size_t head_lines = SIZE_MAX;
 388   char const *outfile = NULL;
 389   char *random_source = NULL;
 390   char eolbyte = '\n';
 391   char **input_lines = NULL;
 392   bool use_reservoir_sampling = false;
 393   bool repeat = false;
 394
 395   int optc;
 396   int n_operands;
 397   char **operand;
 398   size_t n_lines;
 399   char **line = NULL;
 400   struct linebuffer *reservoir = NULL;
 401   struct randint_source *randint_source;
 402   size_t *permutation = NULL;
 403   int i;
 404
 405   initialize_main (&argc, &argv);
 406   set_program_name (argv[0]);
 407   setlocale (LC_ALL, "");
 408   bindtextdomain (PACKAGE, LOCALEDIR);
 409   textdomain (PACKAGE);
 410
 411   atexit (close_stdout);
 412
 413   while ((optc = getopt_long (argc, argv, "ei:n:o:rz", long_opts, NULL)) != -1)
 414     switch (optc)
 415       {
 416       case 'e':
 417         echo = true;
 418         break;
 419
 420       case 'i':
 421         {
 422           char *p = strchr (optarg, '-');
 423           char const *hi_optarg = optarg;
 424           bool invalid = !p;
 425
 426           if (input_range)
 427             error (EXIT_FAILURE, 0, _("multiple -i options specified"));
 428           input_range = true;
 429
 430           if (p)
 431             {
 432               *p = '\0';
 433               lo_input = xdectoumax (optarg, 0, SIZE_MAX, "",
 434                                      _("invalid input range"), 0);
 435               *p = '-';
 436               hi_optarg = p + 1;
 437             }
 438
 439           hi_input = xdectoumax (hi_optarg, 0, SIZE_MAX, "",
 440                                  _("invalid input range"), 0);
 441
 442           n_lines = hi_input - lo_input + 1;
 443           invalid |= ((lo_input <= hi_input) == (n_lines == 0));
 444           if (invalid)
 445             error (EXIT_FAILURE, errno, "%s: %s", _("invalid input range"),
 446                    quote (optarg));
 447         }
 448         break;
 449
 450       case 'n':
 451         {
 452           unsigned long int argval;
 453           strtol_error e = xstrtoul (optarg, NULL, 10, &argval, NULL);
 454
 455           if (e == LONGINT_OK)
 456             head_lines = MIN (head_lines, argval);
 457           else if (e != LONGINT_OVERFLOW)
 458             error (EXIT_FAILURE, 0, _("invalid line count: %s"),
 459                    quote (optarg));
 460         }
 461         break;
 462
 463       case 'o':
 464         if (outfile && !STREQ (outfile, optarg))
 465           error (EXIT_FAILURE, 0, _("multiple output files specified"));
 466         outfile = optarg;
 467         break;
 468
 469       case RANDOM_SOURCE_OPTION:
 470         if (random_source && !STREQ (random_source, optarg))
 471           error (EXIT_FAILURE, 0, _("multiple random sources specified"));
 472         random_source = optarg;
 473         break;
 474
 475       case 'r':
 476         repeat = true;
 477         break;
 478
 479       case 'z':
 480         eolbyte = '\0';
 481         break;
 482
 483       case_GETOPT_HELP_CHAR;
 484       case_GETOPT_VERSION_CHAR (PROGRAM_NAME, AUTHORS);
 485       default:
 486         usage (EXIT_FAILURE);
 487       }
 488
 489   n_operands = argc - optind;
 490   operand = argv + optind;
 491
 492   /* Check invalid usage.  */
 493   if (echo && input_range)
 494     {
 495       error (0, 0, _("cannot combine -e and -i options"));
 496       usage (EXIT_FAILURE);
 497     }
 498   if (input_range ? 0 < n_operands : !echo && 1 < n_operands)
 499     {
 500       error (0, 0, _("extra operand %s"), quote (operand[!input_range]));
 501       usage (EXIT_FAILURE);
 502     }
 503
 504   /* Prepare input.  */
 505   if (echo)
 506     {
 507       input_from_argv (operand, n_operands, eolbyte);
 508       n_lines = n_operands;
 509       line = operand;
 510     }
 511   else if (input_range)
 512     {
 513       n_lines = hi_input - lo_input + 1;
 514       line = NULL;
 515     }
 516   else
 517     {
 518       /* If an input file is specified, re-open it as stdin.  */
 519       if (n_operands == 1)
 520         if (! (STREQ (operand[0], "-") || ! head_lines
 521                || freopen (operand[0], "r", stdin)))
 522           error (EXIT_FAILURE, errno, "%s", quotef (operand[0]));
 523
 524       fadvise (stdin, FADVISE_SEQUENTIAL);
 525
 526       if (! repeat && head_lines != SIZE_MAX
 527           && (! head_lines || input_size () > RESERVOIR_MIN_INPUT))
 528         {
 529           use_reservoir_sampling = true;
 530           n_lines = SIZE_MAX;   /* unknown number of input lines, for now.  */
 531         }
 532       else
 533         {
 534           n_lines = read_input (stdin, eolbyte, &input_lines);
 535           line = input_lines;
 536         }
 537     }
 538
 539   if (! repeat)
 540     head_lines = MIN (head_lines, n_lines);
 541
 542   randint_source = randint_all_new (random_source,
 543                                     (use_reservoir_sampling || repeat
 544                                      ? SIZE_MAX
 545                                      : randperm_bound (head_lines, n_lines)));
 546   if (! randint_source)
 547     error (EXIT_FAILURE, errno, "%s", quotef (random_source));
 548
 549   if (use_reservoir_sampling)
 550     {
 551       /* Instead of reading the entire file into 'line',
 552          use reservoir-sampling to store just "head_lines" random lines.  */
 553       n_lines = read_input_reservoir_sampling (stdin, eolbyte, head_lines,
 554                                                randint_source, &reservoir);
 555       head_lines = n_lines;
 556     }
 557
 558   /* Close stdin now, rather than earlier, so that randint_all_new
 559      doesn't have to worry about opening something other than
 560      stdin.  */
 561   if (! (echo || input_range)
 562       && (fclose (stdin) != 0))
 563     error (EXIT_FAILURE, errno, _("read error"));
 564
 565   if (!repeat)
 566     permutation = randperm_new (randint_source, head_lines, n_lines);
 567
 568   if (outfile && ! freopen (outfile, "w", stdout))
 569     error (EXIT_FAILURE, errno, "%s", quotef (outfile));
 570
 571   /* Generate output according to requested method */
 572   if (repeat)
 573     {
 574       if (head_lines == 0)
 575         i = 0;
 576       else
 577         {
 578           if (n_lines == 0)
 579             error (EXIT_FAILURE, 0, _("no lines to repeat"));
 580           if (input_range)
 581             i = write_random_numbers (randint_source, head_lines,
 582                                       lo_input, hi_input, eolbyte);
 583           else
 584             i = write_random_lines (randint_source, head_lines, line, n_lines);
 585         }
 586     }
 587   else
 588     {
 589       if (use_reservoir_sampling)
 590         i = write_permuted_output_reservoir (n_lines, reservoir, permutation);
 591       else if (input_range)
 592         i = write_permuted_numbers (head_lines, lo_input,
 593                                     permutation, eolbyte);
 594       else
 595         i = write_permuted_lines (head_lines, line, permutation);
 596     }
 597
 598   if (i != 0)
 599     error (EXIT_FAILURE, errno, _("write error"));
 600
 601 #ifdef lint
 602   free (permutation);
 603   randint_all_free (randint_source);
 604   if (input_lines)
 605     {
 606       free (input_lines[0]);
 607       free (input_lines);
 608     }
 609   if (reservoir)
 610     {
 611       size_t j;
 612       for (j = 0; j < n_lines; ++j)
 613         freebuffer (&reservoir[j]);
 614       free (reservoir);
 615     }
 616 #endif
 617
 618   return EXIT_SUCCESS;
 619 }