src/shuf.c

   1 /* Shuffle lines of text.
   2
   3    Copyright (C) 2006-2024 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 <getopt.h>
  24 #include "system.h"
  25
  26 #include "fadvise.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     line = nullptr;
 513   else
 514     {
 515       /* If an input file is specified, re-open it as stdin.  */
 516       if (n_operands == 1
 517           && ! (STREQ (operand[0], "-")
 518                 || freopen (operand[0], "r", stdin)))
 519         error (EXIT_FAILURE, errno, "%s", quotef (operand[0]));
 520
 521       fadvise (stdin, FADVISE_SEQUENTIAL);
 522
 523       if (repeat || head_lines == SIZE_MAX
 524           || input_size () <= RESERVOIR_MIN_INPUT)
 525         {
 526           n_lines = read_input (stdin, eolbyte, &input_lines);
 527           line = input_lines;
 528         }
 529       else
 530         {
 531           use_reservoir_sampling = true;
 532           n_lines = SIZE_MAX;   /* unknown number of input lines, for now.  */
 533         }
 534     }
 535
 536   /* The adjusted head line count; can be less than HEAD_LINES if the
 537      input is small and if not repeating.  */
 538   size_t ahead_lines = repeat || head_lines < n_lines ? head_lines : n_lines;
 539
 540   randint_source = randint_all_new (random_source,
 541                                     (use_reservoir_sampling || repeat
 542                                      ? SIZE_MAX
 543                                      : randperm_bound (ahead_lines, n_lines)));
 544   if (! randint_source)
 545     error (EXIT_FAILURE, errno, "%s",
 546            quotef (random_source ? random_source : "getrandom"));
 547
 548   if (use_reservoir_sampling)
 549     {
 550       /* Instead of reading the entire file into 'line',
 551          use reservoir-sampling to store just AHEAD_LINES random lines.  */
 552       n_lines = read_input_reservoir_sampling (stdin, eolbyte, ahead_lines,
 553                                                randint_source, &reservoir);
 554       ahead_lines = n_lines;
 555     }
 556
 557   /* Close stdin now, rather than earlier, so that randint_all_new
 558      doesn't have to worry about opening something other than
 559      stdin.  */
 560   if (! (head_lines == 0 || echo || input_range || fclose (stdin) == 0))
 561     error (EXIT_FAILURE, errno, _("read error"));
 562
 563   if (!repeat)
 564     permutation = randperm_new (randint_source, ahead_lines, n_lines);
 565
 566   if (outfile && ! freopen (outfile, "w", stdout))
 567     error (EXIT_FAILURE, errno, "%s", quotef (outfile));
 568
 569   /* Generate output according to requested method */
 570   if (repeat)
 571     {
 572       if (head_lines == 0)
 573         i = 0;
 574       else
 575         {
 576           if (n_lines == 0)
 577             error (EXIT_FAILURE, 0, _("no lines to repeat"));
 578           if (input_range)
 579             i = write_random_numbers (randint_source, ahead_lines,
 580                                       lo_input, hi_input, eolbyte);
 581           else
 582             i = write_random_lines (randint_source, ahead_lines, line, n_lines);
 583         }
 584     }
 585   else
 586     {
 587       if (use_reservoir_sampling)
 588         i = write_permuted_output_reservoir (n_lines, reservoir, permutation);
 589       else if (input_range)
 590         i = write_permuted_numbers (ahead_lines, lo_input,
 591                                     permutation, eolbyte);
 592       else
 593         i = write_permuted_lines (ahead_lines, line, permutation);
 594     }
 595
 596   if (i != 0)
 597     write_error ();
 598
 599   main_exit (EXIT_SUCCESS);
 600 }