shuf: add --repetition to support repetition in output

main(): Process new option.  Replace input_numbers_option_used()
with a local variable.  Re-organize argument processing.
usage(): Describe the new option.
(write_random_numbers): A new function to generate a
permutation of the specified input range with repetition.
(write_random_lines): Likewise for stdin and --echo.
(write_permuted_numbers):  New function refactored from
write_permuted_output().
(write_permuted_lines): Likewise.
* tests/misc/shuf.sh: Add tests for --repetitions option.
* doc/coreutils.texi: Mention --repetitions, add examples.
* TODO: Mention an optimization to avoid needing to
read all of the input into memory with --repetitions.
* NEWS: Mention new shuf option.

4 files changed, 236 insertions, 56 deletions

diff --git a/NEWS b/NEWS
index 53c9f4aef..5abaf81f1 100644
--- a/NEWS
+++ b/NEWS
@@ -46,6 +46,9 @@ GNU coreutils NEWS                                    -*- outline -*-
   csplit accepts a new option: --suppressed-matched, to elide the lines
   used to identify the split points.
 
+  shuf accepts a new option: --repetitions (-r), to allow repetitions
+  of input items in the permuted output.
+
 ** Changes in behavior
 
   stdbuf now requires at least one buffering mode option to be specified,
diff --git a/doc/coreutils.texi b/doc/coreutils.texi
index b3233f602..ca10a16ff 100644
--- a/doc/coreutils.texi
+++ b/doc/coreutils.texi
@@ -4945,6 +4945,16 @@ commands like @code{shuf -o F <F} and @code{cat F | shuf -o F}.
 Use @var{file} as a source of random data used to determine which
 permutation to generate.  @xref{Random sources}.
 
+@item -r
+@itemx --repetitions
+@opindex -r
+@opindex --repetitions
+@cindex allowing repetitions in output values
+Changes the default behaviour of @command{shuf}, allowing repetition in
+output values (in which case, @option{--head-count} can be larger
+than the number of input values).  If @option{--head-count} is not
+specified, output a single random value.
+
 @zeroTerminatedOption
 
 @end table
@@ -5004,6 +5014,33 @@ general, if there are @var{n} input lines, there are @var{n}! (i.e.,
 @var{n} factorial, or @var{n} * (@var{n} - 1) * @dots{} * 1) possible
 output permutations.
 
+@noindent
+To output 50 random numbers between 0 and 9, use:
+
+@example
+shuf --repetitions --input-range 0-9 --head-count 50
+@end example
+
+@noindent
+or (using short options):
+
+@example
+shuf -r -i0-9 -n50
+@end example
+
+@noindent
+To simulate 100 coin flips, use:
+
+@example
+shuf -r -n100 -e Head Tail
+@end example
+
+@noindent
+or
+@example
+printf '%s\n' Head Tail | shuf -r -n100
+@end example
+
 @exitstatus
 
 
diff --git a/src/shuf.c b/src/shuf.c
index 0fabb0bfe..a8c1d8530 100644
--- a/src/shuf.c
+++ b/src/shuf.c
@@ -76,6 +76,8 @@ Write a random permutation of the input lines to standard output.\n\
   -n, --head-count=COUNT    output at most COUNT lines\n\
   -o, --output=FILE         write result to FILE instead of standard output\n\
       --random-source=FILE  get random bytes from FILE\n\
+  -r, --repetitions         output COUNT items, allowing repetition.\n\
+                              -n 1 is implied if not specified.\n\
   -z, --zero-terminated     end lines with 0 byte, not newline\n\
 "), stdout);
       fputs (HELP_OPTION_DESCRIPTION, stdout);
@@ -104,18 +106,13 @@ static struct option const long_opts[] =
   {"head-count", required_argument, NULL, 'n'},
   {"output", required_argument, NULL, 'o'},
   {"random-source", required_argument, NULL, RANDOM_SOURCE_OPTION},
+  {"repetitions", no_argument, NULL, 'r'},
   {"zero-terminated", no_argument, NULL, 'z'},
   {GETOPT_HELP_OPTION_DECL},
   {GETOPT_VERSION_OPTION_DECL},
   {0, 0, 0, 0},
 };
 
-static bool
-input_numbers_option_used (size_t lo_input, size_t hi_input)
-{
-  return ! (lo_input == SIZE_MAX && hi_input == 0);
-}
-
 static void
 input_from_argv (char **operand, int n_operands, char eolbyte)
 {
@@ -303,27 +300,81 @@ read_input (FILE *in, char eolbyte, char ***pline)
   return n_lines;
 }
 
