shuf: fix memory and CPU usage with -i LOW-HIGH and small -n COUNT

coreutils/shuf.c

@@ -37,27 +37,31 @@
 #define OPT_z		(1 << 4)
 #define OPT_STR		"ei:n:o:z"
 
+/* Returns a random number in [0, n) */
+static unsigned random_below(unsigned n)
+{
+	unsigned r = rand();
+	/* RAND_MAX can be as small as 32767 */
+	if (n > RAND_MAX)
+		r ^= rand() << 15;
+	return r % n;
+//TODO: the above method is seriously non-uniform when n is very large.
+//For example, with n of          0xf0000000,
+//values of (r % n) in [0, 0x0fffffff] range
+//are more likely: e.g. r=1 and r=0xf0000001 both map to 1,
+//whereas only one value, r=0xefffffff, maps to 0xefffffff.
+}
+
 /*
  * Use the Fisher-Yates shuffle algorithm on an array of lines.
  * If the required number of output lines is less than the total
  * we can stop shuffling early.
  */
 static void shuffle_lines(char **lines, unsigned numlines, unsigned outlines)
 {
-	srand(monotonic_us());
-
 	while (outlines != 0) {
 		char *tmp;
-		unsigned r = rand();
-		/* RAND_MAX can be as small as 32767 */
-		if (numlines > RAND_MAX)
-			r ^= rand() << 15;
-		r %= numlines;
-//TODO: the above method is seriously non-uniform when numlines is very large.
-//For example, with numlines of   0xf0000000,
-//values of (r % numlines) in [0, 0x0fffffff] range
-//are more likely: e.g. r=1 and r=0xf0000001 both map to 1,
-//whereas only one value, r=0xefffffff, maps to 0xefffffff.
+		unsigned r = random_below(numlines);
 		numlines--;
 		tmp = lines[numlines];
 		lines[numlines] = lines[r];
@@ -161,10 +165,8 @@ int shuf_main(int argc, char **argv)
 		}
 
 		numlines = hi + 1;
-		lines = xmalloc((size_t)numlines * sizeof(lines[0]));
-		for (i = 0; i < numlines; i++) {
-			lines[i] = (char*)(uintptr_t)i;
-		}
+		/* lines[] is allocated below, when outlines is known */
+		lines = NULL;
 	} else {
 		/* default - read lines from stdin or the input file */
 		FILE *fp;
@@ -196,7 +198,41 @@ int shuf_main(int argc, char **argv)
 			outlines = numlines;
 	}
 
-	shuffle_lines(lines, numlines, outlines);
+	srand(monotonic_us());
+
+	if ((opts & OPT_i)
+	 && (unsigned long long)outlines * outlines / 2 < numlines
+	) {
+		/* Do not create a "virtual line" for each number in the range:
+		 * a large range with a small -n COUNT would use lots of memory
+		 * and time just to output a few numbers (and worse,
+		 * e.g. "shuf -i 1-2222222222 -n 1" would fail to allocate
+		 * ~17 gigabytes). Instead, pick COUNT distinct random numbers
+		 * from the range. Expected number of comparisons below
+		 * is less than outlines^2 / 2 < numlines - cheaper than
+		 * creating and shuffling the full array.
+		 */
+		lines = xmalloc((size_t)outlines * sizeof(lines[0]));
+		for (i = 0; i < outlines; i++) {
+			unsigned j;
+			uintptr_t v;
+ again:
+			v = random_below(numlines);
+			for (j = 0; j < i; j++)
+				if ((uintptr_t)lines[j] == v)
+					goto again; /* duplicate, pick another */
+			lines[i] = (char*)v;
+		}
+		numlines = outlines;
+	} else {
+		if (opts & OPT_i) {
+			lines = xmalloc((size_t)numlines * sizeof(lines[0]));
+			for (i = 0; i < numlines; i++) {
+				lines[i] = (char*)(uintptr_t)i;
+			}
+		}
+		shuffle_lines(lines, numlines, outlines);
+	}
 
 	if (opts & OPT_o)
 		xmove_fd(xopen(opt_o_str, O_WRONLY|O_CREAT|O_TRUNC), STDOUT_FILENO);

testsuite/shuf.tests

@@ -0,0 +1,55 @@
+#!/bin/sh
+
+# Licensed under GPLv2, see file LICENSE in this source tree.
+
+. ./testing.sh
+
+# testing "test name" "commands" "expected result" "file input" "stdin"
+
+testing "shuf -i 1-10 prints each number once" \
+	"shuf -i 1-10 | sort -n" \
+	"1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n" "" ""
+
+testing "shuf -i 1-10 -n 10 prints each number once" \
+	"shuf -i 1-10 -n 10 | sort -n" \
+	"1\n2\n3\n4\n5\n6\n7\n8\n9\n10\n" "" ""
+
+testing "shuf -i L-L prints the single number" \
+	"shuf -i 5-5" \
+	"5\n" "" ""
+
+testing "shuf -i with -n larger than range" \
+	"shuf -i 5-5 -n 99" \
+	"5\n" "" ""
+
+# Must not try to allocate an array for the whole range
+# (~17 gigabytes here): www.github.com/mirror/busybox/issues/109
+testing "shuf -i with huge range prints distinct numbers" \
+	"shuf -i 1-2222222222 -n 4 | sort -un | wc -l" \
+	"4\n" "" ""
+
+testing "shuf -i with huge range prints numbers within it" \
+	"shuf -i 2000000000-2222222222 -n 3 | awk '\$1 < 2000000000 || \$1 > 2222222222'" \
+	"" "" ""
+
+testing "shuf -i with huge range and -n 0 prints nothing" \
+	"shuf -i 1-2222222222 -n 0" \
+	"" "" ""
+
+testing "shuf -e prints all args" \
+	"shuf -e a b c | sort" \
+	"a\nb\nc\n" "" ""
+
+testing "shuf prints all stdin lines" \
+	"shuf | sort" \
+	"x\ny\nz\n" "" "x\ny\nz\n"
+
+testing "shuf -z ends lines with NUL" \
+	"shuf -z -i 7-7 | tr '\\0' N" \
+	"7N" "" ""
+
+testing "shuf rejects reversed range" \
+	"shuf -i 5-1 2>&1; echo \$?" \
+	"shuf: bad range '5-1'\n1\n" "" ""
+
+exit $FAILCOUNT