lib: Fix bugs in B-tree

This fixes some bugs in the B-tree implementation. The test has also
been rewritten to be more thorough.

2 files changed, 62 insertions, 26 deletions

diff --git a/src/lib/btree.c b/src/lib/btree.c
index 10a900d6..48df8e39 100644
--- a/src/lib/btree.c
+++ b/src/lib/btree.c
@@ -115,18 +115,23 @@ static struct btnode * btnode_create(size_t k)
 
 static void btnode_destroy(struct btnode * node)
 {
-        size_t i = 0;
-
         assert(node);
 
-        for (i = 0; !node->leaf && i <= node->used; ++i)
-                btnode_destroy(node->children[i]);
-
         free(node->children);
         free(node->keyvals);
         free(node);
 }
 
+static void btnode_destroy_subtree(struct btnode * node)
+{
+        size_t i;
+
+        for (i = 0; !node->leaf && i <= node->used; ++i)
+                btnode_destroy_subtree(node->children[i]);
+
+        btnode_destroy(node);
+}
+
 static int btnode_insert(struct btnode *  node,
                          struct key_val   kv,
                          struct key_val * med,
@@ -211,15 +216,15 @@ void merge(struct btnode * node,
         if (!chld->leaf)
                 memmove(&chld->children[node->k / 2],
                         &next->children[0],
-                        sizeof(*next->children) * next->used + 1);
+                        sizeof(*next->children) * (next->used + 1));
 
         memmove(&node->keyvals[i],
                 &node->keyvals[i + 1],
-                sizeof(*node->keyvals) * node->used - i - 1);
+                sizeof(*node->keyvals) * (node->used - i - 1));
 
         memmove(&node->children[i + 1],
                 &node->children[i + 2],
-                sizeof(*node->children) * node->used - i);
+                sizeof(*node->children) * (node->used - i));
 
         chld->used += next->used + 1;
         node->used--;
