#include <assert.h>
#include <err.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#include "container_of.h"
#include "cradix.h"

#define	arraycount(A)	(sizeof(A)/sizeof(*(A)))

struct node {
	uint64_t key;
	struct cradix_node cradix;
};

static uint64_t
getkey(const struct cradix *C __unused, const struct cradix_node *N)
{
	const struct node *n = const_container_of(N, struct node, cradix);

	return n->key;
}

static const struct cradix_ops ops = {
	.key = &getkey,
};

static void
test_lookup_delete_empty(void)
{
	struct cradix cradix, *C = &cradix;

	cradix_init(C, &ops);
	assert(cradix_lookup(C, 0) == NULL);
	assert(cradix_delete_key(C, 0, NULL) == NULL);
	cradix_destroy(C);
}

static void
test_one_node(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node = { 0x0001020304050607 };

	cradix_init(C, &ops);
	assert(cradix_insert(C, &node.cradix, NULL) == &node.cradix);
	assert(cradix_lookup(C, 0x0001020304050607) == &node.cradix);
	assert(cradix_lookup(C, 0x0001020304050608) == NULL);
	assert(cradix_delete_key(C, 0x0001020304050608, NULL) == NULL);
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) == &node.cradix);
	cradix_destroy(C);
}

static void
test_two_nodes(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0f01020304050607 },
	};

	cradix_init(C, &ops);
	assert(cradix_insert(C, &node[0].cradix, NULL) == &node[0].cradix);
	assert(cradix_insert(C, &node[1].cradix, NULL) == &node[1].cradix);
	assert(cradix_lookup(C, 0x0001020304050607) == &node[0].cradix);
	assert(cradix_lookup(C, 0x0001020304050608) == NULL);
	assert(cradix_lookup(C, 0x0f01020304050607) == &node[1].cradix);
	assert(cradix_lookup(C, 0x0f01020304050608) == NULL);
	assert(cradix_lookup(C, 0x0801020304050608) == NULL);
	assert(cradix_delete_key(C, 0x0801020304050608, NULL) == NULL);
	assert(cradix_delete_key(C, 0x0001020304050608, NULL) == NULL);
	cradix_delete(C, &node[0].cradix, NULL);
	assert(cradix_delete_key(C, 0x0f01020304050608, NULL) == NULL);
	assert(cradix_delete_key(C, 0x0f01020304050607, NULL) ==
	    &node[1].cradix);
	cradix_destroy(C);
}

static void
test_lookup_tombstone(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0801020304050607 },
		{ 0x0f01020304050607 },
	};

	cradix_init(C, &ops);
	assert(cradix_insert(C, &node[0].cradix, NULL) == &node[0].cradix);
	assert(cradix_insert(C, &node[1].cradix, NULL) == &node[1].cradix);
	assert(cradix_insert(C, &node[2].cradix, NULL) == &node[2].cradix);
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) ==
	    &node[0].cradix);
	assert(cradix_lookup(C, 0x0001020304050607) == NULL);
	assert(cradix_lookup(C, 0x0801020304050607) == &node[1].cradix);
	assert(cradix_lookup(C, 0x0f01020304050607) == &node[2].cradix);
	assert(cradix_delete_key(C, 0x0801020304050607, NULL) ==
	    &node[1].cradix);
	assert(cradix_delete_key(C, 0x0f01020304050607, NULL) ==
	    &node[2].cradix);
	cradix_destroy(C);
}

static void
test_collision(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
		{ 0x0001020304050607 },
	};

	cradix_init(C, &ops);
	assert(cradix_insert(C, &node[0].cradix, NULL) == &node[0].cradix);
	assert(cradix_insert(C, &node[1].cradix, NULL) == &node[1].cradix);
	assert(cradix_insert(C, &node[2].cradix, NULL) == &node[0].cradix);
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) ==
	    &node[0].cradix);
	assert(cradix_delete_key(C, 0x0001020304050608, NULL) ==
	    &node[1].cradix);
	cradix_destroy(C);
}

static void
test_insert_post_nested(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0801020304050607 },
		{ 0x0001020304050608 },
		{ 0x0801020304050608 },
	};
	unsigned i;

	cradix_init(C, &ops);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_insert(C, &node[i].cradix, NULL) ==
		    &node[i].cradix);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_lookup(C, node[i].key) == &node[i].cradix);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_delete_key(C, node[i].key, NULL) ==
		    &node[i].cradix);
	cradix_destroy(C);
}

