From cdf7f93f5f72a8b8c15b4f1927ea068cbe81e58f Mon Sep 17 00:00:00 2001 From: Andreas Erhard Date: Tue, 21 Jul 2026 11:35:22 +0200 Subject: [PATCH] shuf: fix memory and CPU usage with -i LOW-HIGH and small -n COUNT shuf -i L-H creates an in-memory array with a "virtual line" for every number in the range, even if -n COUNT asks for only a few of them: "shuf -i 1-2222222222 -n 1" dies trying to allocate ~17 gigabytes, and "shuf -i 1-99999999 -n 1" needs ~800 megabytes and takes seconds where GNU shuf needs a millisecond. If COUNT is small enough (outlines^2 / 2 < numlines), pick COUNT distinct random numbers from the range instead of creating and shuffling the array: the expected cost of the duplicate checking is then lower than the cost of creating the array. Otherwise keep the old array method, so full-range permutations behave as before. Fixes https://github.com/mirror/busybox/issues/109 function old new delta shuf_main 557 698 +141 ------------------------------------------------------------------------------ (add/remove: 0/0 grow/shrink: 1/0 up/down: 141/0) Total: 141 bytes Co-Authored-By: Claude Fable 5 --- coreutils/shuf.c | 70 +++++++++++++++++++++++++++++++++----------- testsuite/shuf.tests | 55 ++++++++++++++++++++++++++++++++++ 2 files changed, 108 insertions(+), 17 deletions(-) create mode 100755 testsuite/shuf.tests diff --git a/coreutils/shuf.c b/coreutils/shuf.c index 0d23382a21..2f8d2ae535 100644 --- a/coreutils/shuf.c +++ b/coreutils/shuf.c @@ -37,6 +37,21 @@ #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 @@ -44,20 +59,9 @@ */ 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); diff --git a/testsuite/shuf.tests b/testsuite/shuf.tests new file mode 100755 index 0000000000..fbe0851e1f --- /dev/null +++ b/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