@@ -262,6 +267,7 @@ static void fill(struct btnode * node,
         /* Feed from next sibling. */
         if (i != node->used && node->children[i + 1]->used >= node->k / 2) {
                 struct btnode * next = node->children[i + 1];
+
                 chld->keyvals[chld->used] = node->keyvals[i];
 
                 if (!chld->leaf)
@@ -294,11 +300,11 @@ static void fill(struct btnode * node,
 static struct key_val btnode_pred(struct btnode * node,
                                   size_t          i)
 {
-        struct btnode * pred = node->children[i - 1];
+        struct btnode * pred = node->children[i];
         while (!pred->leaf)
                 pred = pred->children[pred->used];
 
-        return pred->keyvals[pred->used -1];
+        return pred->keyvals[pred->used - 1];
 }
 
 static struct key_val btnode_succ(struct btnode * node,
@@ -314,6 +320,7 @@ static int btnode_delete(struct btnode * node,
                          uint32_t        key)
 {
         size_t i;
+        int ret = 0;
 
         assert(node);
 
@@ -329,15 +336,15 @@ static int btnode_delete(struct btnode * node,
                 } else {
                         if (node->children[i]->used >= node->k / 2) {
                                 node->keyvals[i] = btnode_pred(node, i);
-                                btnode_delete(node->children[i],
-                                              node->keyvals[i].key);
+                                ret = btnode_delete(node->children[i],
+                                                    node->keyvals[i].key);
                         } else if (node->children[i + 1]->used >= node->k / 2) {
                                 node->keyvals[i] = btnode_succ(node, i);
-                                btnode_delete(node->children[i + 1],
-                                              node->keyvals[i].key);
+                                ret = btnode_delete(node->children[i + 1],
+                                                    node->keyvals[i].key);
                         } else {
                                 merge(node, i);
-                                btnode_delete(node, key);
+                                ret = btnode_delete(node, key);
                         }
                 }
         } else {
@@ -345,16 +352,16 @@ static int btnode_delete(struct btnode * node,
                         return -1; /* value not in tree */
                 } else {
                         bool flag = (i == node->used ? true : false);
-                        if (node->children[i]->used > node->children[i]->k / 2)
+                        if (node->children[i]->used < node->children[i]->k / 2)
                                 fill(node, i);
                         if (flag && i > node->used)
-                                btnode_delete(node->children[i - 1], key);
+                                ret = btnode_delete(node->children[i - 1], key);
                         else
-                                btnode_delete(node->children[i + 1], key);
+                                ret = btnode_delete(node->children[i], key);
                 }
         }
 
-        return 0;
+        return ret;
 }
 
 struct btree * btree_create(size_t k)
@@ -378,7 +385,7 @@ void btree_destroy(struct btree * tree)
                 return;
 
         if (tree->root != NULL)
-                btnode_destroy(tree->root);
+                btnode_destroy_subtree(tree->root);
 
         free(tree);
 }
@@ -435,20 +442,25 @@ int btree_insert(struct btree * tree,
 int btree_remove(struct btree * tree,
                  uint32_t       key)
 {
+        struct btnode * prev_root;
+
         if (tree == NULL)
                 return -EINVAL;
 
         if (tree->root == NULL)
                 return 0;
 
-        btnode_delete(tree->root, key);
+        if (btnode_delete(tree->root, key))
+                return -1;
 
         if (tree->root->used == 0) {
                 if (tree->root->leaf) {
                         btnode_destroy(tree->root);
                         tree->root = NULL;
                 } else {
+                        prev_root = tree->root;
                         tree->root = tree->root->children[0];
+                        btnode_destroy(prev_root);
                 }
         }
 
diff --git a/src/lib/tests/btree_test.c b/src/lib/tests/btree_test.c
index 6981f63a..a6344060 100644
--- a/src/lib/tests/btree_test.c
+++ b/src/lib/tests/btree_test.c
@@ -25,6 +25,7 @@
 
 #include <stdio.h>
 #include <stdlib.h>
+#include <string.h>
 
 #define MAX_BTREE_KEY 10000
 
@@ -33,11 +34,15 @@ int btree_test(int     argc,
 {
         struct btree * tree;
 
+        int vals[MAX_BTREE_KEY];
         int i;
+        int j;
 
         (void) argc;
         (void) argv;
 
+        memset(vals, 0, MAX_BTREE_KEY * sizeof(int));
+
         tree = btree_create(32);
         if (tree == NULL)
                 return -1;
@@ -47,34 +52,53 @@ int btree_test(int     argc,
                 return -1;
         }
 
-        for(i = 0; i < MAX_BTREE_KEY; ++i)
+        for (i = 0; i < MAX_BTREE_KEY; ++i)
                 if (btree_insert(tree, i, &argv)) {
                         printf("Failed to add element.\n");
                         btree_destroy(tree);
                         return -1;
                 }
 
-        for (i = 0; i < MAX_BTREE_KEY / 10; ++i)
+        for (i = 0; i < MAX_BTREE_KEY; ++i)
                 if (btree_search(tree, rand() % MAX_BTREE_KEY) != &argv) {
                         printf("Added element not in tree.\n");
                         btree_destroy(tree);
                         return -1;
                 }
 
-        for (i = 0; i < MAX_BTREE_KEY / 10; ++i)
+        for (i = 0; i < MAX_BTREE_KEY; ++i)
                 if (btree_remove(tree, i)) {
-                        printf("Failed to remove element.\n");
+                        printf("Failed to remove element %d.\n", i);
                         btree_destroy(tree);
                         return -1;
                 }
 
         for (i = 0; i < MAX_BTREE_KEY / 10; ++i)
-                if (btree_search(tree, rand() % MAX_BTREE_KEY / 10) != &argv) {
+                if (btree_search(tree, rand() % MAX_BTREE_KEY / 10) != NULL) {
                         printf("Removed element found in tree.\n");
                         btree_destroy(tree);
                         return -1;
                 }
 
+        for (i = 0; i < MAX_BTREE_KEY; ++i)
+                if (btree_insert(tree, i, &argv)) {
+                        printf("Failed to add element.\n");
+                        btree_destroy(tree);
+                        return -1;
+                }
+
+        for (i = 0; i < MAX_BTREE_KEY; ++i) {
+                j = rand() % MAX_BTREE_KEY;
+                if (vals[j] != -1) {
+                        if (btree_remove(tree, j)) {
+                                printf("Failed to remove element %d.\n", j);
+                                btree_destroy(tree);
+                                return -1;
+                        }
+                }
+                vals[j] = -1;
+        }
+
         btree_destroy(tree);
 
         return 0;