/* 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"

/* Perform a rotate-and-fix of a node which is too heavy on one side.
 * Updates the given pointer with a new root, and returns the reduction
 * in height due to the rotation (0 or 1).
 */
static int rotate_and_fix(struct tavl_node *arr, tavl_index_t *ptr, int dir)
{
	const int opposite = dir ^ 1;
	const tavl_index_t x = (*ptr) & TAVL_INDEX;
	const tavl_index_t z = arr[x].down[dir] & TAVL_INDEX;

	/* Simple rotation */
	if (!(arr[z].down[opposite] & TAVL_HEAVY)) {
		const tavl_index_t a = arr[x].down[opposite] & TAVL_INDEX;
		const tavl_index_t b = arr[z].down[opposite] & TAVL_INDEX;
		const tavl_index_t c = arr[z].down[dir] & TAVL_INDEX;
		const tavl_index_t heavy = ((~arr[z].down[dir]) & TAVL_HEAVY);

		arr[x].down[opposite] = a;
		arr[x].down[dir] = b | heavy;

		arr[z].down[opposite] = x | heavy;
		arr[z].down[dir] = c;

		*ptr = ((*ptr) & TAVL_HEAVY) | z;
		return heavy ? 0 : 1;
	}

	/* Double-rotation */
	const tavl_index_t y = arr[z].down[opposite] & TAVL_INDEX;
	const tavl_index_t a = arr[x].down[opposite] & TAVL_INDEX;
	const tavl_index_t b = arr[y].down[opposite] & TAVL_INDEX;
	const tavl_index_t c = arr[y].down[dir] & TAVL_INDEX;
	const tavl_index_t d = arr[z].down[dir] & TAVL_INDEX;
	const tavl_index_t h1 = arr[y].down[opposite] & TAVL_HEAVY;
	const tavl_index_t h2 = arr[y].down[dir] & TAVL_HEAVY;

	arr[x].down[opposite] = a | h2;
	arr[x].down[dir] = b;

	arr[z].down[opposite] = c;
	arr[z].down[dir] = d | h1;

	arr[y].down[opposite] = x;
	arr[y].down[dir] = z;

	*ptr = ((*ptr) & TAVL_HEAVY) | y;
	return 1;
}

/* Maximum possible tree height.
 *
 * Let F(0) = 0, F(1) = 1, F(n) = F(n-1) + F(n-2)
 *
 * Then the number of nodes and the height satisfy:
 *
 *   n >= F(h+3) - 1
 *
 * Where h is the number of nodes in the path.
 */
#if TAVL_MAX_NODES < 46367
#define MAX_PATH 21
#else
#error "check MAX_PATH"
#endif

/* Propagate a height change up a path and fix unbalanced nodes along
 * the way.
 */
static void fix_path(struct tavl_node *arr, tavl_index_t *top,
		     const uint8_t *dirs, int l, int delta_h)
{
	tavl_index_t *ptrs[MAX_PATH];

	l--;
	ptrs[0] = top;
	for (int i = 0; i < l; i++)
		ptrs[i + 1] = &arr[(*ptrs[i]) & TAVL_INDEX].down[dirs[i]];

	while (l >= 0) {
		if (!delta_h)
			return;

		const tavl_index_t p = (*ptrs[l]) & TAVL_INDEX;
		const int dir = dirs[l];
		const int opposite = dir ^ 1;

		if (delta_h > 0) {
			if (arr[p].down[dir] & TAVL_HEAVY) {
				delta_h = 1 - rotate_and_fix(arr, ptrs[l], dir);
			} else if (arr[p].down[opposite] & TAVL_HEAVY) {
				arr[p].down[opposite] &= ~TAVL_HEAVY;
				delta_h = 0;
			} else {
				arr[p].down[dir] |= TAVL_HEAVY;
			}
		} else {
			if (arr[p].down[dir] & TAVL_HEAVY) {
				arr[p].down[dir] &= ~TAVL_HEAVY;
			} else if (arr[p].down[opposite] & TAVL_HEAVY) {
				delta_h = -rotate_and_fix(arr, ptrs[l],
					opposite);
			} else {
				arr[p].down[opposite] |= TAVL_HEAVY;
				delta_h = 0;
			}
		}

		l--;
	}
}

