Add rbtree

base/rbtree.c

@@ -0,0 +1,467 @@
+/*
+ * =============================================================================
+ *
+ *       Filename:  rbtree.c
+ *
+ *    Description:  rbtree(Red-Black tree) implementation adapted from linux
+ *                  kernel thus can be used in userspace c program.
+ *
+ *        Created:  09/02/2012 11:38:12 PM
+ *
+ *         Author:  Fu Haiping (forhappy), haipingf@gmail.com
+ *        Company:  ICT ( Institute Of Computing Technology, CAS )
+ *
+ * =============================================================================
+ */
+
+/*
+  Red Black Trees
+  (C) 1999  Andrea Arcangeli <andrea@suse.de>
+  (C) 2002  David Woodhouse <dwmw2@infradead.org>
+
+  This program is free software; you can redistribute it and/or modify
+  it under the terms of the GNU General Public License as published by
+  the Free Software Foundation; either version 2 of the License, or
+  (at your option) any later version.
+
+  This program is distributed in the hope that it will be useful,
+  but WITHOUT ANY WARRANTY; without even the implied warranty of
+  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+  GNU General Public License for more details.
+
+  You should have received a copy of the GNU General Public License
+  along with this program; if not, write to the Free Software
+  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
+
+  linux/lib/rbtree.c
+*/
+
+#include "rbtree.h"
+
+static void __rb_rotate_left(struct rb_node *node, struct rb_root *root)
+{
+    struct rb_node *right = node->rb_right;
+    struct rb_node *parent = rb_parent(node);
+
+    if ((node->rb_right = right->rb_left))
+        rb_set_parent(right->rb_left, node);
+    right->rb_left = node;
+
+    rb_set_parent(right, parent);
+
+    if (parent)
+    {
+        if (node == parent->rb_left)
+            parent->rb_left = right;
+        else
+            parent->rb_right = right;
+    }
+    else
+        root->rb_node = right;
+    rb_set_parent(node, right);
+}
+
+static void __rb_rotate_right(struct rb_node *node, struct rb_root *root)
+{
+    struct rb_node *left = node->rb_left;
+    struct rb_node *parent = rb_parent(node);
+
+    if ((node->rb_left = left->rb_right))
+        rb_set_parent(left->rb_right, node);
+    left->rb_right = node;
+
+    rb_set_parent(left, parent);
+
+    if (parent)
+    {
+        if (node == parent->rb_right)
+            parent->rb_right = left;
+        else
+            parent->rb_left = left;
+    }
+    else
+        root->rb_node = left;
+    rb_set_parent(node, left);
+}
+
+void rb_insert_color(struct rb_node *node, struct rb_root *root)
+{
+    struct rb_node *parent, *gparent;
+
+    while ((parent = rb_parent(node)) && rb_is_red(parent))
+    {
+        gparent = rb_parent(parent);
+
+        if (parent == gparent->rb_left)
+        {
+            {
+                struct rb_node *uncle = gparent->rb_right;
+                if (uncle && rb_is_red(uncle))
+                {
+                    rb_set_black(uncle);
+                    rb_set_black(parent);
+                    rb_set_red(gparent);
+                    node = gparent;
+                    continue;
+                }
+            }
+
+            if (parent->rb_right == node)
+            {
+                struct rb_node *tmp;
+                __rb_rotate_left(parent, root);
+                tmp = parent;
+                parent = node;
+                node = tmp;
+            }
+
+            rb_set_black(parent);
+            rb_set_red(gparent);
+            __rb_rotate_right(gparent, root);
+        } else {
+            {
+                struct rb_node *uncle = gparent->rb_left;
+                if (uncle && rb_is_red(uncle))
+                {
+                    rb_set_black(uncle);
+                    rb_set_black(parent);
+                    rb_set_red(gparent);
+                    node = gparent;
+                    continue;
+                }
+            }
+
+            if (parent->rb_left == node)
+            {
+                struct rb_node *tmp;
+                __rb_rotate_right(parent, root);
+                tmp = parent;
+                parent = node;
+                node = tmp;
+            }
+
+            rb_set_black(parent);
+            rb_set_red(gparent);
+            __rb_rotate_left(gparent, root);
+        }
+    }
+
+    rb_set_black(root->rb_node);
+}
+
+static void __rb_erase_color(struct rb_node *node, struct rb_node *parent,
+                 struct rb_root *root)
+{
+    struct rb_node *other;
+
+    while ((!node || rb_is_black(node)) && node != root->rb_node)
+    {
+        if (parent->rb_left == node)
+        {
+            other = parent->rb_right;
+            if (rb_is_red(other))
+            {
+                rb_set_black(other);
+                rb_set_red(parent);
+                __rb_rotate_left(parent, root);
+                other = parent->rb_right;
+            }
+            if ((!other->rb_left || rb_is_black(other->rb_left)) &&
+                (!other->rb_right || rb_is_black(other->rb_right)))
+            {
+                rb_set_red(other);
+                node = parent;
+                parent = rb_parent(node);
+            }
+            else
+            {
+                if (!other->rb_right || rb_is_black(other->rb_right))
+                {
+                    rb_set_black(other->rb_left);
+                    rb_set_red(other);
+                    __rb_rotate_right(other, root);
+                    other = parent->rb_right;
+                }
+                rb_set_color(other, rb_color(parent));
+                rb_set_black(parent);
+                rb_set_black(other->rb_right);
+                __rb_rotate_left(parent, root);
+                node = root->rb_node;
+                break;
+            }
+        }
+        else
+        {
+            other = parent->rb_left;
+            if (rb_is_red(other))
+            {
+                rb_set_black(other);
+                rb_set_red(parent);
+                __rb_rotate_right(parent, root);
+                other = parent->rb_left;
+            }
+            if ((!other->rb_left || rb_is_black(other->rb_left)) &&
+                (!other->rb_right || rb_is_black(other->rb_right)))
+            {
+                rb_set_red(other);
+                node = parent;
+                parent = rb_parent(node);
+            }
+            else
+            {
+                if (!other->rb_left || rb_is_black(other->rb_left))
+                {
+                    rb_set_black(other->rb_right);
+                    rb_set_red(other);
+                    __rb_rotate_left(other, root);
+                    other = parent->rb_left;
+                }
+                rb_set_color(other, rb_color(parent));
+                rb_set_black(parent);
+                rb_set_black(other->rb_left);
+                __rb_rotate_right(parent, root);
+                node = root->rb_node;
+                break;
+            }
+        }
+    }
+    if (node)
+        rb_set_black(node);
+}
+
+void rb_erase(struct rb_node *node, struct rb_root *root)
+{
+    struct rb_node *child, *parent;
+    int color;
+
+    if (!node->rb_left)
+        child = node->rb_right;
+    else if (!node->rb_right)
+        child = node->rb_left;
+    else
+    {
+        struct rb_node *old = node, *left;
+
+        node = node->rb_right;
+        while ((left = node->rb_left) != NULL)
+            node = left;
+
+        if (rb_parent(old)) {
+            if (rb_parent(old)->rb_left == old)
+                rb_parent(old)->rb_left = node;
+            else
+                rb_parent(old)->rb_right = node;
+        } else
+            root->rb_node = node;
+
+        child = node->rb_right;
+        parent = rb_parent(node);
+        color = rb_color(node);
+
+        if (parent == old) {
+            parent = node;
+        } else {
+            if (child)
+                rb_set_parent(child, parent);
+            parent->rb_left = child;
+
+            node->rb_right = old->rb_right;
+            rb_set_parent(old->rb_right, node);
+        }
+
+        node->rb_parent_color = old->rb_parent_color;
+        node->rb_left = old->rb_left;
+        rb_set_parent(old->rb_left, node);
+
+        goto color;
+    }
+
+    parent = rb_parent(node);
+    color = rb_color(node);
+
+    if (child)
+        rb_set_parent(child, parent);
+    if (parent)
+    {
+        if (parent->rb_left == node)
+            parent->rb_left = child;
+        else
+            parent->rb_right = child;
+    }
+    else
+        root->rb_node = child;
+
+ color:
+    if (color == RB_BLACK)
+        __rb_erase_color(child, parent, root);
+}
+
+static void rb_augment_path(struct rb_node *node, rb_augment_f func, void *data)
+{
+    struct rb_node *parent;
+
+up:
+    func(node, data);
+    parent = rb_parent(node);
+    if (!parent)
+        return;
+
+    if (node == parent->rb_left && parent->rb_right)
+        func(parent->rb_right, data);
+    else if (parent->rb_left)
+        func(parent->rb_left, data);
+
+    node = parent;
+    goto up;
+}
+
+/*
+ * after inserting @node into the tree, update the tree to account for
+ * both the new entry and any damage done by rebalance
+ */
+void rb_augment_insert(struct rb_node *node, rb_augment_f func, void *data)
+{
+    if (node->rb_left)
+        node = node->rb_left;
+    else if (node->rb_right)
+        node = node->rb_right;
+
+    rb_augment_path(node, func, data);
+}
+
+/*
+ * before removing the node, find the deepest node on the rebalance path
+ * that will still be there after @node gets removed
+ */
+struct rb_node *rb_augment_erase_begin(struct rb_node *node)
+{
+    struct rb_node *deepest;
+
+    if (!node->rb_right && !node->rb_left)
+        deepest = rb_parent(node);
+    else if (!node->rb_right)
+        deepest = node->rb_left;
+    else if (!node->rb_left)
+        deepest = node->rb_right;
+    else {
+        deepest = rb_next(node);
+        if (deepest->rb_right)
+            deepest = deepest->rb_right;
+        else if (rb_parent(deepest) != node)
+            deepest = rb_parent(deepest);
+    }
+
+    return deepest;
+}
+
+/*
+ * after removal, update the tree to account for the removed entry
+ * and any rebalance damage.
+ */
+void rb_augment_erase_end(struct rb_node *node, rb_augment_f func, void *data)
+{
+    if (node)
+        rb_augment_path(node, func, data);
+}
+
+/*
+ * This function returns the first node (in sort order) of the tree.
+ */
+struct rb_node *rb_first(const struct rb_root *root)
+{
+    struct rb_node  *n;
+
+    n = root->rb_node;
+    if (!n)
+        return NULL;
+    while (n->rb_left)
+        n = n->rb_left;
+    return n;
+}
+
+struct rb_node *rb_last(const struct rb_root *root)
+{
+    struct rb_node  *n;
+
+    n = root->rb_node;
+    if (!n)
+        return NULL;
+    while (n->rb_right)
+        n = n->rb_right;
+    return n;
+}
+
+struct rb_node *rb_next(const struct rb_node *node)
+{
+    struct rb_node *parent;
+
+    if (rb_parent(node) == node)
+        return NULL;
+
+    /* If we have a right-hand child, go down and then left as far
+       as we can. */
+    if (node->rb_right) {
+        node = node->rb_right;
+        while (node->rb_left)
+            node=node->rb_left;
+        return (struct rb_node *)node;
+    }
+
+    /* No right-hand children.  Everything down and left is
+       smaller than us, so any 'next' node must be in the general
+       direction of our parent. Go up the tree; any time the
+       ancestor is a right-hand child of its parent, keep going
+       up. First time it's a left-hand child of its parent, said
+       parent is our 'next' node. */
+    while ((parent = rb_parent(node)) && node == parent->rb_right)
+        node = parent;
+
+    return parent;
+}
+
+struct rb_node *rb_prev(const struct rb_node *node)
+{
+    struct rb_node *parent;
+
+    if (rb_parent(node) == node)
+        return NULL;
+
+    /* If we have a left-hand child, go down and then right as far
+       as we can. */
+    if (node->rb_left) {
+        node = node->rb_left;
+        while (node->rb_right)
+            node=node->rb_right;
+        return (struct rb_node *)node;
+    }
+
+    /* No left-hand children. Go up till we find an ancestor which
+       is a right-hand child of its parent */
+    while ((parent = rb_parent(node)) && node == parent->rb_left)
+        node = parent;
+
+    return parent;
+}
+
+void rb_replace_node(struct rb_node *victim, struct rb_node *newer,
+             struct rb_root *root)
+{
+    struct rb_node *parent = rb_parent(victim);
+
+    /* Set the surrounding nodes to point to the replacement */
+    if (parent) {
+        if (victim == parent->rb_left)
+            parent->rb_left = newer;
+        else
+            parent->rb_right = newer;
+    } else {
+        root->rb_node = newer;
+    }
+    if (victim->rb_left)
+        rb_set_parent(victim->rb_left, newer);
+    if (victim->rb_right)
+        rb_set_parent(victim->rb_right, newer);
+
+    /* Copy the pointers/colour from the victim to the replacement */
+    *newer = *victim;
+}

