fix(tree): working balanced remove

src/tree/vieter_tree.c

@@ -48,7 +48,7 @@ vieter_tree_error vieter_tree_remove(void **out, vieter_tree *tree,
     return vieter_tree_not_present;
   }
 
-  vieter_tree_error res = vieter_tree_node_remove(out, tree->root, key);
+  vieter_tree_error res = vieter_tree_node_remove(out, &tree->root, key);
 
   if (res != vieter_tree_ok) {
     return res;

src/tree/vieter_tree_balancing.c

@@ -33,6 +33,9 @@ bool vieter_tree_node_validate(vieter_tree_node *node,
                                    expected_black_nodes);
 }
 
+/*
+ * This function should be rewritten to use tree rotations instead.
+ */
 vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
   vieter_tree_node *parent = node->parent;
   vieter_tree_node *grand_parent = parent->parent;
@@ -144,7 +147,9 @@ vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root,
   if (dir) {
     vieter_tree_node_set_child(old_root, new_root->children[1], false);
     vieter_tree_node_set_child(new_root, old_root, true);
-  } else {
+  }
+  // Left rotation
+  else {
     vieter_tree_node_set_child(old_root, new_root->children[0], true);
     vieter_tree_node_set_child(new_root, old_root, false);
   }
@@ -152,7 +157,32 @@ vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root,
   return new_root;
 }
 
-void vieter_tree_node_remove_balanced(vieter_tree_node *node) {
+/*
+ * This function is currently implemented by very literally following the
+ * Wikipedia pseudocode. It's honestly not too bad, and I couldn't be bothered
+ * to properly research how to implement red-black removal (how is this so hard
+ * to find?)
+ *
+ * https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Removal_of_a_black_non-root_leaf
+ */
+vieter_tree_node *vieter_tree_node_remove_balanced(vieter_tree_node *node) {
+  vieter_tree_node *out;
+
+  if (node->parent == NULL) {
+    out = node->children[0] != NULL ? node->children[0] : node->children[1];
+
+    vieter_tree_node_free(node);
+
+    // This only happens when the root was the only element in the tree
+    if (out == NULL)
+      return out;
+
+    vieter_tree_node_set(out, vieter_tree_node_black, true);
+    out->parent = NULL;
+
+    return out;
+  }
+
   // A red node can only have 0 or 2 children. The node we receive only has
   // one child at most, so we know if it's red that it doesn't have any
   // children. A black node that has a single (right) child can be replaced by
@@ -160,6 +190,8 @@ void vieter_tree_node_remove_balanced(vieter_tree_node *node) {
   // replaced by its right child (even if it's NULL).
   if (!vieter_tree_node_get(node, vieter_tree_node_black) ||
       node->children[1] != NULL) {
+    out = node->children[1];
+
     vieter_tree_node_set_child(
         node->parent, node->children[1],
         vieter_tree_node_get(node, vieter_tree_node_right));
@@ -169,12 +201,13 @@ void vieter_tree_node_remove_balanced(vieter_tree_node *node) {
 
     vieter_tree_node_free(node);
 
-    return;
+    return out;
   }
 
   // The complicated case is when we want to remove a black leaf
   // https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Removal_of_a_black_non-root_leaf
 
+  out = node;
   vieter_tree_node *parent = node->parent;
   vieter_tree_node *sibling, *close_nephew, *distant_nephew;
   bool dir = vieter_tree_node_get(node, vieter_tree_node_right);
@@ -209,13 +242,14 @@ void vieter_tree_node_remove_balanced(vieter_tree_node *node) {
     // Case 2
     vieter_tree_node_set(sibling, vieter_tree_node_black, false);
     node = parent;
+    out = node;
   } while ((parent = node->parent) != NULL);
 
   // Case 1
-  return;
+  return out;
 
 case3:
-  vieter_tree_node_rotate(parent, dir);
+  out = vieter_tree_node_rotate(parent, dir);
   vieter_tree_node_set(parent, vieter_tree_node_black, false);
   vieter_tree_node_set(sibling, vieter_tree_node_black, true);
   sibling = close_nephew;
@@ -234,7 +268,8 @@ case3:
 case4:
   vieter_tree_node_set(sibling, vieter_tree_node_black, false);
   vieter_tree_node_set(parent, vieter_tree_node_black, true);
-  return;
+
+  return out;
 
 case5:
   vieter_tree_node_rotate(sibling, 1 - dir);
@@ -244,9 +279,11 @@ case5:
   sibling = close_nephew;
 
 case6:
-  vieter_tree_node_rotate(parent, dir);
+  out = vieter_tree_node_rotate(parent, dir);
   vieter_tree_node_set(sibling, vieter_tree_node_black,
                        vieter_tree_node_get(parent, vieter_tree_node_black));
   vieter_tree_node_set(parent, vieter_tree_node_black, true);
   vieter_tree_node_set(distant_nephew, vieter_tree_node_black, true);
+
+  return out;
 }

src/tree/vieter_tree_balancing.h

@@ -3,6 +3,26 @@
 
 #include "vieter_tree_node.h"
 
+/*
+ * Check whether the given tree is a valid red-black tree.
+ *
+ * @param node root of the (sub)tree
+ * @param passed_black_nodes how many black nodes the recursion has already seen
+ * at this point. This should be initialized as 0 for the topmost call.
+ * @param expected_black_nodes the correct amount of black nodes to expect when
+ * a NULL child is encountered.
+ * @return true if the tree is valid, false otherwise
+ */
+bool vieter_tree_node_validate(vieter_tree_node *node,
+                               uint64_t passed_black_nodes,
+                               uint64_t expected_black_nodes);
+
+/*
+ * Balance a path of 3 nodes into a complete binary tree, with a red root and
+ * black children.
+ */
+vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node);
+
 /*
  * Ensure the tree remains a valid red-black tree after having inserting the
  * node.
@@ -12,10 +32,15 @@ void vieter_tree_node_balance_after_insert(vieter_tree_node *node);
 /*
  * Remove the given node, ensuring the tree remains a valid red-black tree.
  *
- * @param node node to remove. This should have at most a single child, namely
- * the right one.
+ * @param node node to remove. This should have at most a single child. If node
+ * isn't the root, this should be the right child, otherwise it can be either.
+ * @return root of the subtree that this function operated on. This can
+ * sometimes be the new root of the entire tree.
  */
-void vieter_tree_node_remove_balanced(vieter_tree_node *node);
+vieter_tree_node *vieter_tree_node_remove_balanced(vieter_tree_node *node);
 
+/*
+ * Perform a tree rotation of the subtree with the given root.
+ */
 vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root, bool dir);
 #endif

src/tree/vieter_tree_node.c

@@ -109,16 +109,17 @@ vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root,
   return vieter_tree_ok;
 }
 
-vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
-                                          uint64_t key) {
+vieter_tree_error
+vieter_tree_node_remove(void **out, vieter_tree_node **root_ptr, uint64_t key) {
   vieter_tree_node *target;
-  vieter_tree_error res = vieter_tree_node_search_node(&target, root, key);
+  vieter_tree_error res = vieter_tree_node_search_node(&target, *root_ptr, key);
 
   if (res != vieter_tree_ok) {
     return res;
   }
 
   *out = target->data;
+  vieter_tree_node *possible_new_root;
 
   if (target->children[0] != NULL && target->children[1] != NULL) {
     vieter_tree_node *replacement = target->children[1];
@@ -130,9 +131,15 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
     target->key = replacement->key;
     target->data = replacement->data;
 
-    vieter_tree_node_remove_balanced(replacement);
+    possible_new_root = vieter_tree_node_remove_balanced(replacement);
   } else {
-    vieter_tree_node_remove_balanced(target);
+    possible_new_root = vieter_tree_node_remove_balanced(target);
+  }
+
+  if (possible_new_root == NULL) {
+    *root_ptr = NULL;
+  } else if (possible_new_root->parent == NULL) {
+    *root_ptr = possible_new_root;
   }
 
   return vieter_tree_ok;

src/tree/vieter_tree_node.h

@@ -33,7 +33,7 @@ vieter_tree_error vieter_tree_node_search_node(vieter_tree_node **out, vieter_tr
 
 vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root, uint64_t key);
 
-vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root, uint64_t key);
+vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node **root_ptr, uint64_t key);
 
 void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace, vieter_tree_node *replacement);
 

test/tree/test_balancing.c

@@ -3,6 +3,9 @@
 #include "vieter_tree_node.h"
 #include "vieter_tree_balancing.h"
 
+// This uses the examples from wikipedia
+// https://en.wikipedia.org/wiki/Tree_rotation
+
 void test_rotate_right() {
     vieter_tree_node *a = vieter_tree_node_init();
     vieter_tree_node *b = vieter_tree_node_init();
@@ -24,6 +27,12 @@ void test_rotate_right() {
     TEST_CHECK(new_root->children[1] == q);
     TEST_CHECK(new_root->children[1]->children[0] == b);
     TEST_CHECK(new_root->children[1]->children[1] == c);
+
+    vieter_tree_node_free(a);
+    vieter_tree_node_free(b);
+    vieter_tree_node_free(p);
+    vieter_tree_node_free(c);
+    vieter_tree_node_free(q);
 }
 
 void test_rotate_left() {
@@ -47,6 +56,12 @@ void test_rotate_left() {
     TEST_CHECK(new_root->children[1] == c);
     TEST_CHECK(new_root->children[0]->children[0] == a);
     TEST_CHECK(new_root->children[0]->children[1] == b);
+
+    vieter_tree_node_free(a);
+    vieter_tree_node_free(b);
+    vieter_tree_node_free(p);
+    vieter_tree_node_free(c);
+    vieter_tree_node_free(q);
 }
 
 TEST_LIST = {

test/tree/test_tree.c

@@ -42,7 +42,7 @@ void test_remove() {
 
     void *out;
 
-    for (uint64_t i = 0; i < 100; i++) {
+    for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_ok);
         TEST_CHECK(vieter_tree_remove(&out, tree, i) == vieter_tree_ok);
         TEST_CHECK(vieter_tree_validate(tree));