coreutils/shuf.c

   1 /* vi: set sw=4 ts=4: */
   2 /*
   3  * shuf: Write a random permutation of the input lines to standard output.
   4  *
   5  * Copyright (C) 2014 by Bartosz Golaszewski <bartekgola@gmail.com>
   6  *
   7  * Licensed under GPLv2 or later, see file LICENSE in this source tree.
   8  */
   9 //config:config SHUF
  10 //config:       bool "shuf (5.4 kb)"
  11 //config:       default y
  12 //config:       help
  13 //config:       Generate random permutations
  14
  15 //applet:IF_SHUF(APPLET_NOEXEC(shuf, shuf, BB_DIR_USR_BIN, BB_SUID_DROP, shuf))
  16
  17 //kbuild:lib-$(CONFIG_SHUF) += shuf.o
  18
  19 //usage:#define shuf_trivial_usage
  20 //usage:       "[-e|-i L-H] [-n NUM] [-o FILE] [-z] [FILE|ARG...]"
  21 //usage:#define shuf_full_usage "\n\n"
  22 //usage:       "Randomly permute lines\n"
  23 //usage:     "\n        -e      Treat ARGs as lines"
  24 //usage:     "\n        -i L-H  Treat numbers L-H as lines"
  25 //usage:     "\n        -n NUM  Output at most NUM lines"
  26 //usage:     "\n        -o FILE Write to FILE, not standard output"
  27 //usage:     "\n        -z      End lines with zero byte, not newline"
  28
  29 #include "libbb.h"
  30
  31 /* This is a NOEXEC applet. Be very careful! */
  32
  33 #define OPT_e           (1 << 0)
  34 #define OPT_i           (1 << 1)
  35 #define OPT_n           (1 << 2)
  36 #define OPT_o           (1 << 3)
  37 #define OPT_z           (1 << 4)
  38 #define OPT_STR         "ei:n:o:z"
  39
  40 /*
  41  * Use the Fisher-Yates shuffle algorithm on an array of lines.
  42  */
  43 static void shuffle_lines(char **lines, unsigned numlines)
  44 {
  45         unsigned i;
  46         unsigned r;
  47         char *tmp;
  48
  49         srand(monotonic_us());
  50
  51         for (i = numlines-1; i > 0; i--) {
  52                 r = rand();
  53                 /* RAND_MAX can be as small as 32767 */
  54                 if (i > RAND_MAX)
  55                         r ^= rand() << 15;
  56                 r %= i + 1;
  57                 tmp = lines[i];
  58                 lines[i] = lines[r];
  59                 lines[r] = tmp;
  60         }
  61 }
  62
  63 int shuf_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
  64 int shuf_main(int argc, char **argv)
  65 {
  66         unsigned opts;
  67         char *opt_i_str, *opt_n_str, *opt_o_str;
  68         unsigned i;
  69         char **lines;
  70         unsigned numlines;
  71         char eol;
  72
  73         opts = getopt32(argv, "^"
  74                         OPT_STR
  75                         "\0" "e--i:i--e"/* mutually exclusive */,
  76                         &opt_i_str, &opt_n_str, &opt_o_str
  77         );
  78
  79         argc -= optind;
  80         argv += optind;
  81
  82         /* Prepare lines for shuffling - either: */
  83         if (opts & OPT_e) {
  84                 /* make lines from command-line arguments */
  85                 numlines = argc;
  86                 lines = argv;
  87         } else
  88         if (opts & OPT_i) {
  89                 /* create a range of numbers */
  90                 char *dash;
  91                 unsigned lo, hi;
  92
  93                 dash = strchr(opt_i_str, '-');
  94                 if (!dash) {
  95                         bb_error_msg_and_die("bad range '%s'", opt_i_str);
  96                 }
  97                 *dash = '\0';
  98                 lo = xatou(opt_i_str);
  99                 hi = xatou(dash + 1);
 100                 *dash = '-';
 101                 if (hi < lo) {
 102                         bb_error_msg_and_die("bad range '%s'", opt_i_str);
 103                 }
 104
 105                 numlines = (hi+1) - lo;
 106                 lines = xmalloc(numlines * sizeof(lines[0]));
 107                 for (i = 0; i < numlines; i++) {
 108                         lines[i] = (char*)(uintptr_t)lo;
 109                         lo++;
 110                 }
 111         } else {
 112                 /* default - read lines from stdin or the input file */
 113                 FILE *fp;
 114
 115                 if (argc > 1)
 116                         bb_show_usage();
 117
 118                 fp = xfopen_stdin(argv[0] ? argv[0] : "-");
 119                 lines = NULL;
 120                 numlines = 0;
 121                 for (;;) {
 122                         char *line = xmalloc_fgetline(fp);
 123                         if (!line)
 124                                 break;
 125                         lines = xrealloc_vector(lines, 6, numlines);
 126                         lines[numlines++] = line;
 127                 }
 128                 fclose_if_not_stdin(fp);
 129         }
 130
 131         if (numlines != 0)
 132                 shuffle_lines(lines, numlines);
 133
 134         if (opts & OPT_o)
 135                 xmove_fd(xopen(opt_o_str, O_WRONLY|O_CREAT|O_TRUNC), STDOUT_FILENO);
 136
 137         if (opts & OPT_n) {
 138                 unsigned maxlines;
 139                 maxlines = xatou(opt_n_str);
 140                 if (numlines > maxlines)
 141                         numlines = maxlines;
 142         }
 143
 144         eol = '\n';
 145         if (opts & OPT_z)
 146                 eol = '\0';
 147
 148         for (i = 0; i < numlines; i++) {
 149                 if (opts & OPT_i)
 150                         printf("%u%c", (unsigned)(uintptr_t)lines[i], eol);
 151                 else
 152                         printf("%s%c", lines[i], eol);
 153         }
 154
 155         fflush_stdout_and_exit(EXIT_SUCCESS);
 156 }