base/rbtree.h

@@ -0,0 +1,207 @@
+/*
+ * =============================================================================
+ *
+ *       Filename:  rbtree.h
+ *
+ *    Description:  rbtree(Red-Black tree) implementation adapted from linux
+ *                  kernel thus can be used in userspace c program.
+ *
+ *        Created:  09/02/2012 11:36:11 PM
+ *
+ *         Author:  Fu Haiping (forhappy), haipingf@gmail.com
+ *        Company:  ICT ( Institute Of Computing Technology, CAS )
+ *
+ * =============================================================================
+ */
+
+/*
+  Red Black Trees
+  (C) 1999  Andrea Arcangeli <andrea@suse.de>
+
+  This program is free software; you can redistribute it and/or modify
+  it under the terms of the GNU General Public License as published by
+  the Free Software Foundation; either version 2 of the License, or
+  (at your option) any later version.
+
+  This program is distributed in the hope that it will be useful,
+  but WITHOUT ANY WARRANTY; without even the implied warranty of
+  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+  GNU General Public License for more details.
+
+  You should have received a copy of the GNU General Public License
+  along with this program; if not, write to the Free Software
+  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
+
+  linux/include/linux/rbtree.h
+
+  To use rbtrees you'll have to implement your own insert and search cores.
+  This will avoid us to use callbacks and to drop drammatically performances.
+  I know it's not the cleaner way,  but in C (not in C++) to get
+  performances and genericity...
+
+  Some example of insert and search follows here. The search is a plain
+  normal search over an ordered tree. The insert instead must be implemented
+  in two steps: First, the code must insert the element in order as a red leaf
+  in the tree, and then the support library function rb_insert_color() must
+  be called. Such function will do the not trivial work to rebalance the
+  rbtree, if necessary.
+
+-----------------------------------------------------------------------
+static inline struct page * rb_search_page_cache(struct inode * inode,
+                         unsigned long offset)
+{
+    struct rb_node * n = inode->i_rb_page_cache.rb_node;
+    struct page * page;
+
+    while (n)
+    {
+        page = rb_entry(n, struct page, rb_page_cache);
+
+        if (offset < page->offset)
+            n = n->rb_left;
+        else if (offset > page->offset)
+            n = n->rb_right;
+        else
+            return page;
+    }
+    return NULL;
+}
+
+static inline struct page * __rb_insert_page_cache(struct inode * inode,
+                           unsigned long offset,
+                           struct rb_node * node)
+{
+    struct rb_node ** p = &inode->i_rb_page_cache.rb_node;
+    struct rb_node * parent = NULL;
+    struct page * page;
+
+    while (*p)
+    {
+        parent = *p;
+        page = rb_entry(parent, struct page, rb_page_cache);
+
+        if (offset < page->offset)
+            p = &(*p)->rb_left;
+        else if (offset > page->offset)
+            p = &(*p)->rb_right;
+        else
+            return page;
+    }
+
+    rb_link_node(node, parent, p);
+
+    return NULL;
+}
+
+static inline struct page * rb_insert_page_cache(struct inode * inode,
+                         unsigned long offset,
+                         struct rb_node * node)
+{
+    struct page * ret;
+    if ((ret = __rb_insert_page_cache(inode, offset, node)))
+        goto out;
+    rb_insert_color(node, &inode->i_rb_page_cache);
+ out:
+    return ret;
+}
+-----------------------------------------------------------------------
+*/
+
+#ifndef _LINUX_RBTREE_H
+#define _LINUX_RBTREE_H
+
+#ifndef NULL
+#ifdef __cplusplus
+    #define NULL    0
+#else
+    #define NULL    ((void*)0)
+#endif
+#endif
+
+#ifndef offsetof
+#define offsetof(type, member) \
+((size_t)(&((type*)0)->member))
+#endif
+
+#ifndef container_of
+#define container_of(ptr, type, member) \
+((type*)((char*)(ptr) - offsetof(type, member)))
+#endif
+
+struct rb_node
+{
+    unsigned long  rb_parent_color;
+#define RB_RED      0
+#define RB_BLACK    1
+    struct rb_node *rb_right;
+    struct rb_node *rb_left;
+};
+
+struct rb_root
+{
+    struct rb_node *rb_node;
+};
+
+
+#define rb_parent(r)   ((struct rb_node *)((r)->rb_parent_color & ~3))
+#define rb_color(r)   ((r)->rb_parent_color & 1)
+#define rb_is_red(r)   (!rb_color(r))
+#define rb_is_black(r) rb_color(r)
+#define rb_set_red(r)  do { (r)->rb_parent_color &= ~1; } while (0)
+#define rb_set_black(r)  do { (r)->rb_parent_color |= 1; } while (0)
+
+static inline void rb_set_parent(struct rb_node *rb, struct rb_node *p)
+{
+    rb->rb_parent_color = (rb->rb_parent_color & 3) | (unsigned long)p;
+}
+static inline void rb_set_color(struct rb_node *rb, int color)
+{
+    rb->rb_parent_color = (rb->rb_parent_color & ~1) | color;
+}
+
+#define RB_ROOT (struct rb_root) { NULL, }
+#define rb_entry(ptr, type, member) container_of(ptr, type, member)
+
+#define RB_EMPTY_ROOT(root) ((root)->rb_node == NULL)
+#define RB_EMPTY_NODE(node) (rb_parent(node) == node)
+#define RB_CLEAR_NODE(node) (rb_set_parent(node, node))
+
+static inline void rb_init_node(struct rb_node *rb)
+{
+    rb->rb_parent_color = 0;
+    rb->rb_right = NULL;
+    rb->rb_left = NULL;
+    RB_CLEAR_NODE(rb);
+}
+
+void rb_insert_color(struct rb_node *, struct rb_root *);
+void rb_erase(struct rb_node *, struct rb_root *);
+
+typedef void (*rb_augment_f)(struct rb_node *node, void *data);
+
+void rb_augment_insert(struct rb_node *node,
+                  rb_augment_f func, void *data);
+struct rb_node *rb_augment_erase_begin(struct rb_node *node);
+void rb_augment_erase_end(struct rb_node *node,
+                 rb_augment_f func, void *data);
+
+/* Find logical next and previous nodes in a tree */
+struct rb_node *rb_next(const struct rb_node *);
+struct rb_node *rb_prev(const struct rb_node *);
+struct rb_node *rb_first(const struct rb_root *);
+struct rb_node *rb_last(const struct rb_root *);
+
+/* Fast replacement of a single node without remove/rebalance/add/rebalance */
+void rb_replace_node(struct rb_node *victim, struct rb_node *newer,
+                struct rb_root *root);
+
+static inline void rb_link_node(struct rb_node * node, struct rb_node * parent,
+                struct rb_node ** rb_link)
+{
+    node->rb_parent_color = (unsigned long )parent;
+    node->rb_left = node->rb_right = NULL;
+
+    *rb_link = node;
+}
+
+#endif  /* _LINUX_RBTREE_H */

