From 1d458c68a47afb979ca7741e845529a5d38d9617 Mon Sep 17 00:00:00 2001
From: Chewing_Bever <roosensjef@gmail.com>
Date: Sun, 29 Jan 2023 21:34:44 +0100
Subject: [PATCH] feat(tree): start of balanced removal

---
 src/tree/vieter_tree_balancing.c | 23 ++++++++++++++++++-
 src/tree/vieter_tree_balancing.h |  3 ++-
 src/tree/vieter_tree_node.c      | 39 ++++++++++++--------------------
 src/tree/vieter_tree_node.h      |  1 +
 4 files changed, 39 insertions(+), 27 deletions(-)

diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c
index b160563..e2a9d18 100644
--- a/src/tree/vieter_tree_balancing.c
+++ b/src/tree/vieter_tree_balancing.c
@@ -124,4 +124,25 @@ void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
   }
 }
 
-void vieter_tree_node_remove_balanced(vieter_tree_node *node) {}
+void vieter_tree_node_remove_balanced(vieter_tree_node *node) {
+  // 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
+  // this child, granted it becomes black as well. Either way, the node can be
+  // replaced by its right child (even if it's NULL).
+  if (!vieter_tree_node_get(node, vieter_tree_node_black) ||
+      node->children[1] != NULL) {
+    vieter_tree_node_set_child(
+        node->parent, node->children[1],
+        vieter_tree_node_get(node, vieter_tree_node_right));
+
+    if (node->children[1] != NULL)
+      vieter_tree_node_set(node->children[1], vieter_tree_node_black, true);
+
+    vieter_tree_node_free(node);
+
+    return;
+  }
+
+  // The complicated case is when we want to remove a black leaf
+}
diff --git a/src/tree/vieter_tree_balancing.h b/src/tree/vieter_tree_balancing.h
index 007af34..befc035 100644
--- a/src/tree/vieter_tree_balancing.h
+++ b/src/tree/vieter_tree_balancing.h
@@ -12,7 +12,8 @@ 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 node should have at least one NULL child.
+ * @param node node to remove. This should have at most a single child, namely
+ * the right one.
  */
 void vieter_tree_node_remove_balanced(vieter_tree_node *node);
 
diff --git a/src/tree/vieter_tree_node.c b/src/tree/vieter_tree_node.c
index 2f45b82..3f09f3e 100644
--- a/src/tree/vieter_tree_node.c
+++ b/src/tree/vieter_tree_node.c
@@ -36,6 +36,16 @@ void vieter_tree_node_set_children(vieter_tree_node *parent,
   }
 }
 
+void vieter_tree_node_set_child(vieter_tree_node *parent,
+                                vieter_tree_node *child, bool right) {
+  parent->children[right] = child;
+
+  if (child != NULL) {
+    child->parent = parent;
+    vieter_tree_node_set(child, vieter_tree_node_right, right);
+  }
+}
+
 void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace,
                                          vieter_tree_node *replacement) {
   to_replace->key = replacement->key;
@@ -110,40 +120,19 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
 
   *out = target->data;
 
-  if (target->children[0] == NULL && target->children[1] == NULL) {
-    vieter_tree_node_add_child(target->parent, target->key, NULL);
-
-    vieter_tree_node_free(target);
-  } else if ((target->children[0] == NULL) ^ (target->children[1] == NULL)) {
-    vieter_tree_node *child =
-        target->children[0] != NULL ? target->children[0] : target->children[1];
-
-    if (target->parent != NULL) {
-      vieter_tree_node_add_child(target->parent, child->key, child);
-      vieter_tree_node_free(target);
-    } else {
-      vieter_tree_node_replace_with_child(target, child);
-      vieter_tree_node_free(child);
-    }
-
-  } else {
+  if (target->children[0] != NULL && target->children[1] != NULL) {
     vieter_tree_node *replacement = target->children[1];
 
     while (replacement->children[0] != NULL) {
       replacement = replacement->children[0];
     }
 
-    // We use replacement->key here because the right child can be NULL, so
-    // reading its key isn't safe. Using replacement->key however, the child
-    // will still get placed into the right location because of how binary
-    // trees work.
-    vieter_tree_node_add_child(replacement->parent, replacement->key,
-                               replacement->children[1]);
-
     target->key = replacement->key;
     target->data = replacement->data;
 
-    vieter_tree_node_free(replacement);
+    vieter_tree_node_remove_balanced(replacement);
+  } else {
+    vieter_tree_node_remove_balanced(target);
   }
 
   return vieter_tree_ok;
diff --git a/src/tree/vieter_tree_node.h b/src/tree/vieter_tree_node.h
index 1ebb15a..f164704 100644
--- a/src/tree/vieter_tree_node.h
+++ b/src/tree/vieter_tree_node.h
@@ -43,6 +43,7 @@ bool vieter_tree_node_get(vieter_tree_node *node, vieter_tree_node_flag flag);
 void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
                                 vieter_tree_node *child);
 void vieter_tree_node_set_children(vieter_tree_node *parent, vieter_tree_node **children);
+void vieter_tree_node_set_child(vieter_tree_node *parent, vieter_tree_node *child, bool right);
 
 bool vieter_tree_node_validate(vieter_tree_node *node, uint64_t passed_black_nodes, uint64_t expected_black_nodes);