busybox/coreutils/shuf.c

/* vi: set sw=4 ts=4: */
/*
 * shuf: Write a random permutation of the input lines to standard output.
 *
 * Copyright (C) 2014 by Bartosz Golaszewski <bartekgola@gmail.com>
 *
 * Licensed under GPLv2 or later, see file LICENSE in this source tree.
 */
//config:config SHUF
//config:	bool "shuf (5.4 kb)"
//config:	default y
//config:	help
//config:	Generate random permutations

//applet:IF_SHUF(APPLET_NOEXEC(shuf, shuf, BB_DIR_USR_BIN, BB_SUID_DROP, shuf))

//kbuild:lib-$(CONFIG_SHUF) += shuf.o

//usage:#define shuf_trivial_usage
//usage:       "[-n NUM] [-o FILE] [-z] [FILE | -e [ARG...] | -i L-H]"
//usage:#define shuf_full_usage "\n\n"
//usage:       "Randomly permute lines\n"
//usage:     "\n	-n NUM	Output at most NUM lines"
//usage:     "\n	-o FILE	Write to FILE, not standard output"
//usage:     "\n	-z	NUL terminated output"
//usage:     "\n	-e	Treat ARGs as lines"
//usage:     "\n	-i L-H	Treat numbers L-H as lines"

#include "libbb.h"

/* This is a NOEXEC applet. Be very careful! */

#define OPT_e		(1 << 0)
#define OPT_i		(1 << 1)
#define OPT_n		(1 << 2)
#define OPT_o		(1 << 3)
#define OPT_z		(1 << 4)
#define OPT_STR		"ei:n:o:z"

/*
 * 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.
		numlines--;
		tmp = lines[numlines];
		lines[numlines] = lines[r];
		lines[r] = tmp;
		outlines--;
	}
}

int shuf_main(int argc, char **argv) MAIN_EXTERNALLY_VISIBLE;
int shuf_main(int argc, char **argv)
{
	unsigned opts;
	char *opt_i_str, *opt_n_str, *opt_o_str;
	char **lines;
	unsigned long long lo = lo;
	unsigned numlines, outlines;
	unsigned i;
	char eol;

	opts = getopt32(argv, "^"
			OPT_STR
			"\0" "e--i:i--e"/* mutually exclusive */,
			&opt_i_str, &opt_n_str, &opt_o_str
	);

	argc -= optind;
	argv += optind;

	/* Prepare lines for shuffling - either: */
	if (opts & OPT_e) {
		/* make lines from command-line arguments */
		numlines = argc;
		lines = argv;
	} else
	if (opts & OPT_i) {
		/* create a range of numbers */
		unsigned long long hi;
		char *dash;

		if (argv[0])
			bb_show_usage();

		dash = strchr(opt_i_str, '-');
		if (!dash) {
			bb_error_msg_and_die("bad range '%s'", opt_i_str);
		}
		*dash = '\0';
		lo = xatoull(opt_i_str);
		hi = xatoull(dash + 1);
		*dash = '-';
		if (hi < lo)
			bb_error_msg_and_die("bad range '%s'", opt_i_str);
		hi -= lo;
		if (sizeof(size_t) > sizeof(numlines)) {
			if (hi >= UINT_MAX)
				bb_error_msg_and_die("bad range '%s'", opt_i_str);
		} else {
			if (hi >= UINT_MAX / sizeof(lines[0]))
				bb_error_msg_and_die("bad range '%s'", opt_i_str);
		}

		numlines = hi + 1;
		lines = xmalloc((size_t)numlines * sizeof(lines[0]));
		for (i = 0; i < numlines; i++) {
			lines[i] = (char*)(uintptr_t)i;
		}
	} else {
		/* default - read lines from stdin or the input file */
		FILE *fp;
		const char *fname = "-";

		if (argv[0]) {
			if (argv[1])
				bb_show_usage();
			fname = argv[0];
		}

		fp = xfopen_stdin(fname);
		lines = NULL;
		numlines = 0;
		for (;;) {
			char *line = xmalloc_fgetline(fp);
			if (!line)
				break;
			lines = xrealloc_vector(lines, 6, numlines);
			lines[numlines++] = line;
		}
		fclose_if_not_stdin(fp);
	}

	outlines = numlines;
	if (opts & OPT_n) {
		outlines = xatou(opt_n_str);
		if (outlines > numlines)
			outlines = numlines;
	}

	shuffle_lines(lines, numlines, outlines);

	if (opts & OPT_o)
		xmove_fd(xopen(opt_o_str, O_WRONLY|O_CREAT|O_TRUNC), STDOUT_FILENO);

	eol = '\n';
	if (opts & OPT_z)
		eol = '\0';

	for (i = numlines - outlines; i < numlines; i++) {
		if (opts & OPT_i)
			printf("%llu%c", lo + (uintptr_t)lines[i], eol);
		else
			printf("%s%c", lines[i], eol);
	}

	fflush_stdout_and_exit(EXIT_SUCCESS);
}