/*
 * avl_tree.c
 *
 * Intrusive, nonrecursive AVL tree data structure (self-balancing binary search
 * tree), implementation file.
 *
 * Author:  Eric Biggers
 * Year:    2014
 *
 * This file is placed into the public domain.  You can do whatever you want
 * with it.
 */

#include <wimlib/avl_tree.h>

/* Starts an in-order traversal of the tree: returns the least-valued node, or
 * NULL if the tree is empty.  */
struct avl_tree_node *
avl_tree_first_in_order(const struct avl_tree_node *root)
{
	const struct avl_tree_node *first = root;

	if (first)
		while (first->left)
			first = first->left;
	return (struct avl_tree_node *)first;
}

/* Continues an in-order traversal of the tree: returns the next-greatest-valued
 * node, or NULL if there is none.  */
struct avl_tree_node *
avl_tree_next_in_order(const struct avl_tree_node *prev)
{
	const struct avl_tree_node *next;

	if (prev->right)
		for (next = prev->right;
		     next->left;
		     next = next->left)
			;
	else
		for (next = avl_get_parent(prev);
		     next && prev == next->right;
		     prev = next, next = avl_get_parent(next))
			;
	return (struct avl_tree_node *)next;
}

/* Starts a postorder traversal of the tree.  */
struct avl_tree_node *
avl_tree_first_in_postorder(const struct avl_tree_node *root)
{
	const struct avl_tree_node *first = root;

	if (first)
		while (first->left || first->right)
			first = first->left ? first->left : first->right;

	return (struct avl_tree_node *)first;
}

/* Continues a postorder traversal of the tree.  @prev will not be deferenced as
 * it's allowed that its memory has been freed; @prev_parent must be its saved
 * parent node.  Returns NULL if there are no more nodes (i.e. @prev was the
 * root of the tree).  */
struct avl_tree_node *
avl_tree_next_in_postorder(const struct avl_tree_node *prev,
			   const struct avl_tree_node *prev_parent)
{
	const struct avl_tree_node *next = prev_parent;

	if (next && prev == next->left && next->right)
		for (next = next->right;
		     next->left || next->right;
		     next = next->left ? next->left : next->right)
			;
	return (struct avl_tree_node *)next;
}

/* Get the left child (sign < 0) or the right child (sign > 0)
 * Note: for all calls of this, 'sign' is constant at compilation time and the
 * compiler can remove the conditional.  */
static AVL_INLINE struct avl_tree_node *
avl_get_child(const struct avl_tree_node *parent, int sign)
{
	if (sign < 0)
		return parent->left;
	else
		return parent->right;
}

/* Set the left child (sign < 0) or the right child (sign > 0)
 * Note: for all calls of this, 'sign' is constant at compilation time and the
 * compiler can remove the conditional.  */
static AVL_INLINE void
avl_set_child(struct avl_tree_node *parent, int sign,
	      struct avl_tree_node *child)
{
	if (sign < 0)
		parent->left = child;
	else
		parent->right = child;
}

/* Sets the parent of the AVL tree @node to @parent.  */
static AVL_INLINE void
avl_set_parent(struct avl_tree_node *node, struct avl_tree_node *parent)
{
	node->parent_balance =
		(node->parent_balance & 3) | (uintptr_t)parent;
}

/* Returns the balance factor of the specified AVL tree node --- that is, the
 * height of its right subtree minus the height of its left subtree.  */
static AVL_INLINE int
avl_get_balance_factor(const struct avl_tree_node *node)
{
	return (int)(node->parent_balance & 3) - 1;
}

/* Sets the balance factor of the specified AVL tree node.  The caller MUST
 * ensure that this number is valid and maintains the AVL tree invariants.  */
static AVL_INLINE void
avl_set_balance_factor(struct avl_tree_node *node, int balance_factor)
{
	node->parent_balance =
		(node->parent_balance & ~3) | (balance_factor + 1);
}

/* Increments the balance factor of the specified AVL tree @node by the
 * specified @amount.  The caller MUST ensure that this number is valid and
 * maintains the AVL tree invariants.  */