+/* output 'n_lines' to stdout from 'line' array,
+   chosen by the indices in 'permutation'.
+   'permutation' and 'line' must have at least 'n_lines' elements.
+   strings in 'line' must include the line-terminator character.  */
 static int
-write_permuted_output (size_t n_lines, char *const *line, size_t lo_input,
-                       size_t const *permutation, char eolbyte)
+write_permuted_lines (size_t n_lines, char *const *line,
+                      size_t const *permutation)
 {
   size_t i;
 
-  if (line)
-    for (i = 0; i < n_lines; i++)
-      {
-        char *const *p = line + permutation[i];
-        size_t len = p[1] - p[0];
-        if (fwrite (p[0], sizeof *p[0], len, stdout) != len)
-          return -1;
-      }
-  else
-    for (i = 0; i < n_lines; i++)
-      {
-        unsigned long int n = lo_input + permutation[i];
-        if (printf ("%lu%c", n, eolbyte) < 0)
-          return -1;
-      }
+  for (i = 0; i < n_lines; i++)
+    {
+      char *const *p = line + permutation[i];
+      size_t len = p[1] - p[0];
+      if (fwrite (p[0], sizeof *p[0], len, stdout) != len)
+        return -1;
+    }
+
+  return 0;
+}
+
+/* output 'n_lines' of numbers to stdout, from 'permutation' array.
+   'permutation' must have at least 'n_lines' elements.  */
+static int
+write_permuted_numbers (size_t n_lines, size_t lo_input,
+                        size_t const *permutation, char eolbyte)
+{
+  size_t i;
+
+  for (i = 0; i < n_lines; i++)
+    {
+      unsigned long int n = lo_input + permutation[i];
+      if (printf ("%lu%c", n, eolbyte) < 0)
+        return -1;
+    }
+
+  return 0;
+}
+
+/* output 'count' numbers to stdout, chosen randomly from range
+   lo_input to hi_input.  */
+static int
+write_random_numbers (struct randint_source *s, size_t count,
+                      size_t lo_input, size_t hi_input, char eolbyte)
+{
+  size_t i;
+  const randint range = hi_input - lo_input + 1;
+
+  for (i = 0; i < count; i++)
+    {
+      randint j = lo_input + randint_choose (s, range);
+      if (printf ("%lu%c", j, eolbyte) < 0)
+        return -1;
+    }
+
+  return 0;
+}
+
+/* output 'count' lines to stdout from 'lines' array.
+   'lines' must have at least 'n_lines' element in it.
+   strings in 'line' must include the line-terminator character.  */
+static int
+write_random_lines (struct randint_source *s, size_t count,
+                    char *const *lines, size_t n_lines)
+{
+  size_t i;
+
+  for (i = 0; i < count; i++)
+    {
+      const randint j = randint_choose (s, n_lines);
+      char *const *p = lines + j;
+      size_t len = p[1] - p[0];
+      if (fwrite (p[0], sizeof *p[0], len, stdout) != len)
+        return -1;
+    }
 
   return 0;
 }
