src/shuf.c

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