static AVL_INLINE void
avl_adjust_balance_factor(struct avl_tree_node *node, int amount)
{
	node->parent_balance += amount;
}

/*
 * Template for performing a rotation ---
 *
 * Clockwise/Right (sign > 0):
 *
 *           X             X
 *           |             |
 *           A             B
 *          / \           / \
 *         B   C   =>    D   A
 *        / \               / \
 *       D   E             E   C
 *
 * Counterclockwise/Left (sign < 0)- same, except left and right are reversed:
 *
 *           X             X
 *           |             |
 *           A             B
 *          / \           / \
 *         C   B   =>    A   D
 *            / \       / \
 *           E   D     C   E
 *
 * This updates pointers but not balance factors!
 */
static AVL_INLINE void
avl_rotate(struct avl_tree_node * const A, const int sign)
{
	struct avl_tree_node * const B = avl_get_child(A, -sign);
	struct avl_tree_node * const C = avl_get_child(A, +sign);
	struct avl_tree_node * const E = avl_get_child(B, +sign);
	struct avl_tree_node * const X = avl_get_parent(A);

	avl_set_child(A, -sign, E);
	avl_set_child(A, +sign, C);
	avl_set_parent(A, B);

	avl_set_child(B, +sign, A);
	avl_set_parent(B, X);

	if (E)
		avl_set_parent(E, A);

	if (X) {
		if (avl_get_child(X, +sign) == A)
			avl_set_child(X, +sign, B);
		else
			avl_set_child(X, -sign, B);
	}
}

/* See description in avl_handle_subtree_growth()  */
static AVL_INLINE struct avl_tree_node *
avl_do_double_rotate(struct avl_tree_node * const B,
		     struct avl_tree_node * const A, const int sign)
{

	struct avl_tree_node * const E = avl_get_child(B, -sign);
	const int e = avl_get_balance_factor(E);

	avl_rotate(B, +sign);
	avl_rotate(A, -sign);
	avl_set_balance_factor(A, ((sign * e <= 0) ? 0 : -e));
	avl_set_balance_factor(B, ((sign * e >= 0) ? 0 : -e));
	avl_set_balance_factor(E, 0);

	return E;
}

