/* libheap/tests/test_heap.c - Tests for libheap
 *
 * Copyright (C) 2025  Luke T. Shumaker <lukeshu@lukeshu.com>
 * SPDX-License-Identifier: AGPL-3.0-or-later
 */

#include <string.h> /* for memcmp() */
#include <stdlib.h> /* for exit() */

#include <libheap/inspect.h>
#include <libmisc/assert.h>
#include <libmisc/heap.h>
#include <libmisc/macro.h>

#define LOG_NAME TEST
#include <libmisc/log.h>

#include <valgrind/valgrind.h>

static struct {
	void     *heap_base;
	size_t    heap_size;
	void    (*abort)(void);
} globals;

////////////////////////////////////////////////////////////////////////////////

struct heap {
	void    *base;
	size_t   size;
};

struct heap heap_get(void) {
	return (struct heap){
		.base = globals.heap_base,
		.size = globals.heap_size,
	};
}

void __lm_abort(void) {
	globals.abort();
	__builtin_unreachable();
}

////////////////////////////////////////////////////////////////////////////////

#define test_assert(expr) do {                                                                      \
	if (!(expr)) {                                                                              \
		log_errorln("test failure: ", __FILE__, ":", __LINE__, ":", __func__, ": ", #expr); \
		exit(1);                                                                            \
	}                                                                                           \
} while (0)

#define test_assert_heap(...) do {                                                              \
	uintptr_t exp_heap[LM_ARRAY_LEN(heap)] = { __VA_ARGS__ };                               \
	VALGRIND_DISABLE_ERROR_REPORTING;                                                       \
	if (memcmp(heap, exp_heap, sizeof(heap))) {                                             \
		log_errorln("[  ]\t", (ljust, sizeof(uintptr_t)*2+2, ' ', "exp"), "\t==\tact"); \
		for (size_t i = 0; i < LM_ARRAY_LEN(heap); i++)                                 \
			log_errorln("[", i, "]\t",                                              \
			            (ptr, (void*)exp_heap[i]),                                  \
			            "\t", (str, exp_heap[i] == heap[i] ? "  " : "!="), "\t",    \
			            (ptr, (void*)heap[i]));                                     \
		VALGRIND_ENABLE_ERROR_REPORTING;                                                \
		log_errorln("test failure: ", __FILE__, ":", __LINE__, ":", __func__, ": ^^^"); \
		exit(1);                                                                        \
	}                                                                                       \
	VALGRIND_ENABLE_ERROR_REPORTING;                                                        \
} while (0)

void test_abort(void) {
	test_assert(!"aborted");
}

static_assert(sizeof(void*) == sizeof(size_t));
static_assert(sizeof(void*) == sizeof(uintptr_t));
static_assert(sizeof(void*) == sizeof(cid_t));

