/* Tiny AVL trees
 * Copyright (C) 2020 Tirotech Ltd
 *
 * Author: Daniel Beer <daniel.beer@tirotech.co.nz>
 *
 * Permission to use, copy, modify, and/or distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include "tavl.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <limits.h>
#include <assert.h>

#define N	10000

static struct tavl_node arr[N];
static tavl_index_t root = TAVL_NONE;
static tavl_index_t freelist = TAVL_NONE;

static void freelist_push(tavl_index_t i)
{
	arr[i].down[0] = freelist;
	freelist = i;
}

static tavl_index_t freelist_pop(void)
{
	const tavl_index_t r = freelist;

	assert(r != TAVL_NONE);
	freelist = arr[r].down[0];

	return r;
}

static tavl_index_t put(tavl_key_t k)
{
	tavl_index_t i = tavl_find(arr, root, k);

	if (i != TAVL_NONE)
		return i;

	i = freelist_pop();
	arr[i].key = k;

	tavl_index_t r = tavl_insert(arr, &root, i);
	assert(r == TAVL_NONE);
	return i;
}

static void del(tavl_key_t k)
{
	tavl_index_t i = tavl_delete(arr, &root, k);

	if (i != TAVL_NONE) {
		assert(arr[i].key == k);
		freelist_push(i);
	}
}

static int tree_count(tavl_index_t i)
{
	if (i == TAVL_NONE)
		return 0;

	return 1 + tree_count(arr[i].down[0] & TAVL_INDEX)
		 + tree_count(arr[i].down[1] & TAVL_INDEX);
}

static int check_invariant(tavl_index_t n, tavl_key_t a, tavl_key_t b)
{
	if (n == TAVL_NONE)
		return 0;

	const tavl_key_t k = arr[n].key;

	assert(k >= a);
	assert(k <= b);

	const int l = check_invariant(arr[n].down[0] & TAVL_INDEX,
		a, k - 1);
	const int r = check_invariant(arr[n].down[1] & TAVL_INDEX,
		k + 1, b);

	assert(abs(l - r) <= 1);
	assert(!(arr[n].down[0] & arr[n].down[1] & TAVL_HEAVY));
	assert(l <= r || (arr[n].down[0] & TAVL_HEAVY));
	assert(l >= r || (arr[n].down[1] & TAVL_HEAVY));

	const int max = l > r ? l : r;

	return max + 1;
}

static void check_state(void)
{
	int c = 0;

	for (tavl_index_t i = freelist; i != TAVL_NONE; i = arr[i].down[0])
		c++;

	c += tree_count(root);
	assert(c == N);

	check_invariant(root, 0, UINT32_MAX);
}

/************************************************************************
 * Test schedule
 */

struct test_item {
	tavl_key_t	key;
	int		member;
};

static struct test_item tests[N];
static tavl_key_t values[N];

static void shuffle(void)
{
	for (int i = 0; i < N; i++) {
		int j = random() % (N - i) + i;
		struct test_item tmp;

		memcpy(&tmp, &tests[i], sizeof(struct test_item));
		memcpy(&tests[i], &tests[j], sizeof(struct test_item));
		memcpy(&tests[j], &tmp, sizeof(struct test_item));
	}
}

static void check_membership(void)
{
	for (int i = 0; i < N; i++) {
		const struct test_item *t = &tests[i];
		const tavl_index_t n = tavl_find(arr, root, t->key);

		if (t->member) {
			assert(n != TAVL_NONE);
			assert(arr[n].key == t->key);
			assert(values[n] == t->key);
		} else {
			assert(n == TAVL_NONE);
		}
	}
}

static void dump_recurse(int level, tavl_index_t n)
{
	for (int i = 0; i < level; i++)
		printf(" ");

	if (n == TAVL_NONE) {
		printf("NIL\n");
		return;
	}

	printf("%08x", arr[n].key);

	if (arr[n].down[0] & TAVL_HEAVY)
		printf(" LEFT-HEAVY");
	if (arr[n].down[1] & TAVL_HEAVY)
		printf(" RIGHT-HEAVY");

	if ((arr[n].down[0] == TAVL_NONE) && (arr[n].down[1] == TAVL_NONE)) {
		printf(" LEAF\n");
		return;
	}

	printf("\n");

	dump_recurse(level + 1, arr[n].down[0] & TAVL_INDEX);
	dump_recurse(level + 1, arr[n].down[1] & TAVL_INDEX);
}

static void dump(const char *label)
{
	printf("\n%s:\n", label);
	dump_recurse(1, root);
}

int main(void)
{
	printf("TAVL_INDEX=%04x, TAVL_HEAVY=%04x\n", TAVL_INDEX, TAVL_HEAVY);
	assert(!(TAVL_INDEX & TAVL_HEAVY));
	assert(TAVL_INDEX < TAVL_HEAVY);

	srandom(1);

	printf("init\n");
	for (int i = 0; i < N; i++) {
		freelist_push(i);
		tests[i].key = i * 5;
	}

	printf("grow\n");
	shuffle();

	for (int i = 0; i < N; i++) {
		struct test_item *t = &tests[i];

		values[put(t->key)] = t->key;
		t->member = 1;
	}

	check_state();
	check_membership();

	for (int k = 0; k < 10; k++) {
		printf("thrash (%d)\n", k);
		shuffle();

		for (int i = 0; i < N; i++) {
			struct test_item *t = &tests[i];

			if (random() & 1) {
				del(t->key);
				t->member = 0;
			} else {
				values[put(t->key)] = t->key;
				t->member = 1;
			}
		}

		check_state();
		check_membership();
	}

	printf("shrink\n");
	shuffle();

	for (int i = 0; i < N; i++) {
		struct test_item *t = &tests[i];

		del(t->key);
		t->member = 0;

		if (i + 50 == N)
			dump("50 deletions remain");
	}

	check_state();
	check_membership();
	return 0;
}