static AVL_INLINE bool
avl_handle_subtree_growth(struct avl_tree_node * const node,
			  struct avl_tree_node * const parent,
			  const int sign)
{
	/* Height of @node subtree of @parent has increased by 1.
	 * Adjust @parent's balance factor and check whether rotations need to
	 * be done.  */

	const int old_balance_factor = avl_get_balance_factor(parent);
	const int new_balance_factor = old_balance_factor + sign;

	if (old_balance_factor == 0) {
		/* Parent has increased in height but is still sufficiently
		 * balanced.  Continue up the tree.  */
		avl_adjust_balance_factor(parent, sign);
		return false;
	}

	if (new_balance_factor == 0) {
		/* Parent is balanced; nothing more to to.  */
		avl_adjust_balance_factor(parent, sign);
		return true;
	}

	/* FROM THIS POINT ONWARDS THE COMMENTS ASSUME sign < 0.
	 * The other case is symmetric --- that is, the rotations done are the
	 * the mirror images, all the balance factors are inverted, and left and
	 * right pointers are otherwise reversed.  */

	/* Parent is left-heavy (balance_factor == -2).  */

	if (sign * avl_get_balance_factor(node) > 0) {

		/* Child node (@node, also B below) is also left-heavy.
		 * It must have balance_factor == -1.
		 * Do a clockwise ("right") rotation rooted at
		 * @parent (A below):
		 *
		 *           A              B
		 *          / \           /   \
		 *         B   C   =>    D     A
		 *        / \           / \   / \
		 *       D   E         F   G E   C
		 *      / \
		 *     F   G
		 *
		 * Before the rotation:
		 *	balance(A) = -2
		 *	balance(B) = -1
		 * Let x = height(C).  Then:
		 *	height(B) = x + 2
		 *	height(D) = x + 1
		 *	height(E) = x
		 *	max(height(F), height(G)) = x.
		 *
		 * After the rotation:
		 *	height(D) = max(height(F), height(G)) + 1
		 *		  = x + 1
		 *	height(A) = max(height(E), height(C)) + 1
		 *		  = max(x, x) + 1 = x + 1
		 *	balance(B) = 0
		 *	balance(A) = 0
		 */

		avl_rotate(parent, -sign);

		avl_set_balance_factor(node, 0);   /* B */
		avl_set_balance_factor(parent, 0); /* A */
	} else {
		/* Child node (@node, also B below) is right-heavy.
		 * It must have balance_factor == +1.
		 * Do a counterclockwise ("left") rotation rooted at child node
		 * (B below), then a clockwise ("right") rotation rooted at
		 * parent node (A below).
		 *
		 *           A             A           E
		 *          / \           / \        /   \
		 *         B   C   =>    E   C  =>  B     A
		 *        / \           / \        / \   / \
		 *       D   E         B   G      D   F G   C
		 *          / \       / \
		 *         F   G     D   F
		 *
		 * Before the rotation:
		 *	balance(A) = -2
		 *	balance(B) = +1
		 * Let x = height(C).  Then:
		 *	height(A)
		 *	height(B) = x + 2
		 *	height(E) = x + 1
		 *	height(D) = x
		 *	max(height(F), height(G)) = x
		 *
		 * After both rotations:
		 *	height(A) = max(height(G), height(C)) + 1
		 *		  = x + 1
		 *	balance(A) = balance(E{orig}) >= 0 ? 0 : -balance(E{orig})
		 *	height(B) = max(height(D), height(F)) + 1
		 *		  = x + 1
		 *	balance(B) = balance(E{orig} <= 0) ? 0 : -balance(E{orig})
		 *
		 *	height(E) = x + 2
		 *	balance(E) = 0
		 */
		avl_do_double_rotate(node, parent, sign);
	}
	return true;
}

/* Rebalance the tree after insertion of the specified node.  */
void
avl_tree_rebalance_after_insert(struct avl_tree_node **root_ptr,
				struct avl_tree_node *inserted)
{
	struct avl_tree_node *node, *parent;
	bool done = false;

	inserted->left = NULL;
	inserted->right = NULL;

	for (node = inserted, parent = avl_get_parent(node);
	     parent && !done;
	     node = parent, parent = avl_get_parent(parent))
	{
		/* Height of @node subtree has increased by 1  */

		if (node == parent->left)
			done = avl_handle_subtree_growth(node, parent, -1);
		else
			done = avl_handle_subtree_growth(node, parent, +1);
	}
	/* Due to rotations, *root_ptr may no longer be the root of the tree  */
	while (avl_get_parent(*root_ptr))
		*root_ptr = avl_get_parent(*root_ptr);
}

static AVL_INLINE struct avl_tree_node *
avl_handle_subtree_shrink(struct avl_tree_node *parent,
			  const int sign,
			  bool * const left_deleted_ret)
{
	struct avl_tree_node *node;

	const int old_balance_factor = avl_get_balance_factor(parent);
	const int new_balance_factor = old_balance_factor + sign;

	if (old_balance_factor == 0) {
		/* Prior to the deletion, the subtree rooted at
		 * @parent was perfectly balanced.  It's now
		 * unbalanced by 1, but that's okay and its height
		 * hasn't changed.  Nothing more to do.  */
		avl_adjust_balance_factor(parent, sign);
		return NULL;
	} else if (new_balance_factor == 0) {
		/* The subtree rooted at @parent is now perfectly
		 * balanced, whereas before the deletion it was
		 * unbalanced by 1.  Its height must have decreased
		 * by 1.  No rotation is needed at this location,
		 * but continue up the tree.  */
		avl_adjust_balance_factor(parent, sign);
		node = parent;
	} else {
		/* The subtree rooted at @parent is now significantly
		 * unbalanced (by 2 in some direction).  */
		node = avl_get_child(parent, sign);

		/* The rotations below are similar to those done during
		 * insertion.  The only new case is the one where the
		 * child node has a balance factor of 0.  */

		if (sign * avl_get_balance_factor(node) >= 0) {
			avl_rotate(parent, -sign);

			if (avl_get_balance_factor(node) == 0) {
				avl_set_balance_factor(node,   -sign);
				avl_set_balance_factor(parent, +sign);
				/* Height is unchanged; nothing more to do.  */
				return NULL;
			} else {
				avl_set_balance_factor(node, 0);
				avl_set_balance_factor(parent, 0);
			}
		} else {
			node = avl_do_double_rotate(node, parent, sign);
		}
	}
	parent = avl_get_parent(node);
	if (parent)
		*left_deleted_ret = (node == parent->left);
	return parent;
}

