// SPDX-License-Identifier: GPL-2.0-only
/*
 * Test cases for the min max heap.
 */

#include <kunit/test.h>
#include <linux/min_heap.h>
#include <linux/module.h>
#include <linux/random.h>

struct min_heap_test_case {
	const char *str;
	bool min_heap;
};

static struct min_heap_test_case min_heap_cases[] = {
	{
		.str = "min",
		.min_heap = true,
	},
	{
		.str = "max",
		.min_heap = false,
	},
};

KUNIT_ARRAY_PARAM_DESC(min_heap, min_heap_cases, str);

DEFINE_MIN_HEAP(int, min_heap_test);

static bool less_than(const void *lhs, const void *rhs, void __always_unused *args)
{
	return *(int *)lhs < *(int *)rhs;
}

static bool greater_than(const void *lhs, const void *rhs, void __always_unused *args)
{
	return *(int *)lhs > *(int *)rhs;
}

static void pop_verify_heap(struct kunit *test,
			    bool min_heap,
			    struct min_heap_test *heap,
			    const struct min_heap_callbacks *funcs)
{
	int *values = heap->data;
	int last;

	last = values[0];
	min_heap_pop_inline(heap, funcs, NULL);
	while (heap->nr > 0) {
		if (min_heap)
			KUNIT_EXPECT_LE(test, last, values[0]);
		else
			KUNIT_EXPECT_GE(test, last, values[0]);
		last = values[0];
		min_heap_pop_inline(heap, funcs, NULL);
	}
}

static void test_heapify_all(struct kunit *test)
{
	const struct min_heap_test_case *params = test->param_value;
	int values[] = { 3, 1, 2, 4, 0x8000000, 0x7FFFFFF, 0,
			 -3, -1, -2, -4, 0x8000000, 0x7FFFFFF };
	struct min_heap_test heap = {
		.data = values,
		.nr = ARRAY_SIZE(values),
		.size =  ARRAY_SIZE(values),
	};
	struct min_heap_callbacks funcs = {
		.less = params->min_heap ? less_than : greater_than,
		.swp = NULL,
	};
	int i;

	/* Test with known set of values. */
	min_heapify_all_inline(&heap, &funcs, NULL);
	pop_verify_heap(test, params->min_heap, &heap, &funcs);

	/* Test with randomly generated values. */
	heap.nr = ARRAY_SIZE(values);
	for (i = 0; i < heap.nr; i++)
		values[i] = get_random_u32();

	min_heapify_all_inline(&heap, &funcs, NULL);
	pop_verify_heap(test, params->min_heap, &heap, &funcs);
}

static void test_heap_push(struct kunit *test)
{
	const struct min_heap_test_case *params = test->param_value;
	const int data[] = { 3, 1, 2, 4, 0x80000000, 0x7FFFFFFF, 0,
			     -3, -1, -2, -4, 0x80000000, 0x7FFFFFFF };
	int values[ARRAY_SIZE(data)];
	struct min_heap_test heap = {
		.data = values,
		.nr = 0,
		.size =  ARRAY_SIZE(values),
	};
	struct min_heap_callbacks funcs = {
		.less = params->min_heap ? less_than : greater_than,
		.swp = NULL,
	};
	int i, temp;

	/* Test with known set of values copied from data. */
	for (i = 0; i < ARRAY_SIZE(data); i++)
		min_heap_push_inline(&heap, &data[i], &funcs, NULL);

	pop_verify_heap(test, params->min_heap, &heap, &funcs);

	/* Test with randomly generated values. */
	while (heap.nr < heap.size) {
		temp = get_random_u32();
		min_heap_push_inline(&heap, &temp, &funcs, NULL);
	}
	pop_verify_heap(test, params->min_heap, &heap, &funcs);
}

static void test_heap_pop_push(struct kunit *test)
{
	const struct min_heap_test_case *params = test->param_value;
	const int data[] = { 3, 1, 2, 4, 0x80000000, 0x7FFFFFFF, 0,
			     -3, -1, -2, -4, 0x80000000, 0x7FFFFFFF };
	int values[ARRAY_SIZE(data)];
	struct min_heap_test heap = {
		.data = values,
		.nr = 0,
		.size =  ARRAY_SIZE(values),
	};
	struct min_heap_callbacks funcs = {
		.less = params->min_heap ? less_than : greater_than,
		.swp = NULL,
	};
	int i, temp;

	/* Fill values with data to pop and replace. */
	temp = params->min_heap ? 0x80000000 : 0x7FFFFFFF;
	for (i = 0; i < ARRAY_SIZE(data); i++)
		min_heap_push_inline(&heap, &temp, &funcs, NULL);

	/* Test with known set of values copied from data. */
	for (i = 0; i < ARRAY_SIZE(data); i++)
		min_heap_pop_push_inline(&heap, &data[i], &funcs, NULL);

	pop_verify_heap(test, params->min_heap, &heap, &funcs);

	heap.nr = 0;
	for (i = 0; i < ARRAY_SIZE(data); i++)
		min_heap_push_inline(&heap, &temp, &funcs, NULL);

	/* Test with randomly generated values. */
	for (i = 0; i < ARRAY_SIZE(data); i++) {
		temp = get_random_u32();
		min_heap_pop_push_inline(&heap, &temp, &funcs, NULL);
	}
	pop_verify_heap(test, params->min_heap, &heap, &funcs);
}

static void test_heap_del(struct kunit *test)
{
	const struct min_heap_test_case *params = test->param_value;
	int values[] = { 3, 1, 2, 4, 0x8000000, 0x7FFFFFF, 0,
			 -3, -1, -2, -4, 0x8000000, 0x7FFFFFF };
	struct min_heap_test heap;

	min_heap_init_inline(&heap, values, ARRAY_SIZE(values));
	heap.nr = ARRAY_SIZE(values);
	struct min_heap_callbacks funcs = {
		.less = params->min_heap ? less_than : greater_than,
		.swp = NULL,
	};
	int i;

	/* Test with known set of values. */
	min_heapify_all_inline(&heap, &funcs, NULL);
	for (i = 0; i < ARRAY_SIZE(values) / 2; i++)
		min_heap_del_inline(&heap, get_random_u32() % heap.nr, &funcs, NULL);
	pop_verify_heap(test, params->min_heap, &heap, &funcs);

	/* Test with randomly generated values. */
	heap.nr = ARRAY_SIZE(values);
	for (i = 0; i < heap.nr; i++)
		values[i] = get_random_u32();
	min_heapify_all_inline(&heap, &funcs, NULL);

	for (i = 0; i < ARRAY_SIZE(values) / 2; i++)
		min_heap_del_inline(&heap, get_random_u32() % heap.nr, &funcs, NULL);
	pop_verify_heap(test, params->min_heap, &heap, &funcs);
}

static struct kunit_case min_heap_test_cases[] = {
	KUNIT_CASE_PARAM(test_heapify_all, min_heap_gen_params),
	KUNIT_CASE_PARAM(test_heap_push, min_heap_gen_params),
	KUNIT_CASE_PARAM(test_heap_pop_push, min_heap_gen_params),
	KUNIT_CASE_PARAM(test_heap_del, min_heap_gen_params),
	{},
};

static struct kunit_suite min_heap_test_suite = {
	.name = "min_heap",
	.test_cases = min_heap_test_cases,
};

kunit_test_suite(min_heap_test_suite);

MODULE_DESCRIPTION("Test cases for the min max heap");
MODULE_LICENSE("GPL");