static void
test_insert_pre_nested(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
		{ 0x0801020304050607 },
		{ 0x0801020304050608 },
	};
	unsigned i;

	cradix_init(C, &ops);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_insert(C, &node[i].cradix, NULL) ==
		    &node[i].cradix);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_lookup(C, node[i].key) == &node[i].cradix);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_delete_key(C, node[i].key, NULL) ==
		    &node[i].cradix);
	cradix_destroy(C);
}

static void
test_refill_tombstone(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
		{ 0x0001020304050609 },
		{ 0x0001020304050607 },
	};
	unsigned i;

	cradix_init(C, &ops);
	assert(cradix_insert(C, &node[0].cradix, NULL) == &node[0].cradix);
	assert(cradix_insert(C, &node[1].cradix, NULL) == &node[1].cradix);
	assert(cradix_insert(C, &node[2].cradix, NULL) == &node[2].cradix);
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) ==
	    &node[0].cradix);
	assert(cradix_insert(C, &node[3].cradix, NULL) == &node[3].cradix);
	for (i = 1; i < 4; i++)
		assert(cradix_lookup(C, node[i].key) == &node[i].cradix);
	for (i = 1; i < 4; i++)
		assert(cradix_delete_key(C, node[i].key, NULL) ==
		    &node[i].cradix);
	cradix_destroy(C);
}

static void
test_pick_leaf_fallback(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
		{ 0x0801020304050607 },
		{ 0x0801020304050608 },
		{ 0x0f01020304050607 },
	};
	unsigned i;

	cradix_init(C, &ops);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_insert(C, &node[i].cradix, NULL) ==
		    &node[i].cradix);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_lookup(C, node[i].key) == &node[i].cradix);
	for (i = 0; i < arraycount(node); i++)
		assert(cradix_delete_key(C, node[i].key, NULL) ==
		    &node[i].cradix);
	cradix_destroy(C);
}

static void
test_grow_from_tombstone(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
		{ 0x0001020304050609 },
		{ 0x000102030405060a },
	};
	unsigned i;

	cradix_init(C, &ops);
	assert(cradix_insert(C, &node[0].cradix, NULL) == &node[0].cradix);
	assert(cradix_insert(C, &node[1].cradix, NULL) == &node[1].cradix);
	assert(cradix_insert(C, &node[2].cradix, NULL) == &node[2].cradix);
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) ==
	    &node[0].cradix);
	assert(cradix_insert(C, &node[3].cradix, NULL) == &node[3].cradix);
	for (i = 1; i < arraycount(node); i++)
		assert(cradix_lookup(C, node[i].key) == &node[i].cradix);
	for (i = 1; i < arraycount(node); i++)
		assert(cradix_delete_key(C, node[i].key, NULL) ==
		    &node[i].cradix);
	cradix_destroy(C);
}

static void
test_grow_shrink(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
		{ 0x0001020304050609 },
		{ 0x000102030405060a },
		{ 0x000102030405060b },
		{ 0x000102030405060c },
		{ 0x000102030405060d },
	};
	unsigned i;

	cradix_init(C, &ops);

	/* Insert all but one.  */
	for (i = 0; i < arraycount(node) - 1; i++)
		assert(cradix_insert(C, &node[i].cradix, NULL) ==
		    &node[i].cradix);

	/* Delete all but two of what we inserted.  */
	for (i = 2; i < arraycount(node) - 1; i++)
		assert(cradix_delete_key(C, node[i].key, NULL) ==
		    &node[i].cradix);

	/*
	 * Insert the last one.  The branch is full of tombstones, so
	 * it should shrink.
	 */
	assert(cradix_insert(C, &node[arraycount(node) - 1].cradix, NULL) ==
	    &node[arraycount(node) - 1].cradix);

	/* Confirm they can all be looked up.  */
	assert(cradix_lookup(C, 0x0001020304050607) == &node[0].cradix);
	assert(cradix_lookup(C, 0x0001020304050608) == &node[1].cradix);
	assert(cradix_lookup(C, 0x000102030405060d) ==
	    &node[arraycount(node) - 1].cradix);

	/* Confirm they can all be deleted.  */
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) ==
	    &node[0].cradix);
	assert(cradix_delete_key(C, 0x0001020304050608, NULL) ==
	    &node[1].cradix);
	assert(cradix_delete_key(C, 0x000102030405060d, NULL) ==
	    &node[arraycount(node) - 1].cradix);

	cradix_destroy(C);
}

