1 files changed, 26 insertions, 7 deletions

diff --git a/lib/rbtree.c b/lib/rbtree.c
index dc9caf1..a63df4d 100644
--- a/lib/rbtree.c
+++ b/lib/rbtree.c
@@ -1,4 +1,6 @@
 #include "rbtree.h"
+#include "api/error.h"
+#include "errno.h"
 
 static inline int
 s_direction(const struct rbtree_node *node)
@@ -26,10 +28,27 @@ s_rotate(struct rbtree_node *sub, int dir)
     return new_root;
 }
 
-void
-rbtree_insert(struct rbtree_node *root, struct rbtree_node *node)
+struct rbtree_node *
+rbtree_search(struct rbtree_node *parent, intmax_t key)
 {
-    struct rbtree_node *parent = root;
+    errno = 0;
+    intmax_t parent_key = parent->get_key(parent);
+    if(parent_key == key) return parent;
+    int searchdir = parent_key < key;
+    if(parent->children[searchdir])
+        return rbtree_search(parent->children[searchdir], key);
+    
+    errno = ENOTFOUND;
+    return 0;
+}
+
+struct rbtree_node *
+rbtree_insert(struct rbtree_node *parent, struct rbtree_node *node)
+{
+    node->color = 1;
+    if(parent == 0) return node;
+
+    //Seek to the correct b-tree leaf.
     int direction;
     while(1) {
         direction = parent->get_key(parent) < node->get_key(node);
@@ -37,20 +56,19 @@ rbtree_insert(struct rbtree_node *root, struct rbtree_node *node)
         parent = parent->children[direction];
     }
 
-    node->color = 1;
     parent->children[direction] = node;
     node->parent = parent;
 
     //Rebalance the tree
     do {
         //Tree is already balanced, return
-        if(parent->color == 0) return;
+        if(parent->color == 0) return parent;
         struct rbtree_node *gparent = parent->parent;
 
         //Parent is the root node, set parent to black.
         if(!gparent) {
             parent->color = 0;
-            return;
+            return parent;
         }
 
         direction = s_direction(parent);
@@ -68,7 +86,7 @@ rbtree_insert(struct rbtree_node *root, struct rbtree_node *node)
             s_rotate(gparent, 1-direction);
             parent->color = 0;
             gparent->color = 1;
-            return;
+            return parent;
         }
 
         //Both the parent and uncle are red. Push up the red and iterate up.
@@ -79,6 +97,7 @@ rbtree_insert(struct rbtree_node *root, struct rbtree_node *node)
     } while((parent = node->parent));
 
     //Tree is balanced now.
+    return parent;
 }