src/shuf.c

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