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