/*
 * Copyright (c) 2008 - 2010 Jason Evans <jasone@FreeBSD.org>
 * Copyright (c) 2011 Your File System Inc.
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR `AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
 * IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
 * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
 * DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
 * THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 * THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

/* Left-leaning rbtree implementation. Originally derived from the FreeBSD
 * CPP macro definitions by Jason Evans, but extensively modified to
 *     *) Be function, rather than macro based
 *     *) Use parent pointers, rather than calling an embeded comparison fn
 *     *) Use 'NULL' to represent empty nodes, rather than a magic pointer
 */

#include "config.h"
#include <stdlib.h>
#include "misc/rbtree.h"

/* Update the parent pointers for a node which is being replaced.
 *
 * If the original node had no parent, then it was the root of the tree,
 * and the replacement is too.
 * Otherwise, the original node must have been either the left or right
 * child of its parent, so update the left (or right) pointer to point
 * to the replacement as appropriate.
 */

static inline void update_parent_ptr(struct opr_rbtree *head,
				     struct opr_rbtree_node *old,
				     struct opr_rbtree_node *replacement)
{
	if (old->parent) {
		if (old->parent->left == old)
			old->parent->left = replacement;
		else
			old->parent->right = replacement;
	} else
		head->root = replacement;
}

/* initialize a tree, cmpf may be NULL--if so, do not call internal
 * iterating functions (opr_rbtree_lookup, opr_rbtree_insert) */
void opr_rbtree_init(struct opr_rbtree *head, opr_rbtree_cmpf_t cmpf)
{
	head->cmpf = cmpf;
	head->root = NULL;
	head->size = 0;
	head->gen = 0;
}

struct opr_rbtree_node *opr_rbtree_first(struct opr_rbtree *head)
{
	struct opr_rbtree_node *node;

	node = head->root;
	if (node == NULL)
		return node;

	while (node->left != NULL)
		node = node->left;

	return node;
}

struct opr_rbtree_node *opr_rbtree_last(struct opr_rbtree *head)
{
	struct opr_rbtree_node *node;

	node = head->root;

	if (node == NULL)
		return node;

	while (node->right != NULL)
		node = node->right;

	return node;
}

struct opr_rbtree_node *opr_rbtree_next(struct opr_rbtree_node *node)
{
	struct opr_rbtree_node *parent;

	/* Where there is a right child, the next node is to the right, then
	 * left as far as we can go */
	if (node->right != NULL) {
		node = node->right;
		while (node->left != NULL)
			node = node->left;

		return node;
	}

	/* If there is no right hand child, then the next node is above us.
	 * Whenever our ancestor is a right-hand child, the next node is
	 * further up. When it is a left-hand child, it is our next node
	 */
	while ((parent = node->parent) != NULL && node == parent->right)
		node = parent;

	return parent;
}

struct opr_rbtree_node *opr_rbtree_prev(struct opr_rbtree_node *node)
{
	struct opr_rbtree_node *parent;

	if (node->left != NULL) {
		node = node->left;
		while (node->right != NULL)
			node = node->right;

		return node;
	}

	/* Same ancestor logic as for 'next', but in reverse */
	while ((parent = node->parent) != NULL && node == parent->left)
		node = parent;

	return parent;
}

static inline void rotateright(struct opr_rbtree *head,
			       struct opr_rbtree_node *node)
{
	struct opr_rbtree_node *left = node->left;

	node->left = left->right;
	if (left->right)
		left->right->parent = node;

	left->right = node;
	left->parent = node->parent;

	update_parent_ptr(head, node, left);

	node->parent = left;
}

static inline void rotateleft(struct opr_rbtree *head,
			      struct opr_rbtree_node *node)
{
	struct opr_rbtree_node *right = node->right;

	node->right = right->left;
	if (right->left)
		right->left->parent = node;

	right->left = node;
	right->parent = node->parent;

	update_parent_ptr(head, node, right);

	node->parent = right;
}

static inline void swapnode(struct opr_rbtree_node **a,
			    struct opr_rbtree_node **b)
{
	struct opr_rbtree_node *tmp;

	tmp = *a;
	*a = *b;
	*b = tmp;
}

static void insert_recolour(struct opr_rbtree *head,
			    struct opr_rbtree_node *node)
{
	struct opr_rbtree_node *parent, *gramps;

	while ((parent = node->parent) && parent->red) {
		gramps = parent->parent;

		if (parent == gramps->left) {
			struct opr_rbtree_node *uncle = gramps->right;

			if (uncle && uncle->red) {
				uncle->red = 0;
				parent->red = 0;
				gramps->red = 1;
				node = gramps;
				continue;
			}

			if (parent->right == node) {
				rotateleft(head, parent);
				swapnode(&parent, &node);
			}

			parent->red = 0;
			gramps->red = 1;
			rotateright(head, gramps);
		} else {
			struct opr_rbtree_node *uncle = gramps->left;

			if (uncle && uncle->red) {
				uncle->red = 0;
				parent->red = 0;
				gramps->red = 1;
				node = gramps;
				continue;
			}

			if (parent->left == node) {
				rotateright(head, parent);
				swapnode(&parent, &node);
			}

			parent->red = 0;
			gramps->red = 1;
			rotateleft(head, gramps);
		}
	}

	head->root->red = 0;
}

void opr_rbtree_insert_at(struct opr_rbtree *head,
			  struct opr_rbtree_node *parent,
			  struct opr_rbtree_node **childptr,
			  struct opr_rbtree_node *node)
{
	/* Link node 'node' into the tree at position 'parent', using either the
	 * left or right pointers */

