--- /dev/null
+/*
+ * 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 <afsconfig.h>
+#include <afs/param.h>
+
+#ifdef KERNEL
+# include "afs/sysincludes.h"
+# include "afsincludes.h"
+#else
+# include <roken.h>
+#endif
+
+#include "rbtree.h"
+
+struct {
+ struct opr_rbtree_node *left;
+ struct opr_rbtree_node *right;
+ struct opr_rbtree_node *parent;
+ int red;
+} opr_rbtree_node;
+
+struct {
+ struct opr_rbtree_node *root;
+} opr_rbtree;
+
+/* 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;
+}
+
+void
+opr_rbtree_init(struct opr_rbtree *head)
+{
+ head->root = NULL;
+}
+
+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
+initnode(struct opr_rbtree_node *node)
+{
+ node->left = node->right = node->parent = NULL;
+ node->red = 1;
+}
+
+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(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;
+ *childptr = node;
+
+ /* Rebalance the tree for the newly inserted node */
+ insert_recolour(head, node);
+}
+
+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, *parent;
+ int red;
+
+
+ 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);
+
+ return;
+ }
+
+ 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 the
+ * tree to maintain balance */
+ if (!red)
+ remove_recolour(head, parent, child);
+
+ return;
+ }
+
+ /* 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);
+}
+
+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;
+}
--- /dev/null
+#include <afsconfig.h>
+#include <afs/param.h>
+
+#include <stdlib.h>
+#include <stdio.h>
+
+#include <tap/basic.h>
+
+#include <afs/opr.h>
+#include <opr/rbtree.h>
+
+struct intnode {
+ struct opr_rbtree_node node;
+ int value;
+};
+
+int _checkTree(struct opr_rbtree_node *node)
+{
+ struct opr_rbtree_node *left, *right;
+ int lheight, rheight;
+
+ if (node == NULL)
+ return 1;
+
+ left = node->left;
+ right = node->right;
+
+ /* Leaf nodes are always black */
+ if (node->red && ((left && left->red) || (right &&right->red))) {
+ printf("Consecutive red nodes\n");
+ return 0;
+ }
+
+ /* XXX - Check ordering in nodes */
+
+ lheight = _checkTree(left);
+ rheight = _checkTree(right);
+
+ if (lheight !=0 && rheight !=0 && lheight != rheight) {
+ printf("Left and right branches have differing number of black nodes\n");
+ return 0;
+ }
+
+ if (lheight !=0 && rheight !=0)
+ return node->red ? lheight : lheight+1;
+ else
+ return 0;
+}
+
+int
+checkTree(struct opr_rbtree *head) {
+ return _checkTree(head->root);
+}
+
+int
+insertNode(struct opr_rbtree *head, int value) {
+
+ struct intnode *inode;
+ struct opr_rbtree_node *parent, **childptr;
+
+ inode = malloc(sizeof(struct intnode));
+ inode->value = value;
+
+ childptr = &head->root;
+ parent = NULL;
+
+ while (*childptr) {
+ struct intnode *tnode;
+ parent = *childptr;
+ tnode = opr_containerof(parent, struct intnode, node);
+
+ if (value < tnode->value)
+ childptr = &(*childptr)->left;
+ else if (value > tnode->value)
+ childptr = &(*childptr)->right;
+ else
+ return -1;
+ }
+ opr_rbtree_insert(head, parent, childptr, &inode->node);
+ return 0;
+}
+
+int
+countNodes(struct opr_rbtree *head) {
+ struct opr_rbtree_node *node;
+ int count;
+
+ node = opr_rbtree_first(head);
+ if (node == NULL)
+ return 0;
+
+ count = 1;
+ while ((node = opr_rbtree_next(node)))
+ count++;
+
+ return count;
+}
+
+/* Now, insert 1000 nodes into the tree, making sure with each insertion that
+ * the tree is still valid, and has the correct number of nodes
+ */
+int
+createTree(struct opr_rbtree *head)
+{
+ int counter;
+
+ for (counter = 1000; counter>0; counter--) {
+ while (insertNode(head, random())!=0);
+ if (!checkTree(head)) {
+ printf("Tree check failed at %d\n", 1001 - counter);
+ return 0;
+ }
+ if (countNodes(head) != 1001 - counter ) {
+ printf("Tree has fewer nodes than inserted : %d", 1001 - counter);
+ return 0;
+ }
+ }
+ return 1;
+}
+
+/* Randomly remove nodes from the tree, this is really, really inefficient, but
+ * hey
+ */
+int
+destroyTree(struct opr_rbtree *head)
+{
+ int counter;
+
+ for (counter = 1000; counter>0; counter--) {
+ struct opr_rbtree_node *node;
+ int remove, i;
+
+ remove = random() % counter;
+ node = opr_rbtree_first(head);
+ for (i=0; i<remove; i++)
+ node = opr_rbtree_next(node);
+
+ opr_rbtree_remove(head, node);
+ if (countNodes(head) != counter-1) {
+ printf("Tree has lost nodes after %d deletions", 1001 - counter);
+ return 0;
+ }
+
+ if (!checkTree(head)) {
+ printf("Tree check failed at %d removals\n", 1001 - counter);
+ return 0;
+ }
+ }
+ return 1;
+}
+
+int
+main(void)
+{
+ struct opr_rbtree head;
+
+ plan(3);
+
+ opr_rbtree_init(&head);
+ ok(1, "Initialising the tree works");
+
+ ok(createTree(&head), "Creating tree with 1000 nodes works");
+ ok(destroyTree(&head), "Tree retains consistency as nodes deleted");
+
+ return 0;
+}