src/shuf.c

   1 /* Shuffle lines of text.
   2
   3    Copyright (C) 2006-2013 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 "quotearg.h"
  31 #include "randint.h"
  32 #include "randperm.h"
  33 #include "read-file.h"
  34 #include "stdio--.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_mandatory_arg_note ();
  72
  73       fputs (_("\
  74   -e, --echo                treat each ARG as an input line\n\
  75   -i, --input-range=LO-HI   treat each number LO through HI as an input line\n\
  76   -n, --head-count=COUNT    output at most COUNT lines\n\
  77   -o, --output=FILE         write result to FILE instead of standard output\n\
  78       --random-source=FILE  get random bytes from FILE\n\
  79   -z, --zero-terminated     end lines with 0 byte, not newline\n\
  80 "), stdout);
  81       fputs (HELP_OPTION_DESCRIPTION, stdout);
  82       fputs (VERSION_OPTION_DESCRIPTION, stdout);
  83       fputs (_("\
  84 \n\
  85 With no FILE, or when FILE is -, read standard input.\n\
  86 "), stdout);
  87       emit_ancillary_info ();
  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   {"zero-terminated", no_argument, NULL, 'z'},
 108   {GETOPT_HELP_OPTION_DECL},
 109   {GETOPT_VERSION_OPTION_DECL},
 110   {0, 0, 0, 0},
 111 };
 112
 113 static bool
 114 input_numbers_option_used (size_t lo_input, size_t hi_input)
 115 {
 116   return ! (lo_input == SIZE_MAX && hi_input == 0);
 117 }
 118
 119 static void
 120 input_from_argv (char **operand, int n_operands, char eolbyte)
 121 {
 122   char *p;
 123   size_t size = n_operands;
 124   int i;
 125
 126   for (i = 0; i < n_operands; i++)
 127     size += strlen (operand[i]);
 128   p = xmalloc (size);
 129
 130   for (i = 0; i < n_operands; i++)
 131     {
 132       char *p1 = stpcpy (p, operand[i]);
 133       operand[i] = p;
 134       p = p1;
 135       *p++ = eolbyte;
 136     }
 137
 138   operand[n_operands] = p;
 139 }
 140
 141 /* Return the start of the next line after LINE.  The current line
 142    ends in EOLBYTE, and is guaranteed to end before LINE + N.  */
 143
 144 static char *
 145 next_line (char *line, char eolbyte, size_t n)
 146 {
 147   char *p = memchr (line, eolbyte, n);
 148   return p + 1;
 149 }
 150
 151 /* Return the size of the input if possible or OFF_T_MAX if not.  */
 152
 153 static off_t
 154 input_size (void)
 155 {
 156   off_t file_size;
 157
 158   struct stat stat_buf;
 159   if (fstat (STDIN_FILENO, &stat_buf) != 0)
 160     return OFF_T_MAX;
 161   if (usable_st_size (&stat_buf))
 162     file_size = stat_buf.st_size;
 163   else
 164     return OFF_T_MAX;
 165
 166   off_t input_offset = lseek (STDIN_FILENO, 0, SEEK_CUR);
 167   if (input_offset < 0)
 168     return OFF_T_MAX;
 169
 170   file_size -= input_offset;
 171
 172   return file_size;
 173 }
 174
 175 /* Read all lines and store up to K permuted lines in *OUT_RSRV.
 176    Return the number of lines read, up to a maximum of K.  */
 177
 178 static size_t
 179 read_input_reservoir_sampling (FILE *in, char eolbyte, size_t k,
 180                                struct randint_source *s,
 181                                struct linebuffer **out_rsrv)
 182 {
 183   randint n_lines = 0;
 184   size_t n_alloc_lines = MIN (k, RESERVOIR_LINES_INCREMENT);
 185   struct linebuffer *line = NULL;
 186   struct linebuffer *rsrv;
 187
 188   rsrv = xcalloc (n_alloc_lines, sizeof (struct linebuffer));
 189
 190   /* Fill the first K lines, directly into the reservoir.  */
 191   while (n_lines < k
 192          && (line =
 193              readlinebuffer_delim (&rsrv[n_lines], in, eolbyte)) != NULL)
 194     {
 195       n_lines++;
 196
 197       /* Enlarge reservoir.  */
 198       if (n_lines >= n_alloc_lines)
 199         {
 200           n_alloc_lines += RESERVOIR_LINES_INCREMENT;
 201           rsrv = xnrealloc (rsrv, n_alloc_lines, sizeof (struct linebuffer));
 202           memset (&rsrv[n_lines], 0,
 203                   RESERVOIR_LINES_INCREMENT * sizeof (struct linebuffer));
 204         }
 205     }
 206
 207   /* last line wasn't NULL - so there may be more lines to read.  */
 208   if (line != NULL)
 209     {
 210       struct linebuffer dummy;
 211       initbuffer (&dummy);  /* space for lines not put in reservoir.  */
 212
 213       /* Choose the fate of the next line, with decreasing probability (as
 214          n_lines increases in size).
 215
 216          If the line will be used, store it directly in the reservoir.
 217          Otherwise, store it in dummy space.
 218
 219          With 'struct linebuffer', storing into existing buffer will reduce
 220          re-allocations (will only re-allocate if the new line is longer than
 221          the currently allocated space).  */
 222       do
 223         {
 224           randint j = randint_choose (s, n_lines + 1);  /* 0 .. n_lines.  */
 225           line = (j < k) ? (&rsrv[j]) : (&dummy);
 226         }
 227       while (readlinebuffer_delim (line, in, eolbyte) != NULL && n_lines++);
 228
 229       if (! n_lines)
 230         error (EXIT_FAILURE, EOVERFLOW, _("too many input lines"));
 231
 232       freebuffer (&dummy);
 233     }
 234
 235   /* no more input lines, or an input error.  */
 236   if (ferror (in))
 237     error (EXIT_FAILURE, errno, _("read error"));
 238
 239   *out_rsrv = rsrv;
 240   return MIN (k, n_lines);
 241 }
 242
 243 static int
 244 write_permuted_output_reservoir (size_t n_lines, struct linebuffer *lines,
 245                                  size_t const *permutation)
 246 {
 247   size_t i;
 248
 249   for (i = 0; i < n_lines; i++)
 250     {
 251       const struct linebuffer *p = &lines[permutation[i]];
 252       if (fwrite (p->buffer, sizeof (char), p->length, stdout) != p->length)
 253         return -1;
 254     }
 255
 256   return 0;
 257 }
 258
 259 /* Read data from file IN.  Input lines are delimited by EOLBYTE;
 260    silently append a trailing EOLBYTE if the file ends in some other
 261    byte.  Store a pointer to the resulting array of lines into *PLINE.
 262    Return the number of lines read.  Report an error and exit on
 263    failure.  */
 264
 265 static size_t
 266 read_input (FILE *in, char eolbyte, char ***pline)
 267 {
 268   char *p;
 269   char *buf = NULL;
 270   size_t used;
 271   char *lim;
 272   char **line;
 273   size_t i;
 274   size_t n_lines;
 275
 276   /* TODO: We should limit the amount of data read here,
 277      to less than RESERVOIR_MIN_INPUT.  I.E. adjust fread_file() to support
 278      taking a byte limit.  We'd then need to ensure we handle a line spanning
 279      this boundary.  With that in place we could set use_reservoir_sampling
 280      when used==RESERVOIR_MIN_INPUT, and have read_input_reservoir_sampling()
 281      call a wrapper function to populate a linebuffer from the internal pline
 282      or if none left, stdin.  Doing that would give better performance by
 283      avoiding the reservoir CPU overhead when reading < RESERVOIR_MIN_INPUT
 284      from a pipe, and allow us to dispense with the input_size() function.  */
 285   if (!(buf = fread_file (in, &used)))
 286     error (EXIT_FAILURE, errno, _("read error"));
 287
 288   if (used && buf[used - 1] != eolbyte)
 289     buf[used++] = eolbyte;
 290
 291   lim = buf + used;
 292
 293   n_lines = 0;
 294   for (p = buf; p < lim; p = next_line (p, eolbyte, lim - p))
 295     n_lines++;
 296
 297   *pline = line = xnmalloc (n_lines + 1, sizeof *line);
 298
 299   line[0] = p = buf;
 300   for (i = 1; i <= n_lines; i++)
 301     line[i] = p = next_line (p, eolbyte, lim - p);
 302
 303   return n_lines;
 304 }
 305
 306 static int
 307 write_permuted_output (size_t n_lines, char *const *line, size_t lo_input,
 308                        size_t const *permutation, char eolbyte)
 309 {
 310   size_t i;
 311
 312   if (line)
 313     for (i = 0; i < n_lines; i++)
 314       {
 315         char *const *p = line + permutation[i];
 316         size_t len = p[1] - p[0];
 317         if (fwrite (p[0], sizeof *p[0], len, stdout) != len)
 318           return -1;
 319       }
 320   else
 321     for (i = 0; i < n_lines; i++)
 322       {
 323         unsigned long int n = lo_input + permutation[i];
 324         if (printf ("%lu%c", n, eolbyte) < 0)
 325           return -1;
 326       }
 327
 328   return 0;
 329 }
 330
 331 int
 332 main (int argc, char **argv)
 333 {
 334   bool echo = false;
 335   size_t lo_input = SIZE_MAX;
 336   size_t hi_input = 0;
 337   size_t head_lines = SIZE_MAX;
 338   char const *outfile = NULL;
 339   char *random_source = NULL;
 340   char eolbyte = '\n';
 341   char **input_lines = NULL;
 342   bool use_reservoir_sampling = false;
 343
 344   int optc;
 345   int n_operands;
 346   char **operand;
 347   size_t n_lines;
 348   char **line = NULL;
 349   struct linebuffer *reservoir = NULL;
 350   struct randint_source *randint_source;
 351   size_t *permutation;
 352   int i;
 353
 354   initialize_main (&argc, &argv);
 355   set_program_name (argv[0]);
 356   setlocale (LC_ALL, "");
 357   bindtextdomain (PACKAGE, LOCALEDIR);
 358   textdomain (PACKAGE);
 359
 360   atexit (close_stdout);
 361
 362   while ((optc = getopt_long (argc, argv, "ei:n:o:z", long_opts, NULL)) != -1)
 363     switch (optc)
 364       {
 365       case 'e':
 366         echo = true;
 367         break;
 368
 369       case 'i':
 370         {
 371           unsigned long int argval = 0;
 372           char *p = strchr (optarg, '-');
 373           char const *hi_optarg = optarg;
 374           bool invalid = !p;
 375
 376           if (input_numbers_option_used (lo_input, hi_input))
 377             error (EXIT_FAILURE, 0, _("multiple -i options specified"));
 378
 379           if (p)
 380             {
 381               *p = '\0';
 382               invalid = ((xstrtoul (optarg, NULL, 10, &argval, NULL)
 383                           != LONGINT_OK)
 384                          || SIZE_MAX < argval);
 385               *p = '-';
 386               lo_input = argval;
 387               hi_optarg = p + 1;
 388             }
 389
 390           invalid |= ((xstrtoul (hi_optarg, NULL, 10, &argval, NULL)
 391                        != LONGINT_OK)
 392                       || SIZE_MAX < argval);
 393           hi_input = argval;
 394           n_lines = hi_input - lo_input + 1;
 395           invalid |= ((lo_input <= hi_input) == (n_lines == 0));
 396           if (invalid)
 397             error (EXIT_FAILURE, 0, _("invalid input range %s"),
 398                    quote (optarg));
 399         }
 400         break;
 401
 402       case 'n':
 403         {
 404           unsigned long int argval;
 405           strtol_error e = xstrtoul (optarg, NULL, 10, &argval, NULL);
 406
 407           if (e == LONGINT_OK)
 408             head_lines = MIN (head_lines, argval);
 409           else if (e != LONGINT_OVERFLOW)
 410             error (EXIT_FAILURE, 0, _("invalid line count %s"),
 411                    quote (optarg));
 412         }
 413         break;
 414
 415       case 'o':
 416         if (outfile && !STREQ (outfile, optarg))
 417           error (EXIT_FAILURE, 0, _("multiple output files specified"));
 418         outfile = optarg;
 419         break;
 420
 421       case RANDOM_SOURCE_OPTION:
 422         if (random_source && !STREQ (random_source, optarg))
 423           error (EXIT_FAILURE, 0, _("multiple random sources specified"));
 424         random_source = optarg;
 425         break;
 426
 427       case 'z':
 428         eolbyte = '\0';
 429         break;
 430
 431       case_GETOPT_HELP_CHAR;
 432       case_GETOPT_VERSION_CHAR (PROGRAM_NAME, AUTHORS);
 433       default:
 434         usage (EXIT_FAILURE);
 435       }
 436
 437   n_operands = argc - optind;
 438   operand = argv + optind;
 439
 440   if (echo)
 441     {
 442       if (input_numbers_option_used (lo_input, hi_input))
 443         error (EXIT_FAILURE, 0, _("cannot combine -e and -i options"));
 444       input_from_argv (operand, n_operands, eolbyte);
 445       n_lines = n_operands;
 446       line = operand;
 447     }
 448   else if (input_numbers_option_used (lo_input, hi_input))
 449     {
 450       if (n_operands)
 451         {
 452           error (0, 0, _("extra operand %s"), quote (operand[0]));
 453           usage (EXIT_FAILURE);
 454         }
 455       n_lines = hi_input - lo_input + 1;
 456       line = NULL;
 457     }
 458   else
 459     {
 460       switch (n_operands)
 461         {
 462         case 0:
 463           break;
 464
 465         case 1:
 466           if (! (STREQ (operand[0], "-") || ! head_lines
 467                  || freopen (operand[0], "r", stdin)))
 468             error (EXIT_FAILURE, errno, "%s", operand[0]);
 469           break;
 470
 471         default:
 472           error (0, 0, _("extra operand %s"), quote (operand[1]));
 473           usage (EXIT_FAILURE);
 474         }
 475
 476       fadvise (stdin, FADVISE_SEQUENTIAL);
 477
 478       if (head_lines != SIZE_MAX && (! head_lines
 479                                      || input_size () > RESERVOIR_MIN_INPUT))
 480         {
 481           use_reservoir_sampling = true;
 482           n_lines = SIZE_MAX;   /* unknown number of input lines, for now.  */
 483         }
 484       else
 485         {
 486           n_lines = read_input (stdin, eolbyte, &input_lines);
 487           line = input_lines;
 488         }
 489     }
 490
 491   head_lines = MIN (head_lines, n_lines);
 492
 493   randint_source = randint_all_new (random_source,
 494                                     use_reservoir_sampling ? SIZE_MAX :
 495                                     randperm_bound (head_lines, n_lines));
 496   if (! randint_source)
 497     error (EXIT_FAILURE, errno, "%s", quotearg_colon (random_source));
 498
 499   if (use_reservoir_sampling)
 500     {
 501       /* Instead of reading the entire file into 'line',
 502          use reservoir-sampling to store just "head_lines" random lines.  */
 503       n_lines = read_input_reservoir_sampling (stdin, eolbyte, head_lines,
 504                                                randint_source, &reservoir);
 505       head_lines = n_lines;
 506     }
 507
 508   /* Close stdin now, rather than earlier, so that randint_all_new
 509      doesn't have to worry about opening something other than
 510      stdin.  */
 511   if (! (echo || input_numbers_option_used (lo_input, hi_input))
 512       && (fclose (stdin) != 0))
 513     error (EXIT_FAILURE, errno, _("read error"));
 514
 515   permutation = randperm_new (randint_source, head_lines, n_lines);
 516
 517   if (outfile && ! freopen (outfile, "w", stdout))
 518     error (EXIT_FAILURE, errno, "%s", quotearg_colon (outfile));
 519
 520   if (use_reservoir_sampling)
 521     i = write_permuted_output_reservoir (n_lines, reservoir, permutation);
 522   else
 523     i = write_permuted_output (head_lines, line, lo_input,
 524                                permutation, eolbyte);
 525   if (i != 0)
 526     error (EXIT_FAILURE, errno, _("write error"));
 527
 528 #ifdef lint
 529   free (permutation);
 530   randint_all_free (randint_source);
 531   if (input_lines)
 532     {
 533       free (input_lines[0]);
 534       free (input_lines);
 535     }
 536   if (reservoir)
 537     {
 538       size_t j;
 539       for (j = 0; j < n_lines; ++j)
 540         freebuffer (&reservoir[j]);
 541       free (reservoir);
 542     }
 543 #endif
 544
 545   return EXIT_SUCCESS;
 546 }