utils-std

shuf.c (4636B)

// utils-std: Collection of commonly available Unix tools
// SPDX-FileCopyrightText: 2017 Haelwenn (lanodan) Monnier <contact+utils@hacktivis.me>
// SPDX-License-Identifier: MPL-2.0
#define _POSIX_C_SOURCE 200809L
#include "../config.h"
#include "../libutils/getopt_nolong.h"
#include <errno.h>
#include <stdbool.h>
#include <stdio.h>  // getdelim, fprintf
#include <stdlib.h> // free, malloc, srand, rand, exit, strtoul
#include <string.h> // strerror, memcpy
#include <time.h>   // time
#include <unistd.h> // getopt
#ifdef HAS_GETOPT_LONG
#include <getopt.h>
#endif
// Not a full shuffle, if there is more than 512 lines then last lines are never going to be printed first.
// But this allows bounded memory usage.
// /!\ Make sure to modify the manpage as well if this gets changed /!\
// FIXME: handle newline-less lines
const char *argv0 = "shuf";
#define LINES_LEN 512
static char *lines[LINES_LEN];
static char delim = '\n';
char *line = NULL;
size_t line_len = 0;
unsigned long wrote = 0;
unsigned long write_limit = 0;
static int
shuf(FILE *in, const char *fname)
{
	for(int ln = 0;; ln++)
	{
		errno = 0;
		ssize_t nread = getdelim(&line, &line_len, delim, in);
		if(errno != 0)
		{
			fprintf(stderr,
			        "%s: error: Failed reading line %d from file \"%s\": %s\n",
			        argv0,
			        ln,
			        fname,
			        strerror(errno));
			return 1;
		}
		if(nread <= 0) return 0;
		errno = 0;
		char *dup = malloc(nread + 1);
		if(!dup)
		{
			fprintf(
			    stderr,
			    "%s: error: Failed to allocate %zd bytes of memory for line %d from file \"%s\": %s\n",
			    argv0,
			    nread,
			    ln,
			    fname,
			    strerror(errno));
			return 1;
		}
		memcpy(dup, line, nread);
		dup[nread] = '\0';
		int p = rand() % LINES_LEN;
		if(lines[p] != NULL)
		{
			fputs(lines[p], stdout);
			free(lines[p]);
			lines[p] = NULL;
			wrote++;
			if(write_limit != 0 && write_limit <= wrote) exit(0);
		}
		lines[p] = dup;
	}
}
static void
usage(void)
{
	fputs("\
Usage: shuf [-z] [-n num] [files...]\n\
       shuf -e [-z] [-n num] [string...]\n\
",
	      stderr);
}
int
main(int argc, char *argv[])
{
	bool e_flag = false;
	srand((int)time(NULL));
#ifdef HAS_GETOPT_LONG
	// Strictly for GNUisms compatibility so no long-only options
	// clang-format off
	static struct option opts[] = {
		{"echo", no_argument, NULL, 'e'},
		{"head-count", required_argument, NULL, 'n'},
		{"zero-terminated", no_argument, NULL, 'z'},
		{0, 0, 0, 0},
	};
	// clang-format on
	// Need + as first character to get POSIX-style option parsing
	for(int c = -1; (c = getopt_long(argc, argv, "+:en:z", opts, NULL)) != -1;)
#else
	for(int c = -1; (c = getopt_nolong(argc, argv, ":en:z")) != -1;)
#endif
	{
		char *endptr = NULL;
		switch(c)
		{
		case 'e':
			e_flag = true;
			break;
		case 'n':
			write_limit = strtoul(optarg, &endptr, 0);
			if(errno != 0)
			{
				fprintf(stderr,
				        "%s: error: Failed parsing number for `-n %s`: %s\n",
				        argv0,
				        optarg,
				        strerror(errno));
				return 1;
			}
			if(endptr != NULL && *endptr != 0)
			{
				fprintf(stderr,
				        "%s: error: Found extraneous characters while parsing `-n %s` as a number: %s\n",
				        argv0,
				        optarg,
				        endptr);
				return 1;
			}
			break;
		case 'z':
			delim = '\0';
			break;
		case ':':
			fprintf(stderr, "%s: error: Missing operand for option: '-%c'\n", argv0, optopt);
			usage();
			return 1;
		case '?':
			GETOPT_UNKNOWN_OPT
			usage();
			return 1;
		}
	}
	argc -= optind;
	argv += optind;
	if(e_flag)
	{
		// Fisher-Yates shuffles
		for(int i = 0; i <= argc - 2; i++)
		{
			int p = rand() % argc;
			// swap
			char *tmp = argv[p];
			argv[p] = argv[argc - 1];
			argv[argc - 1] = tmp;
		}
		unsigned long limit = argc;
		if(write_limit != 0 && write_limit < limit) limit = write_limit;
		for(unsigned long i = 0; i < limit; i++)
		{
			printf("%s%c", argv[i], delim);
		}
		return 0;
	}
	for(int i = 0; i < LINES_LEN; i++)
		lines[i] = NULL;
	if(argc <= 0)
	{
		if(shuf(stdin, "<stdin>") != 0) return 1;
	}
	else
	{
		for(int i = 0; i < argc; i++)
		{
			if(strncmp(argv[i], "-", 2) == 0)
			{
				if(shuf(stdin, "<stdin>") != 0) return 1;
				continue;
			}
			FILE *in = fopen(argv[i], "r");
			if(shuf(in, argv[i]) != 0)
			{
				fclose(in);
				return 1;
			}
			fclose(in);
		}
	}
	// inserts are random so iterating on it is fine
	for(int i = 0; i < LINES_LEN; i++)
	{
		if(write_limit != 0 && write_limit <= wrote) break;
		if(lines[i] != NULL)
		{
			fputs(lines[i], stdout);
			free(lines[i]);
			wrote++;
		}
	}
	return 0;
}