/* libmisc/tests/test_rand.c - Tests for <libmisc/rand.h>
 *
 * Copyright (C) 2024  Luke T. Shumaker <lukeshu@lukeshu.com>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 */

#include <stdbool.h>
#include <setjmp.h>

#include <libmisc/rand.h>

#include "test.h"

/* Intercept failures *********************************************************/

jmp_buf *__catch;

void __assert_msg_fail(const char *expr,
                  const char *file, unsigned int line, const char *func,
                  const char *msg) {
	static bool in_fail = false;
	if (__catch)
		longjmp(*__catch, 1);
	if (!in_fail) {
		in_fail = true;
		printf("error: %s:%u:%s(): assertion \"%s\" failed%s%s\n",
		       file, line, func,
		       expr,
		       msg ? ": " : "", msg);
	}
	abort();
}

#define should_abort(cmd) do {                \
		jmp_buf *old_catch = __catch; \
		jmp_buf env;                  \
		__catch = &env;               \
		if (!setjmp(env)) {           \
			cmd;                  \
			__catch = old_catch;  \
			test_assert(false);   \
		} else {                      \
			__catch = old_catch;  \
		}                             \
	} while (0);

/* Actual tests ***************************************************************/

#define ROUNDS 4096
#define MAX_SEE_ALL 128

static void test_n(uint64_t cnt) {
	if (cnt == 0 || cnt > UINT64_C(1)<<63) {
		should_abort(rand_uint63n(cnt));
	} else {
		double sum = 0;
		bool seen[MAX_SEE_ALL] = {0};
		for (int i = 0; i < ROUNDS; i++) {
			uint64_t val = rand_uint63n(cnt);
			sum += ((double)val)/(cnt-1);
			test_assert(val < cnt);
			if (cnt < MAX_SEE_ALL)
				seen[val] = true;
		}
		if (cnt > 1) {
			test_assert(sum/ROUNDS > 0.45);
			test_assert(sum/ROUNDS < 0.55);
		}
		if (cnt < MAX_SEE_ALL) {
			for (uint64_t i = 0; i < cnt; i++)
				test_assert(seen[i]);
		}
	}
}

int main() {
	for (uint8_t i = 0; i < 64; i++)
		test_n(UINT64_C(1)<<i);
	for (uint64_t j = 0; j < MAX_SEE_ALL; j++)
		test_n(j);
	return 0;
}