static void
test_delete_tombstone(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
		{ 0x0001020304050609 },
	};
	unsigned i;

	cradix_init(C, &ops);
	assert(cradix_insert(C, &node[0].cradix, NULL) == &node[0].cradix);
	assert(cradix_insert(C, &node[1].cradix, NULL) == &node[1].cradix);
	assert(cradix_insert(C, &node[2].cradix, NULL) == &node[2].cradix);
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) ==
	    &node[0].cradix);
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) == NULL);
	for (i = 1; i < arraycount(node); i++)
		assert(cradix_lookup(C, node[i].key) == &node[i].cradix);
	for (i = 1; i < arraycount(node); i++)
		assert(cradix_delete_key(C, node[i].key, NULL) ==
		    &node[i].cradix);
	cradix_destroy(C);
}

static uint8_t custom_buf[4096];
static bool custom_allocated = false;
static bool custom_recycled = false;
static size_t custom_size = 0;

static void *
custom_alloc(struct cradix *C, size_t size)
{

	assert(!custom_allocated);
	assert(size > 0);
	custom_allocated = true;
	custom_size = size;
	return custom_buf;
}

static void
custom_free(struct cradix *C __unused, void *ptr __unused,
    size_t size __unused)
{

	abort();
}

static void
custom_recycle(struct cradix *C, void *ptr, size_t size)
{

	assert(custom_allocated);
	assert(size == custom_size);
	custom_allocated = false;
	custom_size = 0;
	custom_recycled = true;
}

static const struct cradix_ops custom_cradix_ops = {
	.key = &getkey,
	.alloc = &custom_alloc,
	.free = &custom_free,
	.recycle = &custom_recycle,
};

static void
test_custom_alloc(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
	};

	cradix_init(C, &custom_cradix_ops);

	assert(cradix_insert(C, &node[0].cradix, NULL) == &node[0].cradix);
	assert(!custom_allocated);
	assert(!custom_recycled);

	assert(cradix_insert(C, &node[1].cradix, NULL) == &node[1].cradix);
	assert(custom_allocated);
	assert(!custom_recycled);

	assert(cradix_delete_key(C, 0x0001020304050607, NULL) ==
	    &node[0].cradix);
	assert(!custom_allocated);
	assert(custom_recycled);

	assert(cradix_delete_key(C, 0x0001020304050608, NULL) ==
	    &node[1].cradix);
	assert(!custom_allocated);
	assert(custom_recycled);

	cradix_destroy(C);
}

/* 0 = no fault, 1 = delayed fault, 2 = fault */
static unsigned inject_alloc_fault = 0;
static unsigned faulty_nalloc = 0;
static unsigned faulty_nfree = 0;

static void *
faulty_alloc(struct cradix *C, size_t size)
{

	if (inject_alloc_fault == 2)
		return NULL;
	if (inject_alloc_fault)
		inject_alloc_fault++;
	faulty_nalloc++;
	return malloc(size);
}

static void
faulty_free(struct cradix *C, void *ptr, size_t size __unused)
{

	faulty_nfree++;
	free(ptr);
}

static const struct cradix_ops faulty_cradix_ops = {
	.key = &getkey,
	.alloc = &faulty_alloc,
	.free = &faulty_free,
};