int main() {
	uintptr_t heap[3*(2+8)] = {};
	globals.heap_base = heap;
	globals.heap_size = sizeof(heap);
	globals.abort = test_abort;

	{
		log_infoln("== test 1: basic init ===========================");
		struct heap_info info;
		heap_usage(&info);

		struct heap_info exp_info = {
			.internal_use = 2*sizeof(void*),
			.allocated = {
				[HEAP_CID_UNALLOCATED] = {
					.sum = sizeof(heap)-2*sizeof(void*),
					.nseg = 1,
				}
			}
		};
		test_assert(memcmp(&info, &exp_info, sizeof(info)) == 0);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 1 */ HEAP_CID_UNALLOCATED,
		);
	}

	{
		log_infoln("== test 2: zero alloc ===========================");
		void *ptr = heap_alloc(0, char);
		test_assert(ptr == &heap[2]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[2],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 3 */ HEAP_CID_UNALLOCATED,
		);

		heap_free(&heap[2]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 1 */ HEAP_CID_UNALLOCATED,
		);
	}

	{
		log_infoln("== test 3: free =================================");
		/* We have 9 scenarios that we want to test:
		 *
		 *     {beg,free,used}|block|{free,used,end}
		 */

		void *ptr = heap_alloc(LM_ARRAY_LEN(heap)-2, uintptr_t);
		test_assert(ptr == &heap[2]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 1 */ HEAP_CID_KERNEL,
		);
		log_infoln("1/9: beg|block|end"); heap_free(ptr);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 1 */ HEAP_CID_UNALLOCATED,
		);

		ptr = heap_alloc(2, uintptr_t);
		test_assert(ptr == &heap[2]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 0,
			/* 3 */ 0,
			/* 4 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 5 */ HEAP_CID_UNALLOCATED,
		);
		log_infoln("2/9: beg|block|free"); heap_free(ptr);
		test_assert_heap(
			(uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			HEAP_CID_UNALLOCATED,
		);

		ptr = heap_alloc(2, uintptr_t);
		test_assert(ptr == &heap[2]);
		void *ptr2 = heap_alloc(2, uintptr_t);
		test_assert(ptr2 == &heap[6]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 0,
			/* 3 */ 0,
			/* 4 */ (uintptr_t)&heap[8],
			/* 5 */ HEAP_CID_KERNEL,
			/* 6 */ 0,
			/* 7 */ 0,
			/* 8 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 9 */ HEAP_CID_UNALLOCATED,
		);
		log_infoln("3/9: beg|block|used"); heap_free(ptr);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_UNALLOCATED,
			/* 2 */ 0,
			/* 3 */ 0,
			/* 4 */ (uintptr_t)&heap[8],
			/* 5 */ HEAP_CID_KERNEL,
			/* 6 */ 0,
			/* 7 */ 0,
			/* 8 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 9 */ HEAP_CID_UNALLOCATED,
		);
		log_infoln("4/9: unused|block|unused"); heap_free(ptr2);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 1 */ HEAP_CID_UNALLOCATED,
		);

		ptr = heap_alloc(2, uintptr_t);
		test_assert(ptr == &heap[2]);
		ptr2 = heap_alloc(LM_ARRAY_LEN(heap)-6, uintptr_t);
		test_assert(ptr2 == &heap[6]);
		heap_free(ptr);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_UNALLOCATED,
			/* 2 */ 0,
			/* 3 */ 0,
			/* 4 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 5 */ HEAP_CID_KERNEL,
		);
		log_infoln("5/9: unused|block|end"); heap_free(ptr2);
		test_assert_heap(
			(uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			HEAP_CID_UNALLOCATED,
		);

		ptr = heap_alloc(2, uintptr_t);
		test_assert(ptr == &heap[2]);
		ptr2 = heap_alloc(2, uintptr_t);
		test_assert(ptr2 == &heap[6]);
		heap_free(ptr);
		ptr = heap_alloc(3, uintptr_t);
		test_assert(ptr == &heap[10]);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[4],
			/* _1 */ HEAP_CID_UNALLOCATED,
			/* _2 */ 0,
			/* _3 */ 0,
			/* _4 */ (uintptr_t)&heap[8],
			/* _5 */ HEAP_CID_KERNEL,
			/* _6 */ 0,
			/* _7 */ 0,
			/* _8 */ (uintptr_t)&heap[13],
			/* _9 */ HEAP_CID_KERNEL,
			/* 10 */ 0,
			/* 11 */ 0,
			/* 12 */ 0,
			/* 13 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 14 */ HEAP_CID_UNALLOCATED,
		);
		log_infoln("6/9: unused|block|used"); heap_free(ptr2);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[8],
			/* _1 */ HEAP_CID_UNALLOCATED,
			/* _2 */ 0,
			/* _3 */ 0,
			/* _4 */ 0,
			/* _5 */ 0,
			/* _6 */ 0,
			/* _7 */ 0,
			/* _8 */ (uintptr_t)&heap[13],
			/* _9 */ HEAP_CID_KERNEL,
			/* 10 */ 0,
			/* 11 */ 0,
			/* 12 */ 0,
			/* 13 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 14 */ HEAP_CID_UNALLOCATED,
		);
		heap_free(ptr);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 1 */ HEAP_CID_UNALLOCATED,
		);

		ptr2 = heap_alloc(2, uintptr_t);
		test_assert(ptr2 == &heap[2]);
		ptr = heap_alloc(2, uintptr_t);
		test_assert(ptr == &heap[6]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 0,
			/* 3 */ 0,
			/* 4 */ (uintptr_t)&heap[8],
			/* 5 */ HEAP_CID_KERNEL,
			/* 6 */ 0,
			/* 7 */ 0,
			/* 8 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 9 */ HEAP_CID_UNALLOCATED,
		);
		log_infoln("7/9: used|block|unused"); heap_free(ptr);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 0,
			/* 3 */ 0,
			/* 4 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 5 */ HEAP_CID_UNALLOCATED,
		);

		ptr = heap_alloc(LM_ARRAY_LEN(heap)-6, uintptr_t);
		test_assert(ptr == &heap[6]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 0,
			/* 3 */ 0,
			/* 4 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 5 */ HEAP_CID_KERNEL,
		);
		log_infoln("8/9: used|block|end"); heap_free(ptr);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 0,
			/* 3 */ 0,
			/* 4 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 5 */ HEAP_CID_UNALLOCATED,
		);

		ptr = heap_alloc(3, uintptr_t);
		test_assert(ptr == &heap[6]);
		ptr2 = heap_alloc(4, uintptr_t);
		test_assert(ptr2 == &heap[11]);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[4],
			/* _1 */ HEAP_CID_KERNEL,
			/* _2 */ 0,
			/* _3 */ 0,
			/* _4 */ (uintptr_t)&heap[9],
			/* _5 */ HEAP_CID_KERNEL,
			/* _6 */ 0,
			/* _7 */ 0,
			/* _8 */ 0,
			/* _9 */ (uintptr_t)&heap[15],
			/* 10 */ HEAP_CID_KERNEL,
			/* 11 */ 0,
			/* 12 */ 0,
			/* 13 */ 0,
			/* 14 */ 0,
			/* 15 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 16 */ HEAP_CID_UNALLOCATED,
		);
		log_infoln("9/9: used|block|used"); heap_free(ptr);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[4],
			/* _1 */ HEAP_CID_KERNEL,
			/* _2 */ 0,
			/* _3 */ 0,
			/* _4 */ (uintptr_t)&heap[9],
			/* _5 */ HEAP_CID_UNALLOCATED,
			/* _6 */ 0,
			/* _7 */ 0,
			/* _8 */ 0,
			/* _9 */ (uintptr_t)&heap[15],
			/* 10 */ HEAP_CID_KERNEL,
			/* 11 */ 0,
			/* 12 */ 0,
			/* 13 */ 0,
			/* 14 */ 0,
			/* 15 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 16 */ HEAP_CID_UNALLOCATED,
		);

		heap_free(&heap[2]);
		heap_free(&heap[11]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 1 */ HEAP_CID_UNALLOCATED,
		);
	}

	{
		log_infoln("== test 4: realloc ==============================");
		/* We have 4 scenarios that we want to test:
		 *  - no-op
		 *  - grow
		 *  - non-overlapping move
		 *  - overlapping move
		 */

		uintptr_t *ptr = heap_alloc(2, uintptr_t);
		test_assert(ptr == &heap[2]);
		ptr[0] = 98;
		ptr[1] = 99;
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 98,
			/* 3 */ 99,
			/* 4 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 5 */ HEAP_CID_UNALLOCATED,
		);

		log_infoln("1/4: no-op");
		ptr = heap_realloc(ptr, 2, uintptr_t);
		test_assert(ptr == &heap[2]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[4],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 98,
			/* 3 */ 99,
			/* 4 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 5 */ HEAP_CID_UNALLOCATED,
		);

		log_infoln("2/4: grow");
		ptr = heap_realloc(ptr, 4, uintptr_t);
		test_assert(ptr == &heap[2]);
		test_assert_heap(
			/* 0 */ (uintptr_t)&heap[6],
			/* 1 */ HEAP_CID_KERNEL,
			/* 2 */ 98,
			/* 3 */ 99,
			/* 4 */ 0,
			/* 5 */ 0,
			/* 6 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 7 */ HEAP_CID_UNALLOCATED,
		);

		log_infoln("3/4: non-overlapping move");
		uintptr_t *ptr2 = heap_alloc(1, uintptr_t);
		test_assert(ptr2 == &heap[8]);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[6],
			/* _1 */ HEAP_CID_KERNEL,
			/* _2 */ 98,
			/* _3 */ 99,
			/* _4 */ 0,
			/* _5 */ 0,
			/* _6 */ (uintptr_t)&heap[9],
			/* _7 */ HEAP_CID_KERNEL,
			/* _8 */ 0,
			/* _9 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 10 */ HEAP_CID_UNALLOCATED,
		);
		ptr = heap_realloc(ptr, 6, uintptr_t);
		test_assert(ptr == &heap[11]);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[6],
			/* _1 */ HEAP_CID_UNALLOCATED,
			/* _2 */ 98,
			/* _3 */ 99,
			/* _4 */ 0,
			/* _5 */ 0,
			/* _6 */ (uintptr_t)&heap[9], /* ptr2 */
			/* _7 */ HEAP_CID_KERNEL,
			/* _8 */ 0,
			/* _9 */ (uintptr_t)&heap[17], /* ptr */
			/* 10 */ HEAP_CID_KERNEL,
			/* 11 */ 98,
			/* 12 */ 99,
			/* 13 */ 0,
			/* 14 */ 0,
			/* 15 */ 0,
			/* 16 */ 0,
			/* 17 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* 18 */ HEAP_CID_UNALLOCATED,
		);

		log_infoln("4/4: overlapping move");
		ptr2[0] = 66;
		ptr = heap_realloc(ptr, LM_ARRAY_LEN(heap)-11, uintptr_t);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[6],
			/* _1 */ HEAP_CID_UNALLOCATED,
			/* _2 */ 98,
			/* _3 */ 99,
			/* _4 */ 0,
			/* _5 */ 0,
			/* _6 */ (uintptr_t)&heap[9], /* ptr2 */
			/* _7 */ HEAP_CID_KERNEL,
			/* _8 */ 66,
			/* _9 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)], /* ptr */
			/* 10 */ HEAP_CID_KERNEL,
			/* 11 */ 98,
			/* 12 */ 99,
		);
		ptr2 = heap_realloc(ptr2, 7, uintptr_t);
		test_assert(ptr2 == &heap[2]);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[9], /* ptr2 */
			/* _1 */ HEAP_CID_KERNEL,
			/* _2 */ 66,
			/* _3 */ 0,
			/* _4 */ 0,
			/* _5 */ 0,
			/* _6 */ 0,
			/* _7 */ 0,
			/* _8 */ 0,
			/* _9 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)], /* ptr */
			/* 10 */ HEAP_CID_KERNEL,
			/* 11 */ 98,
			/* 12 */ 99,
		);

		heap_free(ptr);
		heap_free(ptr2);
		test_assert_heap(
			/* _0 */ (uintptr_t)&heap[LM_ARRAY_LEN(heap)],
			/* _1 */ HEAP_CID_UNALLOCATED,
			/* _2 */ 66,
			/* _3 */ 0,
			/* _4 */ 0,
			/* _5 */ 0,
			/* _6 */ 0,
			/* _7 */ 0,
			/* _8 */ 0,
			/* _9 */ 0,
			/* 10 */ 0,
			/* 11 */ 98,
			/* 12 */ 99,
		);
	}

	log_infoln("== passed =======================================");
	return 0;
}