src/shuf.c

   1 /* Shuffle lines of text.
   2
   3    Copyright (C) 2006, 2007 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 <stdio.h>
  23 #include <sys/types.h>
  24 #include "system.h"
  25
  26 #include "error.h"
  27 #include "getopt.h"
  28 #include "quote.h"
  29 #include "quotearg.h"
  30 #include "randint.h"
  31 #include "randperm.h"
  32 #include "xstrtol.h"
  33
  34 /* The official name of this program (e.g., no `g' prefix).  */
  35 #define PROGRAM_NAME "shuf"
  36
  37 #define AUTHORS "Paul Eggert"
  38
  39 /* The name this program was run with. */
  40 char *program_name;
  41
  42 void
  43 usage (int status)
  44 {
  45   if (status != EXIT_SUCCESS)
  46     fprintf (stderr, _("Try `%s --help' for more information.\n"),
  47              program_name);
  48   else
  49     {
  50       printf (_("\
  51 Usage: %s [OPTION]... [FILE]\n\
  52   or:  %s -e [OPTION]... [ARG]...\n\
  53   or:  %s -i LO-HI [OPTION]...\n\
  54 "),
  55               program_name, program_name, program_name);
  56       fputs (_("\
  57 Write a random permutation of the input lines to standard output.\n\
  58 \n\
  59 "), stdout);
  60       fputs (_("\
  61 Mandatory arguments to long options are mandatory for short options too.\n\
  62 "), stdout);
  63       fputs (_("\
  64   -e, --echo                treat each ARG as an input line\n\
  65   -i, --input-range=LO-HI   treat each number LO through HI as an input line\n\
  66   -n, --head-lines=LINES    output at most LINES lines\n\
  67   -o, --output=FILE         write result to FILE instead of standard output\n\
  68       --random-source=FILE  get random bytes from FILE (default /dev/urandom)\n\
  69   -z, --zero-terminated     end lines with 0 byte, not newline\n\
  70 "), stdout);
  71       fputs (HELP_OPTION_DESCRIPTION, stdout);
  72       fputs (VERSION_OPTION_DESCRIPTION, stdout);
  73       fputs (_("\
  74 \n\
  75 With no FILE, or when FILE is -, read standard input.\n\
  76 "), stdout);
  77       emit_bug_reporting_address ();
  78     }
  79
  80   exit (status);
  81 }
  82
  83 /* For long options that have no equivalent short option, use a
  84    non-character as a pseudo short option, starting with CHAR_MAX + 1.  */
  85 enum
  86 {
  87   RANDOM_SOURCE_OPTION = CHAR_MAX + 1
  88 };
  89
  90 static struct option const long_opts[] =
  91 {
  92   {"echo", no_argument, NULL, 'e'},
  93   {"input-range", required_argument, NULL, 'i'},
  94   {"head-count", required_argument, NULL, 'n'},
  95   {"output", required_argument, NULL, 'o'},
  96   {"random-source", required_argument, NULL, RANDOM_SOURCE_OPTION},
  97   {"zero-terminated", no_argument, NULL, 'z'},
  98   {GETOPT_HELP_OPTION_DECL},
  99   {GETOPT_VERSION_OPTION_DECL},
 100   {0, 0, 0, 0},
 101 };
 102
 103 static bool
 104 input_numbers_option_used (size_t lo_input, size_t hi_input)
 105 {
 106   return ! (lo_input == SIZE_MAX && hi_input == 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.  The current line
 132    ends in EOLBYTE, and is guaranteed to end before LINE + N.  */
 133
 134 static char *
 135 next_line (char *line, char eolbyte, size_t n)
 136 {
 137   char *p = memchr (line, eolbyte, n);
 138   return p + 1;
 139 }
 140
 141 /* Read data from file IN.  Input lines are delimited by EOLBYTE;
 142    silently append a trailing EOLBYTE if the file ends in some other
 143    byte.  Store a pointer to the resulting array of lines into *PLINE.
 144    Return the number of lines read.  Report an error and exit on
 145    failure.  */
 146
 147 static size_t
 148 read_input (FILE *in, char eolbyte, char ***pline)
 149 {
 150   char *p;
 151   char *buf = NULL;
 152   char *lim;
 153   size_t alloc = 0;
 154   size_t used = 0;
 155   size_t next_alloc = (1 << 13) + 1;
 156   size_t bytes_to_read;
 157   size_t nread;
 158   char **line;
 159   size_t i;
 160   size_t n_lines;
 161   int fread_errno;
 162   struct stat instat;
 163
 164   if (fstat (fileno (in), &instat) == 0 && S_ISREG (instat.st_mode))
 165     {
 166       off_t file_size = instat.st_size;
 167       off_t current_offset = ftello (in);
 168       if (0 <= current_offset)
 169         {
 170           off_t remaining_size =
 171             (current_offset < file_size ? file_size - current_offset : 0);
 172           if (SIZE_MAX - 2 < remaining_size)
 173             xalloc_die ();
 174           next_alloc = remaining_size + 2;
 175         }
 176     }
 177
 178   do
 179     {
 180       if (alloc <= used + 1)
 181         {
 182           if (alloc == SIZE_MAX)
 183             xalloc_die ();
 184           alloc = next_alloc;
 185           next_alloc = alloc * 2;
 186           if (next_alloc < alloc)
 187             next_alloc = SIZE_MAX;
 188           buf = xrealloc (buf, alloc);
 189         }
 190
 191       bytes_to_read = alloc - used - 1;
 192       nread = fread (buf + used, sizeof (char), bytes_to_read, in);
 193       used += nread;
 194     }
 195   while (nread == bytes_to_read);
 196
 197   fread_errno = errno;
 198
 199   if (used && buf[used - 1] != eolbyte)
 200     buf[used++] = eolbyte;
 201
 202   lim = buf + used;
 203
 204   n_lines = 0;
 205   for (p = buf; p < lim; p = next_line (p, eolbyte, lim - p))
 206     n_lines++;
 207
 208   *pline = line = xnmalloc (n_lines + 1, sizeof *line);
 209
 210   line[0] = p = buf;
 211   for (i = 1; i <= n_lines; i++)
 212     line[i] = p = next_line (p, eolbyte, lim - p);
 213
 214   errno = fread_errno;
 215   return n_lines;
 216 }
 217
 218 static int
 219 write_permuted_output (size_t n_lines, char * const *line, size_t lo_input,
 220                        size_t const *permutation)
 221 {
 222   size_t i;
 223
 224   if (line)
 225     for (i = 0; i < n_lines; i++)
 226       {
 227         char * const *p = line + permutation[i];
 228         size_t len = p[1] - p[0];
 229         if (fwrite (p[0], sizeof *p[0], len, stdout) != len)
 230           return -1;
 231       }
 232   else
 233     for (i = 0; i < n_lines; i++)
 234       {
 235         unsigned long int n = lo_input + permutation[i];
 236         if (printf ("%lu\n", n) < 0)
 237           return -1;
 238       }
 239
 240   return 0;
 241 }
 242
 243 int
 244 main (int argc, char **argv)
 245 {
 246   bool echo = false;
 247   size_t lo_input = SIZE_MAX;
 248   size_t hi_input = 0;
 249   size_t head_lines = SIZE_MAX;
 250   char const *outfile = NULL;
 251   char *random_source = NULL;
 252   char eolbyte = '\n';
 253   char **input_lines = NULL;
 254
 255   int optc;
 256   int n_operands;
 257   char **operand;
 258   size_t n_lines;
 259   char **line;
 260   struct randint_source *randint_source;
 261   size_t *permutation;
 262
 263   initialize_main (&argc, &argv);
 264   program_name = argv[0];
 265   setlocale (LC_ALL, "");
 266   bindtextdomain (PACKAGE, LOCALEDIR);
 267   textdomain (PACKAGE);
 268
 269   atexit (close_stdout);
 270
 271   while ((optc = getopt_long (argc, argv, "ei:n:o:z", long_opts, NULL)) != -1)
 272     switch (optc)
 273       {
 274       case 'e':
 275         echo = true;
 276         break;
 277
 278       case 'i':
 279         {
 280           unsigned long int argval = 0;
 281           char *p = strchr (optarg, '-');
 282           char const *hi_optarg = optarg;
 283           bool invalid = !p;
 284
 285           if (input_numbers_option_used (lo_input, hi_input))
 286             error (EXIT_FAILURE, 0, _("multiple -i options specified"));
 287
 288           if (p)
 289             {
 290               *p = '\0';
 291               invalid = ((xstrtoul (optarg, NULL, 10, &argval, NULL)
 292                           != LONGINT_OK)
 293                          || SIZE_MAX < argval);
 294               *p = '-';
 295               lo_input = argval;
 296               hi_optarg = p + 1;
 297             }
 298
 299           invalid |= ((xstrtoul (hi_optarg, NULL, 10, &argval, NULL)
 300                        != LONGINT_OK)
 301                       || SIZE_MAX < argval);
 302           hi_input = argval;
 303           n_lines = hi_input - lo_input + 1;
 304           invalid |= ((lo_input <= hi_input) == (n_lines == 0));
 305           if (invalid)
 306             error (EXIT_FAILURE, 0, _("invalid input range %s"),
 307                    quote (optarg));
 308         }
 309         break;
 310
 311       case 'n':
 312         {
 313           unsigned long int argval;
 314           strtol_error e = xstrtoul (optarg, NULL, 10, &argval, NULL);
 315
 316           if (e == LONGINT_OK)
 317             head_lines = MIN (head_lines, argval);
 318           else if (e != LONGINT_OVERFLOW)
 319             error (EXIT_FAILURE, 0, _("invalid line count %s"),
 320                    quote (optarg));
 321         }
 322         break;
 323
 324       case 'o':
 325         if (outfile && !STREQ (outfile, optarg))
 326           error (EXIT_FAILURE, 0, _("multiple output files specified"));
 327         outfile = optarg;
 328         break;
 329
 330       case RANDOM_SOURCE_OPTION:
 331         if (random_source && !STREQ (random_source, optarg))
 332           error (EXIT_FAILURE, 0, _("multiple random sources specified"));
 333         random_source = optarg;
 334         break;
 335
 336       case 'z':
 337         eolbyte = '\0';
 338         break;
 339
 340       case_GETOPT_HELP_CHAR;
 341       case_GETOPT_VERSION_CHAR (PROGRAM_NAME, AUTHORS);
 342       default:
 343         usage (EXIT_FAILURE);
 344       }
 345
 346   n_operands = argc - optind;
 347   operand = argv + optind;
 348
 349   if (echo)
 350     {
 351       if (input_numbers_option_used (lo_input, hi_input))
 352         error (EXIT_FAILURE, 0, _("cannot combine -e and -i options"));
 353       input_from_argv (operand, n_operands, eolbyte);
 354       n_lines = n_operands;
 355       line = operand;
 356     }
 357   else if (input_numbers_option_used (lo_input, hi_input))
 358     {
 359       if (n_operands)
 360         {
 361           error (0, 0, _("extra operand %s\n"), quote (operand[0]));
 362           usage (EXIT_FAILURE);
 363         }
 364       n_lines = hi_input - lo_input + 1;
 365       line = NULL;
 366     }
 367   else
 368     {
 369       switch (n_operands)
 370         {
 371         case 0:
 372           break;
 373
 374         case 1:
 375           if (! (STREQ (operand[0], "-") || freopen (operand[0], "r", stdin)))
 376             error (EXIT_FAILURE, errno, "%s", operand[0]);
 377           break;
 378
 379         default:
 380           error (0, 0, _("extra operand %s"), quote (operand[1]));
 381           usage (EXIT_FAILURE);
 382         }
 383
 384       n_lines = read_input (stdin, eolbyte, &input_lines);
 385       line = input_lines;
 386     }
 387
 388   head_lines = MIN (head_lines, n_lines);
 389
 390   randint_source = randint_all_new (random_source,
 391                                     randperm_bound (head_lines, n_lines));
 392   if (! randint_source)
 393     error (EXIT_FAILURE, errno, "%s", quotearg_colon (random_source));
 394
 395   /* Close stdin now, rather than earlier, so that randint_all_new
 396      doesn't have to worry about opening something other than
 397      stdin.  */
 398   if (! (echo || input_numbers_option_used (lo_input, hi_input))
 399       && (ferror (stdin) || fclose (stdin) != 0))
 400     error (EXIT_FAILURE, errno, _("read error"));
 401
 402   permutation = randperm_new (randint_source, head_lines, n_lines);
 403
 404   if (outfile && ! freopen (outfile, "w", stdout))
 405     error (EXIT_FAILURE, errno, "%s", quotearg_colon (outfile));
 406   if (write_permuted_output (head_lines, line, lo_input, permutation) != 0)
 407     error (EXIT_FAILURE, errno, _("write error"));
 408
 409 #ifdef lint
 410   free (permutation);
 411   randint_all_free (randint_source);
 412   if (input_lines)
 413     {
 414       free (input_lines[0]);
 415       free (input_lines);
 416     }
 417 #endif
 418
 419   return EXIT_SUCCESS;
 420 }