tavl_index_t tavl_find(const struct tavl_node *arr, tavl_index_t n,
		       tavl_key_t k)
{
	while (n != TAVL_NONE) {
		if (arr[n].key == k)
			return n;

		n = arr[n].down[k < arr[n].key ? 0 : 1] & TAVL_INDEX;
	}

	return TAVL_NONE;
}

tavl_index_t tavl_insert(struct tavl_node *arr, tavl_index_t *root,
			 tavl_index_t incoming)
{
	tavl_index_t *top = root;
	const tavl_key_t k = arr[incoming].key;
	uint8_t dir[MAX_PATH];
	int l = 0;

	for (;;) {
		const tavl_index_t n = (*root) & TAVL_INDEX;

		if (n == TAVL_NONE) {
			*root = ((*root) & TAVL_HEAVY) | incoming;
			arr[incoming].down[0] = TAVL_NONE;
			arr[incoming].down[1] = TAVL_NONE;
			fix_path(arr, top, dir, l, 1);
			break;
		}

		if (arr[n].key == k) {
			*root = ((*root) & TAVL_HEAVY) | incoming;
			arr[incoming].down[0] = arr[n].down[0];
			arr[incoming].down[1] = arr[n].down[1];
			return n;
		}

		dir[l] = k < arr[n].key ? 0 : 1;
		root = &arr[n].down[dir[l]];
		l++;
	}

	return TAVL_NONE;
}

tavl_index_t tavl_delete(struct tavl_node *arr, tavl_index_t *root,
			 tavl_key_t k)
{
	tavl_index_t *top = root;
	uint8_t dir[MAX_PATH];
	tavl_index_t n = TAVL_NONE;
	int l = 0;

	/* Find the target node, if it exists */
	for (;;) {
		n = (*root) & TAVL_INDEX;

		if (n == TAVL_NONE)
			return TAVL_NONE;

		if (arr[n].key == k)
			break;

		dir[l] = k < arr[n].key ? 0 : 1;
		root = &arr[n].down[dir[l]];
		l++;
	}

	/* Now we found the node to delete (*root), but it might have
	 * children.
	 */
	if ((arr[n].down[0] == TAVL_NONE) &&
	    (arr[n].down[1] == TAVL_NONE)) {
		/* Easy case: no children */
		*root = ((*root) & TAVL_HEAVY) | TAVL_NONE;
		goto done;
	}

	if (arr[n].down[0] == TAVL_NONE) {
		/* Easy case: only right child */
		*root = ((*root) & TAVL_HEAVY) |
			(arr[n].down[1] & TAVL_INDEX);
		goto done;
	}

	if (arr[n].down[1] == TAVL_NONE) {
		/* Easy case: only left child */
		*root = ((*root) & TAVL_HEAVY) |
			(arr[n].down[0] & TAVL_INDEX);
		goto done;
	}

	tavl_index_t *old_root = root;
	const tavl_index_t old_n = n;
	const tavl_index_t old_n_0 = arr[old_n].down[0];
	const tavl_index_t old_n_1 = arr[old_n].down[1];

	/* Descend to the top of the left subtree */
	dir[l++] = 0;
	n = old_n_0 & TAVL_INDEX;

	/* No right child? */
	if (arr[n].down[1] == TAVL_NONE) {
		const tavl_index_t nl = arr[n].down[0] & TAVL_INDEX;

		arr[n].down[0] = (old_n_0 & TAVL_HEAVY) | nl;
		arr[n].down[1] = old_n_1;
		(*old_root) = ((*old_root) & TAVL_HEAVY) | n;

		n = old_n;
		goto done;
	}

	/* Go as far right as we can */
	while (arr[n].down[1] != TAVL_NONE) {
		dir[l++] = 1;
		root = &arr[n].down[1];
		n = (*root) & TAVL_INDEX;
	}

	/* Swap nodes and then hoist left child, if any */
	const tavl_index_t nl = arr[n].down[0] & TAVL_INDEX;

	arr[n].down[0] = old_n_0;
	arr[n].down[1] = old_n_1;
	(*old_root) = ((*old_root) & TAVL_HEAVY) | n;

	(*root) = ((*root) & TAVL_HEAVY) | nl;
	n = old_n;

done:
	fix_path(arr, top, dir, l, -1);
	return n;
}