/* 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 <libmisc/_intercept.h>

#include "test.h"

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

#ifndef NDEBUG
jmp_buf *__catch;

void __lm_abort(void) {
	if (__catch)
		longjmp(*__catch, 1);
	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);
#endif

/* 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) {
#ifndef NDEBUG
		should_abort(rand_uint63n(cnt));
#else
		return;
#endif
	} else {
		double sum = 0;
		bool seen[MAX_SEE_ALL] = {0};
		for (int i = 0; i < ROUNDS; i++) {
			uint64_t val = rand_uint63n(cnt);
			sum += val;
			test_assert(val < cnt);
			if (cnt < MAX_SEE_ALL)
				seen[val] = true;
		}
		if (cnt > 1) {
			test_assert(sum/ROUNDS > 0.45*(cnt-1));
			test_assert(sum/ROUNDS < 0.55*(cnt-1));
		}
		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;
}