unittest/rbtree_test.c

@@ -0,0 +1,111 @@
+#include <stdio.h>
+#include <string.h>
+
+#include "rbtree.h"
+
+typedef int rbtree_key_type;
+typedef int rbtree_val_type;
+
+struct rbtree_entry {
+    struct rb_node  rb_node;
+    rbtree_key_type key;
+    rbtree_val_type val;
+};
+
+int rbtree_insert(struct rb_root* root, struct rbtree_entry* entry) {
+    printf("insert %d\n", entry->key);
+    struct rb_node** n = &root->rb_node;
+    struct rb_node* parent = NULL;
+    struct rbtree_entry* e = NULL;
+    while (*n) {
+        parent = *n;
+        e = rb_entry(*n, struct rbtree_entry, rb_node);
+        if (entry->key < e->key) {
+            n = &(*n)->rb_left;
+        } else if (entry->key > e->key) {
+            n = &(*n)->rb_right;
+        } else {
+            return -1;
+        }
+    }
+
+    rb_link_node(&entry->rb_node, parent, n);
+    rb_insert_color(&entry->rb_node, root);
+    return 0;
+}
+
+int rbtree_remove(struct rb_root* root, struct rbtree_entry* entry) {
+    printf("remove %d\n", entry->key);
+    rb_erase(&entry->rb_node, root);
+    return 0;
+}
+
+struct rbtree_entry* rbtree_search(struct rb_root* root, const rbtree_key_type* key) {
+    struct rb_node* n = root->rb_node;
+    struct rbtree_entry* e = NULL;
+    while (n) {
+        e = rb_entry(n, struct rbtree_entry, rb_node);
+        if (*key < e->key) {
+            n = n->rb_left;
+        } else if (*key > e->key) {
+            n = n->rb_right;
+        } else {
+            return e;
+        }
+    }
+    return NULL;
+}
+
+void rbtree_entry_print(struct rbtree_entry* entry) {
+    if (entry == NULL) {
+        printf("null\n");
+        return;
+    }
+    printf("%d:%d\n", entry->key, entry->val);
+}
+
+int main() {
+    struct rb_root root = { NULL };
+    struct rbtree_entry* entry = NULL;
+
+    struct rbtree_entry entries[10];
+    for (int i = 0; i < 10; ++i) {
+        memset(&entries[i], 0, sizeof(struct rbtree_entry));
+        entries[i].key = i;
+        entries[i].val = i;
+    }
+
+    rbtree_insert(&root, &entries[1]);
+    rbtree_insert(&root, &entries[2]);
+    rbtree_insert(&root, &entries[3]);
+    rbtree_insert(&root, &entries[7]);
+    rbtree_insert(&root, &entries[8]);
+    rbtree_insert(&root, &entries[9]);
+    rbtree_insert(&root, &entries[4]);
+    rbtree_insert(&root, &entries[5]);
+    rbtree_insert(&root, &entries[6]);
+
+    rbtree_remove(&root, &entries[1]);
+    rbtree_remove(&root, &entries[9]);
+    rbtree_remove(&root, &entries[4]);
+    rbtree_remove(&root, &entries[6]);
+
+    int key = 5;
+    entry = rbtree_search(&root, &key);
+    rbtree_entry_print(entry);
+
+    key = 4;
+    entry = rbtree_search(&root, &key);
+    rbtree_entry_print(entry);
+
+    struct rb_node* node = NULL;
+    // while((node = rb_first(&root))) {
+    while((node = root.rb_node)) {
+        entry = rb_entry(node, struct rbtree_entry, rb_node);
+        rb_erase(node, &root);
+        rbtree_entry_print(entry);
+        memset(entry, 0, sizeof(struct rbtree_entry));
+    }
+
+    return 0;
+}