#include "rbtree.h"
#include "assert.h"
#include "string.h"
#define BLACK 0
#define RED 1
#define LEFT 0
#define left children[LEFT]
#define RIGHT 1
#define right children[RIGHT]
#define CHILD_DIR(c) (c->parent->children[LEFT] == c ? LEFT : RIGHT)
/**Given a root and a direction, rotate the subtree such that the order of the
* elements are unchanged.
* The opposite direction child of the given root is treated as the pivot point
* The dir-most child of the pivot becomes the opposite dir child of the root.
* The pivot becomes the root and vice-versa.
* @param tree to do rotation on.
* @param root of rotation.
* @param dir to rotate by.
* @return new root node.*/
rbnode_t*
s_node_rotate(rbtree_t *tree, rbnode_t *root, bool dir)
{
assert(root != NULL);
rbnode_t *gparent = root->parent;
rbnode_t *sibling = root->children[1-dir];
rbnode_t *child = NULL;
assert(sibling != NULL);
child = sibling->children[dir];
/* Opposite child of root is dir-most child of sibling.*/
root->children[1-dir] = child;
if(child != NULL) child->parent = root;
/* Child of sibling is the root.*/
sibling->children[ dir] = root;
root->parent = sibling;
/* Parent of sibling is the grandparent. */
sibling->parent = gparent;
if(gparent != NULL)
gparent->children[CHILD_DIR(root)] = sibling;
else
tree->root = sibling;
return sibling;
}
/**Search the rbtree for a node that matches the key.
* If no such node is found, returns NULL.
* @param tree to search through.
* @param key to search for.
* @return found node.*/
rbnode_t*
s_find(rbtree_t *tree, intmax_t key)
{
rbnode_t *node = tree->root;
while(node != NULL && node->key != key) {
node = node->children[node->key > key];
}
return node;
}
rbnode_t*
s_closest(rbtree_t *tree, intmax_t key)
{
rbnode_t *node = tree->root;
while(node != NULL && node->key != key) {
rbnode_t *child = node->children[node->key > key];
if(child == NULL) return node;
node = child;
}
return node;
}
/**Find the left-most node for this subtree.
* @param node root of the subtree to search through.
* @return left-most node.*/
rbnode_t*
s_node_min(rbnode_t *node)
{
while(node->left != NULL) node = node->left;
return node;
}
/**Find the successor for the specified node.
* The successor is defined as the left-most node of the right sub-tree.
* (the smallest node larger than the root)
* @param node to search for.
* @return the successor.*/
rbnode_t*
s_node_successor(rbnode_t *node)
{
if(node->right != NULL)
return s_node_min(node->right);
/* If the node has no left child (meaning that the node is larger than it's
* parents), we cycle up through the RHS of the parent tree (such that each
* parent is smaller than the child) until we reach a point by which the
* child node is smaller than the parent.*/
rbnode_t *parent = node->parent;
while(parent != NULL && node == parent->right) {
node = parent;
parent = node->parent;
}
return parent;
}
/**Copy the data from one node onto another, keeping parent and child relations.
* @prarm tree for metadata purposes.
* @param dest destination node.
* @param src source node.*/
void
s_node_copy_value(rbtree_t *tree, rbnode_t *dest, rbnode_t *src)
{
dest->key = src->key;
memcpy(dest->data, src->data, tree->objw);
}
/**Perform tree adjustments after insertion to rebalance the rbtree.
* @param tree to rebalance.
* @param node that was inserted.
* @param parent of the node.
* @param dir that the node was inserted in.*/
void
s_node_insert_fix(
rbtree_t *tree,
rbnode_t *node,
rbnode_t *parent,
bool dir)
{
rbnode_t *gparent;
rbnode_t *uncle;
while((parent = node->parent) != NULL) {
if(parent->color == BLACK) return; /* Tree is valid. no fixes needed.*/
/*Parent is known to be red.*/
if((gparent = parent->parent) == NULL) {
/*Parent is red and the root. We can just fix it and return.*/
parent->color = BLACK;
return;
}
/* Parent is red and has a parent.
* We now need to fix the parent.*/
dir = CHILD_DIR(parent);
uncle = gparent->children[1-dir];
if(uncle == NULL || uncle->color == BLACK) {
/* Parent is red, but uncle (opposite child of gparent) is black.
* We need to rotate the gparent node such that the parent takes
* its place. However, if the current node is an inner child of
* gparent, we need to rotate N into P.*/
if(node == parent->children[1-dir]) {
/* Node is an inner child. We need to rotate N and P first.*/
s_node_rotate(tree, parent, dir);
node = parent;
parent = gparent->children[dir];
}
/* N is not an inner child, we can rotate the tree properly.*/
s_node_rotate(tree, gparent, 1-dir);
parent->color = BLACK;
gparent->color = RED;
return;
}
}
}
/**Insert a node into the rbtree on the parent and direction given.
* @param tree to insert into.
* @param node to insert.
* @param parent node to insert onto.
* @param dir direction to insert node.*/
void
s_node_insert(rbtree_t *tree, rbnode_t *node, rbnode_t *parent, bool dir)
{
node->color = RED;
node->left = NULL;
node->right = NULL;
node->parent = parent;
if(parent == NULL) {
tree->root = node;
return;
}
parent->children[dir] = node;
s_node_insert_fix(tree, node, parent, dir);
}
void
s_node_remove_fixup(rbtree_t *tree, rbnode_t *node)
{
bool dir;
rbnode_t *parent = node->parent;
rbnode_t *sibling;
rbnode_t *cnephew;
rbnode_t *dnephew;
dir = CHILD_DIR(node);
parent->children[dir] = NULL;
while((parent = node->parent) != NULL)
{
dir = CHILD_DIR(node);
sibling = parent->children[1-dir];
cnephew = parent->children[ dir];
dnephew = parent->children[1-dir];
if(sibling->color == RED)
goto case3;
if(dnephew != NULL && dnephew->color == RED)
goto case6;
if(cnephew != NULL && cnephew->color == RED)
goto case5;
if(parent->color == RED)
goto case4;
/* The parent, both nephews, and sibling are black. We need to paint
* the sibling black and move up a level.*/
sibling->color = RED;
node = parent;
}
return;
case3:
/* The sibling is red, which means both nephews and the parent are black.
* We rotate the subtree so that the parent moves to become the siblings
* left child. The tree is then repainted.*/
s_node_rotate(tree, parent, dir);
parent->color = RED;
sibling->color = BLACK;
/* Now we are working on the sibling subtree.*/
sibling = cnephew;
dnephew = sibling->children[1-dir];
if(dnephew != NULL && dnephew->color == BLACK)
goto case6;
cnephew = sibling->children[dir];
if(cnephew != NULL && cnephew->color == RED)
goto case5;
case4:
sibling->color = RED;
parent->color = BLACK;
return;
case5:
s_node_rotate(tree, sibling, 1-dir);
sibling->color = RED;
cnephew->color = BLACK;
dnephew = sibling;
sibling = cnephew;
case6:
s_node_rotate(tree, parent, dir);
sibling->color = parent->color;
parent->color = BLACK;
dnephew->color = BLACK;
}
void
s_node_remove(rbtree_t *tree, rbnode_t *node)
{
rbnode_t *parent = node->parent;
if(node->left != NULL && node->right != NULL) {
/* Node is a full tree. We can replace the node with it's successor to
* keep the balance.
* Needs to recurse if the successor has a right-most child, which is
* dealt with in the next case.*/
rbnode_t *successor = s_node_successor(node);
s_node_copy_value(tree, node, successor);
s_node_remove(tree, successor);
return;
} else if(node->left == NULL && node->right == NULL) {
if(parent == NULL) {
tree->root = NULL;
kfree(node);
return;
}
if(node->color == RED) {
parent->children[CHILD_DIR(node)] = NULL;
kfree(node);
return;
}
s_node_remove_fixup(tree, node);
} else {
/* Node has a single child. Because the child must be red as per spec,
* and this node must be black, we can simply replace the node
* with it's child and color it black.*/
rbnode_t *child = node->left == NULL ? node->right : node->left;
memcpy(node, child, sizeof(rbnode_t) + tree->objw);
node->parent = parent;
node->color = BLACK;
kfree(child);
return;
}
}
void
__rbtree_new(rbtree_t *tree, size_t objw)
{
*tree = (rbtree_t) {
.root = NULL,
.objw = objw,
};
}
void*
rbtree_find(rbtree_t *tree, intmax_t key)
{
rbnode_t *found = s_find(tree, key);
if(found == NULL) return NULL;
return found->data;
}
rbnode_t*
s_node_new(rbtree_t *tree, intmax_t key, void *data)
{
rbnode_t *node = kmalloc(sizeof(rbnode_t) + tree->objw);
if(data != NULL) memcpy(node->data, data, tree->objw);
node->key = key;
return node;
}
void*
rbtree_insert(rbtree_t *tree, intmax_t key, void *data)
{
assert(tree != NULL);
rbnode_t *node = s_closest(tree, key);
if(node != NULL && node->key == key) {
memcpy(node->data, data, tree->objw);
} else if(node != NULL) {
rbnode_t *child = s_node_new(tree, key, data);
s_node_insert(tree, child, node, node->key > key);
return child->data;
} else {
node = s_node_new(tree, key, data);
node->parent = NULL;
node->children[0] = node->children[1] = NULL;
node->color = BLACK;
tree->root = node;
}
return node->data;
}
void *rbtree_reserve(rbtree_t *tree, intmax_t key)
{
assert(tree != NULL);
rbnode_t *node = s_closest(tree, key);
if(node != NULL && node->key == key) {
return node->data;
} else if(node != NULL) {
rbnode_t *child = s_node_new(tree, key, NULL);
s_node_insert(tree, child, node, node->key > key);
return child->data;
} else {
node = s_node_new(tree, key, NULL);
tree->root = node;
}
return node->data;
}
