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