/*-
 * SPDX-License-Identifier: BSD-2-Clause
 *
 * Copyright (c) 2024 Robert Clausecker <fuz@FreeBSD.org>
 *
 * biastest.c -- bias test for arc4random_uniform().
 *
 * The default configuration of this test has an upper bound of
 * (3/4) * UINT32_MAX, which should give a high amount of bias in
 * an incorrect implementation.  If the range reduction is
 * implemented correctly, the parameters of the statistic should
 * closely match the expected values.  If not, they'll differ.
 *
 * For memory usage reasons, we use an uchar to track the number of
 * observations per bucket.  If the number of tries is much larger
 * than upper_bound, the buckets likely overflow.  This is detected
 * by the test, but will lead to incorrect results.
 */

#include <assert.h>
#include <limits.h>
#include <math.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>

static void	collect_sample(unsigned char *, long long, uint32_t);
static void	analyze_sample(const unsigned char *, long long, uint32_t);

static atomic_bool complete = false;
static long long tries = 5ULL << 32;
static atomic_llong tries_done = 0;

static void
usage(const char *argv0)
{
	fprintf(stderr, "usage: %s [-n tries] [-t upper_bound]\n", argv0);
	exit(EXIT_FAILURE);
}

int
main(int argc, char *argv[])
{
	uint32_t threshold = 3UL << 30;
	int ch;
	unsigned char *sample;

	while (ch = getopt(argc, argv, "n:t:"), ch != EOF)
		switch (ch) {
		case 'n':
			tries = atoll(optarg);
			break;

		case 't':
			threshold = (uint32_t)atoll(optarg);
			break;

		default:
			usage(argv[0]);
		}

	if (optind != argc)
		usage(argv[0]);

	if (threshold == 0) {
		fprintf(stderr, "threshold must be between 1 and %lu\n", (unsigned long)UINT32_MAX);
		exit(EXIT_FAILURE);
	}

	sample = calloc(threshold, 1);
	if (sample == NULL) {
		perror("calloc(threshold, 1)");
		return (EXIT_FAILURE);
	}

	collect_sample(sample, tries, threshold);
	analyze_sample(sample, tries, threshold);
}

static void
progress(int signo)
{
	(void)signo;

	if (!complete) {
		fprintf(stderr, "\r%10lld of %10lld samples taken (%5.2f%% done)",
		    tries_done, tries, (tries_done * 100.0) / tries);

		signal(SIGALRM, progress);
		alarm(1);
	}
}

static void
collect_sample(unsigned char *sample, long long tries, uint32_t threshold)
{
	long long i;
	uint32_t x;
	bool overflowed = false;

	progress(SIGALRM);

	for (i = 0; i < tries; i++) {
		x = arc4random_uniform(threshold);
		tries_done++;
		assert(x < threshold);

		if (sample[x] == UCHAR_MAX) {
			if (!overflowed) {
				printf("sample table overflow, results will be incorrect\n");
				overflowed = true;
			}
		} else
			sample[x]++;
	}

	progress(SIGALRM);
	complete = true;
	fputc('\n', stderr);
}

static void
analyze_sample(const unsigned char *sample, long long tries,  uint32_t threshold)
{
	double discrepancy, average, variance, total;
	long long histogram[UCHAR_MAX + 1] = { 0 }, sum, n, median;
	uint32_t i, i_min, i_max;
	int min, max;

	printf("distribution properties:\n");

	/* find median, average, deviation, smallest, and largest bucket */
	total = 0.0;
	for (i = 0; i < threshold; i++) {
		histogram[sample[i]]++;
		total += (double)i * sample[i];
	}

	average = total / tries;

	variance = 0.0;
	median = threshold;
	n = 0;
	i_min = 0;
	i_max = 0;
	min = sample[i_min];
	max = sample[i_max];

	for (i = 0; i < threshold; i++) {
		discrepancy = i - average;
		variance += sample[i] * discrepancy * discrepancy;

		n += sample[i];
		if (median == threshold && n > tries / 2)
			median = i;

		if (sample[i] < min) {
			i_min = i;
			min = sample[i_min];
		} else if (sample[i] > max) {
			i_max = i;
			max = sample[i_max];
		}
	}

	variance /= tries;
	assert(median < threshold);

	printf("\tthreshold:	%lu\n", (unsigned long)threshold);
	printf("\tobservations:	%lld\n", tries);
	printf("\tleast common:	%lu (%d observations)\n", (unsigned long)i_min, min);
	printf("\tmost common:	%lu (%d observations)\n", (unsigned long)i_max, max);
	printf("\tmedian:		%lld (expected %lu)\n", median, (unsigned long)threshold / 2);
	printf("\taverage:	%f (expected %f)\n", average, 0.5 * (threshold - 1));
	printf("\tdeviation:	%f (expected %f)\n\n", sqrt(variance),
	    sqrt(((double)threshold * threshold - 1.0) / 12));

	/* build histogram and analyze it */
	printf("sample properties:\n");

	/* find median, average, and deviation */
	average = (double)tries / threshold;

	variance = 0.0;
	for (i = 0; i < UCHAR_MAX; i++) {
		discrepancy = i - average;
		variance += histogram[i] * discrepancy * discrepancy;
	}

	variance /= threshold;

	n = 0;
	median = UCHAR_MAX + 1;
	for (i = 0; i <= UCHAR_MAX; i++) {
		n += histogram[i];
		if (n >= threshold / 2) {
			median = i;
			break;
		}
	}

	assert(median <= UCHAR_MAX); /* unreachable */

	printf("\tmedian:		%lld\n", median);
	printf("\taverage:	%f\n", average);
	printf("\tdeviation:	%f (expected %f)\n\n", sqrt(variance), sqrt(average * (1.0 - 1.0 / threshold)));

	printf("histogram:\n");
	for (i = 0; i < 256; i++)
		if (histogram[i] != 0)
			printf("\t%3d:\t%lld\n", (int)i, histogram[i]);
}