From 17c2e153851697ce9e2f9969d15a20ffcaad6ff5 Mon Sep 17 00:00:00 2001
From: Jef Roosens <roosensjef@gmail.com>
Date: Mon, 30 Jan 2023 12:28:08 +0100
Subject: [PATCH] feat(tree): first attempt at wikipedia pseudo code

---
 src/tree/vieter_tree_balancing.c | 92 ++++++++++++++++++++++++++++++++
 1 file changed, 92 insertions(+)

diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c
index e2a9d18..710203b 100644
--- a/src/tree/vieter_tree_balancing.c
+++ b/src/tree/vieter_tree_balancing.c
@@ -124,6 +124,22 @@ void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
   }
 }
 
+vieter_tree_node *vieter_tree_node_rotate(vieter_tree_node *old_root,
+                                          bool dir) {
+  vieter_tree_node *new_root = old_root->children[1 - dir];
+
+  // Right rotation
+  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 {
+    vieter_tree_node_set_child(old_root, new_root->children[0], true);
+    vieter_tree_node_set_child(new_root, old_root, false);
+  }
+
+  return new_root;
+}
+
 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
@@ -145,4 +161,80 @@ void vieter_tree_node_remove_balanced(vieter_tree_node *node) {
   }
 
   // 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
+
+  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);
+  parent->children[dir] = NULL;
+
+  vieter_tree_node_free(node);
+
+  goto start;
+
+  do {
+    dir = vieter_tree_node_get(node, vieter_tree_node_right);
+
+  start:
+    sibling = parent->children[1 - dir];
+    distant_nephew = sibling->children[1 - dir];
+    close_nephew = sibling->children[dir];
+
+    if (!vieter_tree_node_get(sibling, vieter_tree_node_black))
+      goto case3;
+
+    if (distant_nephew != NULL &&
+        !vieter_tree_node_get(distant_nephew, vieter_tree_node_black))
+      goto case6;
+
+    if (close_nephew != NULL &&
+        !vieter_tree_node_get(close_nephew, vieter_tree_node_black))
+      goto case5;
+
+    if (!vieter_tree_node_get(parent, vieter_tree_node_black))
+      goto case4;
+
+    // Case 2
+    vieter_tree_node_set(sibling, vieter_tree_node_black, false);
+    node = parent;
+  } while ((parent = node->parent) != NULL);
+
+  // Case 1
+  return;
+
+case3:
+  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;
+  distant_nephew = sibling->children[1 - dir];
+
+  if (distant_nephew != NULL &&
+      !vieter_tree_node_get(distant_nephew, vieter_tree_node_black))
+    goto case6;
+
+  close_nephew = sibling->children[dir];
+
+  if (close_nephew != NULL &&
+      !vieter_tree_node_get(close_nephew, vieter_tree_node_black))
+    goto case5;
+
+case4:
+  vieter_tree_node_set(sibling, vieter_tree_node_black, false);
+  vieter_tree_node_set(parent, vieter_tree_node_black, true);
+  return;
+
+case5:
+  vieter_tree_node_rotate(sibling, 1 - dir);
+  vieter_tree_node_set(sibling, vieter_tree_node_black, false);
+  vieter_tree_node_set(close_nephew, vieter_tree_node_black, true);
+  distant_nephew = sibling;
+  sibling = close_nephew;
+
+case6:
+  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);
 }