diff options
Diffstat (limited to 'klib/rbtree.c')
| -rw-r--r-- | klib/rbtree.c | 361 |
1 files changed, 361 insertions, 0 deletions
diff --git a/klib/rbtree.c b/klib/rbtree.c new file mode 100644 index 0000000..2920384 --- /dev/null +++ b/klib/rbtree.c @@ -0,0 +1,361 @@ +#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; +} + |