/* Swaps node X, which must have 2 children, with its in-order successor.
 * This adjusts all the necessary pointers and balance factors.  */
static AVL_INLINE void
avl_tree_swap_with_successor(struct avl_tree_node * const X)
{
	struct avl_tree_node * const Y		= avl_tree_next_in_order(X);
	struct avl_tree_node * const Y_right	= Y->right;
	struct avl_tree_node * const Y_parent   = avl_get_parent(Y);
	struct avl_tree_node * const X_parent	= avl_get_parent(X);
	const int Y_balance_factor		= avl_get_balance_factor(Y);

	/* Set child pointer to Y when placed in X's position  */
	if (X_parent) {
		if (X == X_parent->left)
			X_parent->left = Y;
		else
			X_parent->right = Y;
	}

	/* Set child pointer to X when placed in Y's position.  Also set
	 * the right pointer of Y (it may point to X if the nodes are
	 * adjacent)  */
	if (Y_parent != X) {
		if (Y == Y_parent->left)
			Y_parent->left = X;
		else
			Y_parent->right = X;
		Y->right = X->right;
	} else {
		Y->right = X;
	}

	/* Set parent pointer to X when placed in Y's position  */
	if (Y_right)
		avl_set_parent(Y_right, X);

	/* Note: Y has no left child since it is the in-order sucessor of X.  */

	/* Set parent pointers to Y when placed in X's position, as well as X's
	 * parent when placed in Y's position.  */
	avl_set_parent(X->left, Y);
	if (X->right != Y) {
		avl_set_parent(X->right, Y);
		avl_set_parent(X, Y_parent);
	} else {
		avl_set_parent(X, Y);
	}

	avl_set_parent(Y, X_parent);
	Y->left = X->left;
	avl_set_balance_factor(Y, avl_get_balance_factor(X));

	X->left = NULL;
	X->right = Y_right;
	avl_set_balance_factor(X, Y_balance_factor);
}

/* Removes the specified @node from the AVL tree whose root is pointed to by
 * @root_ptr.
 *
 * This *only* unlinks the node and rebalances the tree; it does not free any
 * memory or anything.  */
void
avl_tree_remove(struct avl_tree_node **root_ptr, struct avl_tree_node *node)
{
	struct avl_tree_node *child, *parent;
	bool left_deleted;

	if (node->left && node->right)
		avl_tree_swap_with_successor(node);

	/* Unlink @node  */
	child = node->left ? node->left : node->right;
	parent = avl_get_parent(node);
	if (parent) {
		if (node == parent->left) {
			parent->left = child;
			left_deleted = true;
		} else {
			parent->right = child;
			left_deleted = false;
		}
	} else {
		*root_ptr = child;
	}
	if (child)
		avl_set_parent(child, parent);

	/* Rebalance the tree  */
	while (parent) {
		if (left_deleted)
			parent = avl_handle_subtree_shrink(parent, +1, &left_deleted);
		else
			parent = avl_handle_subtree_shrink(parent, -1, &left_deleted);
	}

	/* Due to rotations, *root_ptr may no longer point to the root of the
	 * tree.  Fix it.  */
	if (*root_ptr)
		while (avl_get_parent(*root_ptr))
			*root_ptr = avl_get_parent(*root_ptr);
}