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.  */
43enum { 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  .*/
51enum { RESERVOIR_MIN_INPUT = 8192 * 1024 };
52
53
54void
55usage (int status)
56{
57  if (status != EXIT_SUCCESS)
58    emit_try_help ();
59  else
60    {
61      printf (_("\
62Usage: %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 (_("\
68Write 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.  */
95enum
96{
97  RANDOM_SOURCE_OPTION = CHAR_MAX + 1
98};
99
100static 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
114static void
115input_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
139static char *
140next_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
148static off_t
149input_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
173static size_t
174read_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
238static int
239write_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
260static size_t
261read_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.  */
305static int
306write_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.  */
324static int
325write_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.  */
342static int
343write_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.  */
362static int
363write_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
380int
381main (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}
620