static void
test_insert_begin_fail(void)
{
	struct cradix cradix, *C = &cradix;
	struct node node[] = {
		{ 0x0001020304050607 },
		{ 0x0001020304050608 },
		{ 0x0001020304050609 },
		{ 0x0801020304050607 },
	};
	struct cradix_insert insert, *I = &insert;
	struct cradix_delete delete, *D = &delete;

	cradix_init(C, &faulty_cradix_ops);

	assert(cradix_insert(C, &node[0].cradix, NULL) == &node[0].cradix);
	assert(cradix_insert(C, &node[1].cradix, NULL) == &node[1].cradix);

	/* Confirm growing the branch fails.  */
	inject_alloc_fault = 2;
	assert(cradix_insert(C, &node[2].cradix, NULL) == NULL);

	/* Confirm it fails.  */
	inject_alloc_fault = 2;
	faulty_nalloc = faulty_nfree = 0;
	assert(!cradix_insert_begin(C, I));
	assert(faulty_nalloc == 0);
	assert(faulty_nfree == 0);

	/* Confirm it runs the error branch.  */
	inject_alloc_fault = 1;
	faulty_nalloc = faulty_nfree = 0;
	assert(!cradix_insert_begin(C, I));
	assert(faulty_nalloc == 1);
	assert(faulty_nfree == 1);

	/* Confirm no allocation is needed after cradix_insert_begin.  */
	inject_alloc_fault = 0;
	assert(cradix_insert_begin(C, I));
	inject_alloc_fault = 2;
	assert(cradix_insert(C, &node[2].cradix, I) == &node[2].cradix);
	cradix_insert_end(C, I);

	/* Confirm making a new branch fails too.  */
	inject_alloc_fault = 2;
	assert(cradix_insert(C, &node[3].cradix, NULL) == NULL);

	/* Confirm making a new branch works with preallocation.  */
	inject_alloc_fault = 0;
	assert(cradix_insert_begin(C, I));
	inject_alloc_fault = 2;
	assert(cradix_insert(C, &node[3].cradix, I) == &node[3].cradix);
	cradix_insert_end(C, I);

	/* Confirm deletion still works despite allocq failure.  */
	inject_alloc_fault = 2;
	assert(cradix_delete_key(C, 0x0001020304050607, NULL) ==
	    &node[0].cradix);
	assert(cradix_delete_key(C, 0x0801020304050607, NULL) ==
	    &node[3].cradix);

	/* Confirm cradix_delete_begin/end work despite alloc failure.  */
	inject_alloc_fault = 2;
	cradix_delete_begin(C, D);
	assert(cradix_delete_key(C, 0x0001020304050608, D) == &node[1].cradix);
	cradix_delete_end(C, D);

	/* Confirm they work if we reuse the transaction record.  */
	inject_alloc_fault = 2;
	cradix_delete_begin(C, D);
	assert(cradix_delete_key(C, 0x0001020304050609, D) == &node[2].cradix);
	cradix_delete_end(C, D);

	cradix_destroy(C);
	inject_alloc_fault = 0;
}

static void
test_misc(void)
{
	struct node nodes[] = {
		{ 0x1a72793825c29927 },
		{ 0x1a2645932258442f },
		{ 0x1a1dbdecdf867a27 },
		{ 0x1be734b8713f0891 },
		{ 0xb3706073f27e7e75 },
		{ 0x1b101bb7da4c1827 },
		{ 0x1bca0c138c8ab327 },
		{ 0x1a0b822870a66433 },
		{ 0x53764233931ab6d8 },
		{ 0x8ea49a59a2c1c0b3 },
		{ 0x8d3c54acd5b76827 },
		{ 0x1ae502fb060ee427 },
		{ 0x1ae502fb060ee428 },
		{ 0x1a1dbdecdf867a26 },
		{ 0x1ae502fb060ee429 },
	};
	struct cradix cradix, *C = &cradix;
	struct cradix_node *N;
	struct cradix_insert insert, *I = &insert;
	unsigned i;

	cradix_init(C, &ops);
	cradix_dump(C);
	for (i = 0; i < arraycount(nodes) - 2; i++) {
		N = cradix_insert(C, &nodes[i].cradix, NULL);
		assert(N == &nodes[i].cradix);
	}
	for (i = 0; i < arraycount(nodes) - 2; i++) {
		N = cradix_lookup(C, nodes[i].key);
		assert(N == &nodes[i].cradix);
	}
	for (i = 0; i < arraycount(nodes) - 2; i += 2) {
		N = cradix_delete_key(C, nodes[i].key, NULL);
		assert(N == &nodes[i].cradix);
	}
	for (i = 0; i < arraycount(nodes) - 2; i++) {
		N = cradix_lookup(C, nodes[i].key);
		if (i % 2)
			assert(N == &nodes[i].cradix);
		else
			assert(N == NULL);
	}
	N = cradix_insert(C, &nodes[arraycount(nodes) - 2].cradix, NULL);
	assert(N == &nodes[arraycount(nodes) - 2].cradix);
	assert(cradix_insert_begin(C, I));
	N = cradix_insert(C, &nodes[arraycount(nodes) - 1].cradix, I);
	assert(N == &nodes[arraycount(nodes) - 1].cradix);
	cradix_insert_end(C, I);
	cradix_dump(C);
}

int
main(void)
{

	test_lookup_delete_empty();
	test_one_node();
	test_two_nodes();
	test_lookup_tombstone();
	test_collision();
	test_insert_post_nested();
	test_insert_pre_nested();
	test_refill_tombstone();
	test_pick_leaf_fallback();
	test_grow_from_tombstone();
	test_grow_shrink();
	test_delete_tombstone();
	test_insert_begin_fail();
	test_custom_alloc();

	test_misc();

	fflush(stdout);
	return ferror(stdout);
}