LXR busybox/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 */
  43static 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
  63int shuf_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
  64int 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}
 157