	node->parent = parent;
	node->left = node->right = NULL;
	node->red = 1;
	node->gen = (uint32_t)(++(head->gen));
	if (!node->gen) {
		/* truncated to zero */
		node->gen = (uint32_t)(++(head->gen));
	}
	*childptr = node;

	/* Rebalance the tree for the newly inserted node */
	insert_recolour(head, node);
}

/* Try to insert node on tree head.  If an existing node compares
 * equal to node, return it--the caller will deal.  Else, return
 * NULL. Do NOT call this if your iteration is external */
struct opr_rbtree_node *opr_rbtree_insert(struct opr_rbtree *head,
					  struct opr_rbtree_node *node)
{
	struct opr_rbtree_node **iter, *parent = NULL;

	iter = &head->root;
	while (*iter) {
		parent = *iter;
		switch (head->cmpf(node, parent)) {
		case -1:
			iter = &parent->left;
			break;
		case 1:
			iter = &parent->right;
			break;
		default:
			return (parent);
			break;
		}
	}			/* while */

	opr_rbtree_insert_at(head, parent, iter, node);
	(head->size)++;

	return (NULL);
}

/* search for node in tree--do NOT call this if your iteration
 * is external  */
struct opr_rbtree_node *opr_rbtree_lookup(struct opr_rbtree *head,
					  struct opr_rbtree_node *node)
{
	struct opr_rbtree_node *iter, *parent = NULL;

	iter = head->root;
	while (iter) {
		parent = iter;
		switch (head->cmpf(node, parent)) {
		case -1:
			iter = parent->left;
			break;
		case 1:
			iter = parent->right;
			break;
		default:
			return (parent);
			break;
		}
	}			/* while */

	return (NULL);
}

static void remove_recolour(struct opr_rbtree *head,
			    struct opr_rbtree_node *parent,
			    struct opr_rbtree_node *node)
{
	struct opr_rbtree_node *other;

	while ((node == NULL || !node->red) && node != head->root) {
		if (parent->left == node) {
			other = parent->right;
			if (other->red) {
				other->red = 0;
				parent->red = 1;
				rotateleft(head, parent);
				other = parent->right;
			}
			if ((other->left == NULL || !other->left->red)
			    && (other->right == NULL || !other->right->red)) {
				other->red = 1;
				node = parent;
				parent = node->parent;
			} else {
				if ((other->right == NULL) ||
				    (!other->right->red)) {
					other->left->red = 0;
					other->red = 1;
					rotateright(head, other);
					other = parent->right;
				}
				other->red = parent->red;
				parent->red = 0;
				other->right->red = 0;
				rotateleft(head, parent);
				node = head->root;
				break;
			}
		} else {
			other = parent->left;
			if (other->red) {
				other->red = 0;
				parent->red = 1;
				rotateright(head, parent);
				other = parent->left;
			}
			if ((other->left == NULL || !other->left->red)
			    && (other->right == NULL || !other->right->red)) {
				other->red = 1;
				node = parent;
				parent = node->parent;
			} else {
				if (other->left == NULL || !other->left->red) {
					other->right->red = 0;
					other->red = 1;
					rotateleft(head, other);
					other = parent->left;
				}
				other->red = parent->red;
				parent->red = 0;
				other->left->red = 0;
				rotateright(head, parent);
				node = head->root;
				break;
			}
		}
	}
	if (node)
		node->red = 0;
}

void opr_rbtree_remove(struct opr_rbtree *head, struct opr_rbtree_node *node)
{
	struct opr_rbtree_node *child = NULL, *parent;
	int red;

	if (!node->gen) {
		/* not in tree */
		return;
	}
	node->gen = 0;

	if (node->left == NULL && node->right == NULL) {
		/* A node with no non-leaf children */
		update_parent_ptr(head, node, NULL);

		if (!node->red)
			remove_recolour(head, node->parent, NULL);

		goto done;
	}

	if (node->left != NULL && node->right != NULL) {
		/* A node with two children.
		 *
		 * Move the next node in the tree (which will be a leaf node)
		 * onto our tree current position, then rebalance as required
		 */
		struct opr_rbtree_node *old, *left;

		old = node;

		/* Set node to the next node in the tree from the current
		 * position, where the next node is the left-most leaf node
		 * in our right child */
		node = node->right;
		while ((left = node->left) != NULL)
			node = left;

		/* Move 'node' into the position occupied by 'old', which is
		 * being removed */
		update_parent_ptr(head, old, node);

		child = node->right;
		parent = node->parent;
		red = node->red;

		/* As we're logically just copying the value, must preserve the
		 * old node's colour */
		node->red = old->red;

		/* ... and the old node's linkage */
		if (parent == old)
			parent = node;
		else {
			if (child)
				child->parent = parent;
			parent->left = child;

			node->right = old->right;
			old->right->parent = node;
		}

		node->parent = old->parent;
		node->left = old->left;
		old->left->parent = node;

		/* If the node being removed was black, then we must recolour
		 * tree to maintain balance */
		if (!red)
			remove_recolour(head, parent, child);

		goto done;
	}

	/* Only remaining option - node with a single child */

	if (node->left == NULL)
		child = node->right;
	else if (node->right == NULL)
		child = node->left;

	child->parent = node->parent;

	update_parent_ptr(head, node, child);

	if (!node->red)
		remove_recolour(head, node->parent, child);

 done:
	(head->size)--;
	(head->gen)++;
	return;
}

void opr_rbtree_replace(struct opr_rbtree *head, struct opr_rbtree_node *old,
			struct opr_rbtree_node *replacement)
{
	/* Update our parent's pointer to us */
	update_parent_ptr(head, old, replacement);

	/* And our children's pointers to us */
	if (old->left)
		old->left->parent = replacement;
	if (old->right)
		old->right->parent = replacement;

	/* Copy over parent, left, right and colour */
	*replacement = *old;
}