@@ -332,6 +383,7 @@ int
 main (int argc, char **argv)
 {
   bool echo = false;
+  bool input_range = false;
   size_t lo_input = SIZE_MAX;
   size_t hi_input = 0;
   size_t head_lines = SIZE_MAX;
@@ -340,6 +392,7 @@ main (int argc, char **argv)
   char eolbyte = '\n';
   char **input_lines = NULL;
   bool use_reservoir_sampling = false;
+  bool repetition = false;
 
   int optc;
   int n_operands;
@@ -348,7 +401,7 @@ main (int argc, char **argv)
   char **line = NULL;
   struct linebuffer *reservoir = NULL;
   struct randint_source *randint_source;
-  size_t *permutation;
+  size_t *permutation = NULL;
   int i;
 
   initialize_main (&argc, &argv);
@@ -359,7 +412,7 @@ main (int argc, char **argv)
 
   atexit (close_stdout);
 
-  while ((optc = getopt_long (argc, argv, "ei:n:o:z", long_opts, NULL)) != -1)
+  while ((optc = getopt_long (argc, argv, "ei:n:o:rz", long_opts, NULL)) != -1)
     switch (optc)
       {
       case 'e':
@@ -373,8 +426,9 @@ main (int argc, char **argv)
           char const *hi_optarg = optarg;
           bool invalid = !p;
 
-          if (input_numbers_option_used (lo_input, hi_input))
+          if (input_range)
             error (EXIT_FAILURE, 0, _("multiple -i options specified"));
+          input_range = true;
 
           if (p)
             {
@@ -424,6 +478,10 @@ main (int argc, char **argv)
         random_source = optarg;
         break;
 
+      case 'r':
+        repetition = true;
+        break;
+
       case 'z':
         eolbyte = '\0';
         break;
@@ -437,46 +495,43 @@ main (int argc, char **argv)
   n_operands = argc - optind;
   operand = argv + optind;
 
+  /* Check invalid usage */
+  if (echo && input_range)
+    {
+      error (0, 0, _("cannot combine -e and -i options"));
+      usage (EXIT_FAILURE);
+    }
+  if ((n_operands>0 && input_range)
+      || (!echo && !input_range && n_operands>=2))
+    {
+      error (0, 0, _("extra operand %s"), quote (operand[1]));
+      usage (EXIT_FAILURE);
+    }
+
+  /* Prepare input */
   if (echo)
     {
-      if (input_numbers_option_used (lo_input, hi_input))
-        error (EXIT_FAILURE, 0, _("cannot combine -e and -i options"));
       input_from_argv (operand, n_operands, eolbyte);
       n_lines = n_operands;
       line = operand;
     }
-  else if (input_numbers_option_used (lo_input, hi_input))
+  else if (input_range)
     {
-      if (n_operands)
-        {
-          error (0, 0, _("extra operand %s"), quote (operand[0]));
-          usage (EXIT_FAILURE);
-        }
       n_lines = hi_input - lo_input + 1;
       line = NULL;
     }
   else
     {
-      switch (n_operands)
-        {
-        case 0:
-          break;
-
-        case 1:
+      /* Input file specified, re-open it as STDIN */
+      if (n_operands==1)
           if (! (STREQ (operand[0], "-") || ! head_lines
                  || freopen (operand[0], "r", stdin)))
             error (EXIT_FAILURE, errno, "%s", operand[0]);
-          break;
-
-        default:
-          error (0, 0, _("extra operand %s"), quote (operand[1]));
-          usage (EXIT_FAILURE);
-        }
 
       fadvise (stdin, FADVISE_SEQUENTIAL);
 
-      if (head_lines != SIZE_MAX && (! head_lines
-                                     || input_size () > RESERVOIR_MIN_INPUT))
+      if (! repetition && head_lines != SIZE_MAX
+          && (! head_lines || input_size () > RESERVOIR_MIN_INPUT))
         {
           use_reservoir_sampling = true;
           n_lines = SIZE_MAX;   /* unknown number of input lines, for now.  */
@@ -488,10 +543,17 @@ main (int argc, char **argv)
         }
     }
 
-  head_lines = MIN (head_lines, n_lines);
+  /* When generating random numbers with repetitions,
+     the default count is one, unless specified by the user.  */
+  if (repetition && head_lines == SIZE_MAX)
+    head_lines = 1 ;
+
+  if (! repetition)
+    head_lines = MIN (head_lines, n_lines);
 
   randint_source = randint_all_new (random_source,
-                                    use_reservoir_sampling ? SIZE_MAX :
+                                    (use_reservoir_sampling || repetition)?
+                                    SIZE_MAX:
                                     randperm_bound (head_lines, n_lines));
   if (! randint_source)
     error (EXIT_FAILURE, errno, "%s", quotearg_colon (random_source));
@@ -508,20 +570,36 @@ main (int argc, char **argv)
   /* Close stdin now, rather than earlier, so that randint_all_new
      doesn't have to worry about opening something other than
      stdin.  */
-  if (! (echo || input_numbers_option_used (lo_input, hi_input))
+  if (! (echo || input_range)
       && (fclose (stdin) != 0))
     error (EXIT_FAILURE, errno, _("read error"));
 
-  permutation = randperm_new (randint_source, head_lines, n_lines);
+  if (!repetition)
+    permutation = randperm_new (randint_source, head_lines, n_lines);
 
   if (outfile && ! freopen (outfile, "w", stdout))
     error (EXIT_FAILURE, errno, "%s", quotearg_colon (outfile));
 
-  if (use_reservoir_sampling)
-    i = write_permuted_output_reservoir (n_lines, reservoir, permutation);
+  /* Generate output according to requested method */
+  if (repetition)
+    {
+      if (input_range)
+        i = write_random_numbers (randint_source, head_lines,
+                                  lo_input, hi_input, eolbyte);
+      else
+        i = write_random_lines (randint_source, head_lines, line, n_lines);
+    }
   else
-    i = write_permuted_output (head_lines, line, lo_input,
-                               permutation, eolbyte);
+    {
+      if (use_reservoir_sampling)
+        i = write_permuted_output_reservoir (n_lines, reservoir, permutation);
+      else if (input_range)
+        i = write_permuted_numbers (head_lines, lo_input,
+                                    permutation, eolbyte);
+      else
+        i = write_permuted_lines (head_lines, line, permutation);
+    }
+
   if (i != 0)
     error (EXIT_FAILURE, errno, _("write error"));
 
diff --git a/tests/misc/shuf.sh b/tests/misc/shuf.sh
index 492fd4186..a25a6f8ad 100755
--- a/tests/misc/shuf.sh
+++ b/tests/misc/shuf.sh
@@ -94,4 +94,66 @@ shuf -i0-9 -o A -o B &&
 shuf -i0-9 --random-source A --random-source B &&
   { fail=1; echo "shuf did not detect multiple --random-source usage.">&2 ; }
 
+# Test --repetition option
+
+# --repetition without count should return one line
+shuf --rep -i0-10 > exp || framework_failure_
+c=$(wc -l < exp) || framework_failure_
+test "$c" -eq 1 || { fail=1; echo "--repetition default count is not 1">&2 ; }
+
+# --repetition can output more values than the input range
+shuf --rep -i0-9 -n1000 > exp || framework_failure_
+c=$(wc -l < exp) || framework_failure_
+test "$c" -eq 1000 || { fail=1; echo "--repetition with --count failed">&2 ; }
+
+# Check output values (this is not bullet-proof, but drawing 1000 values
+# between 0 and 9 should produce all values, unless there's a bug in shuf
+# or a very poor random source, or extremely bad luck)
+c=$(sort -nu exp | paste -s -d ' ') || framework_failure_
+test "$c" = "0 1 2 3 4 5 6 7 8 9" ||
+  { fail=1; echo "--repetition produced bad output">&2 ; }
+
+# check --repetition with non-zero low value
+shuf --rep -i222-233 -n2000 > exp || framework_failure_
+c=$(cat exp | sort -nu | paste -s -d ' ') || framework_failure_
+test "$c" = "222 223 224 225 226 227 228 229 230 231 232 233" ||
+ { fail=1; echo "--repetition produced bad output with non-zero low">&2 ; }
+
+# --repetition,-i,count=0 should not fail and produce no output
+shuf --rep -i0-9 -n0 > exp || framework_failure_
+# file size should be zero (no output from shuf)
+test \! -s exp ||
+  { fail=1; echo "--repetition,-i0-9,-n0 produced bad output">&2 ; }
+
+# --repetition with -e, without count, should return one line
+shuf --rep -e A B C D > exp || framework_failure_
+c=$(cat exp | wc -l) || framework_failure_
+test "$c" -eq 1 ||
+  { fail=1; echo "--repetition,-e default count is not 1">&2 ; }
+
+# --repetition with STDIN, without count, should return one line
+printf "A\nB\nC\nD\nE\n" | shuf --rep > exp || framework_failure_
+c=$(wc -l < exp) || framework_failure_
+test "$c" -eq 1 ||
+  { fail=1; echo "--repetition,STDIN default count is not 1">&2 ; }
+
+# --repetition with STDIN,count - can return move values than input lines
+printf "A\nB\nC\nD\nE\n" | shuf --rep -n2000 > exp || framework_failure_
+c=$(wc -l < exp) || framework_failure_
+test "$c" -eq 2000 ||
+  { fail=1; echo "--repetition,STDIN,count failed">&2 ; }
+
+# Check output values (this is not bullet-proof, but drawing 2000 values
+# between A and E should produce all values, unless there's a bug in shuf
+# or a very poor random source, or extremely bad luck)
+c=$(sort -u exp | paste -s -d ' ') || framework_failure_
+test "$c" = "A B C D E" ||
+  { fail=1; echo "--repetition,STDIN,count produced bad output">&2 ; }
+
+# --repetition,stdin,count=0 should not fail and produce no output
+printf "A\nB\nC\nD\nE\n" | shuf --rep -n0 > exp || framework_failure_
+# file size should be zero (no output from shuf)
+test \! -s exp ||
+  { fail=1; echo "--repetition,STDIN,-n0 produced bad output">&2 ; }
+
 Exit $fail