From 06fca6f5ae68b8fc9e3f5f9790fedec54e70d172 Mon Sep 17 00:00:00 2001
From: Chewing_Bever <roosensjef@gmail.com>
Date: Sun, 29 Jan 2023 17:32:27 +0100
Subject: [PATCH 1/4] refactor(tree): better testing code

---
 src/tree/vieter_tree.c           | 11 +++---
 src/tree/vieter_tree_balancing.c | 68 ++++++++++++++------------------
 src/tree/vieter_tree_balancing.h | 11 ++++++
 src/tree/vieter_tree_internal.h  |  6 ++-
 src/tree/vieter_tree_node.c      |  7 ++--
 src/tree/vieter_tree_node.h      | 22 ++++++-----
 test/tree/test_tree.c            |  6 +--
 7 files changed, 70 insertions(+), 61 deletions(-)

diff --git a/src/tree/vieter_tree.c b/src/tree/vieter_tree.c
index 792ad30..bd0200e 100644
--- a/src/tree/vieter_tree.c
+++ b/src/tree/vieter_tree.c
@@ -16,7 +16,7 @@ vieter_tree_error vieter_tree_insert(vieter_tree *tree, uint64_t key,
     tree->root = vieter_tree_node_init();
     tree->root->key = key;
     tree->root->data = data;
-    vieter_tree_node_set(tree->root, vieter_node_black);
+    vieter_tree_node_set(tree->root, vieter_tree_node_black);
     tree->size = 1;
 
     return vieter_tree_ok;
@@ -99,9 +99,9 @@ end:
   free(tree);
 }
 
-int vieter_tree_validate(vieter_tree *tree) {
+bool vieter_tree_validate(vieter_tree *tree) {
   if (tree->size == 0) {
-    return 0;
+    return true;
   }
 
   // DFS to get expected black nodes
@@ -109,12 +109,13 @@ int vieter_tree_validate(vieter_tree *tree) {
   vieter_tree_node *node = tree->root;
 
   while (node != NULL) {
-    if (vieter_tree_node_get(node, vieter_node_black)) {
+    if (vieter_tree_node_get(node, vieter_tree_node_black)) {
       expected_black_nodes++;
     }
 
     node = node->left;
   }
 
-  return vieter_tree_node_validate(tree->root, 0, expected_black_nodes);
+  return vieter_tree_node_get(tree->root, vieter_tree_node_black) &&
+         vieter_tree_node_validate(tree->root, 0, expected_black_nodes);
 }
diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c
index 70c52e0..d9fc66e 100644
--- a/src/tree/vieter_tree_balancing.c
+++ b/src/tree/vieter_tree_balancing.c
@@ -1,41 +1,31 @@
 #include "vieter_tree_balancing.h"
 
-int vieter_tree_node_validate(vieter_tree_node *node,
-                              uint64_t passed_black_nodes,
-                              uint64_t expected_black_nodes) {
-  if (vieter_tree_node_get(node, vieter_node_black)) {
+bool vieter_tree_node_validate(vieter_tree_node *node,
+                               uint64_t passed_black_nodes,
+                               uint64_t expected_black_nodes) {
+  // Path to NULL children should contain the same amount of black nodes
+  if (node == NULL) {
+    return passed_black_nodes == expected_black_nodes;
+  }
+
+  if (vieter_tree_node_get(node, vieter_tree_node_black)) {
     passed_black_nodes++;
-  } else {
-    // A red node should have black child nodes
-    if ((node->left != NULL &&
-         !vieter_tree_node_get(node->left, vieter_node_black)) ||
-        (node->right != NULL &&
-         !vieter_tree_node_get(node->right, vieter_node_black))) {
-      return 1;
-    }
   }
 
-  // All paths to a NULL child should have the same amount of black nodes
-  if ((node->left == NULL || node->right == NULL) &&
-      passed_black_nodes != expected_black_nodes) {
-    return 1;
-  }
+  // Either the node itself is black, or its children are both either NULL or
+  // black.
+  bool correctly_colored_children =
+      vieter_tree_node_get(node, vieter_tree_node_black) ||
+      ((node->left == NULL ||
+        vieter_tree_node_get(node->left, vieter_tree_node_black)) &&
+       (node->right == NULL ||
+        vieter_tree_node_get(node->right, vieter_tree_node_black)));
 
-  if (node->left != NULL) {
-    if (vieter_tree_node_validate(node->left, passed_black_nodes,
-                                  expected_black_nodes) != 0) {
-      return 1;
-    }
-  }
-
-  if (node->right != NULL) {
-    if (vieter_tree_node_validate(node->right, passed_black_nodes,
-                                  expected_black_nodes) != 0) {
-      return 1;
-    }
-  }
-
-  return 0;
+  return correctly_colored_children &&
+         vieter_tree_node_validate(node->left, passed_black_nodes,
+                                   expected_black_nodes) &&
+         vieter_tree_node_validate(node->right, passed_black_nodes,
+                                   expected_black_nodes);
 }
 
 vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
@@ -106,9 +96,9 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
     children[3] = node->right;
   }
 
-  vieter_tree_node_set(left, vieter_node_black);
-  vieter_tree_node_set(right, vieter_node_black);
-  vieter_tree_node_unset(root, vieter_node_black);
+  vieter_tree_node_set(left, vieter_tree_node_black);
+  vieter_tree_node_set(right, vieter_tree_node_black);
+  vieter_tree_node_unset(root, vieter_tree_node_black);
 
   left->left = children[0];
   left->right = children[1];
@@ -137,14 +127,16 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
 }
 
 void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
-  while (!vieter_tree_node_get(node, vieter_node_black) &&
+  while (!vieter_tree_node_get(node, vieter_tree_node_black) &&
          node->parent != NULL &&
-         !vieter_tree_node_get(node->parent, vieter_node_black)) {
+         !vieter_tree_node_get(node->parent, vieter_tree_node_black)) {
     node = vieter_tree_node_balance(node);
   }
 
   // The root must always be black
   if (node->parent == NULL) {
-    vieter_tree_node_set(node, vieter_node_black);
+    vieter_tree_node_set(node, vieter_tree_node_black);
   }
 }
+
+void vieter_tree_node_remove_balanced(vieter_tree_node *node) {}
diff --git a/src/tree/vieter_tree_balancing.h b/src/tree/vieter_tree_balancing.h
index 4492588..007af34 100644
--- a/src/tree/vieter_tree_balancing.h
+++ b/src/tree/vieter_tree_balancing.h
@@ -3,6 +3,17 @@
 
 #include "vieter_tree_node.h"
 
+/*
+ * Ensure the tree remains a valid red-black tree after having inserting the
+ * node.
+ */
 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.
+ */
+void vieter_tree_node_remove_balanced(vieter_tree_node *node);
+
 #endif
diff --git a/src/tree/vieter_tree_internal.h b/src/tree/vieter_tree_internal.h
index dc4ee81..670d45b 100644
--- a/src/tree/vieter_tree_internal.h
+++ b/src/tree/vieter_tree_internal.h
@@ -1,8 +1,10 @@
 #include "vieter_tree.h"
 
+#include <stdbool.h>
+
 /*
  * Inspect whether the tree is still a valid red-black-tree.
  *
- * @return 0 if correct, 1 if not.
+ * @return true if valid, false otherwise.
  */
-int vieter_tree_validate(vieter_tree *tree);
+bool vieter_tree_validate(vieter_tree *tree);
diff --git a/src/tree/vieter_tree_node.c b/src/tree/vieter_tree_node.c
index 40407a3..b7e82a8 100644
--- a/src/tree/vieter_tree_node.c
+++ b/src/tree/vieter_tree_node.c
@@ -153,14 +153,15 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
   return vieter_tree_ok;
 }
 
-void vieter_tree_node_set(vieter_tree_node *node, vieter_node_flag flag) {
+void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag) {
   node->flags |= flag;
 }
 
-void vieter_tree_node_unset(vieter_tree_node *node, vieter_node_flag flag) {
+void vieter_tree_node_unset(vieter_tree_node *node,
+                            vieter_tree_node_flag flag) {
   node->flags &= ~flag;
 }
 
-bool vieter_tree_node_get(vieter_tree_node *node, vieter_node_flag flag) {
+bool vieter_tree_node_get(vieter_tree_node *node, vieter_tree_node_flag flag) {
   return (node->flags & flag) != 0;
 }
diff --git a/src/tree/vieter_tree_node.h b/src/tree/vieter_tree_node.h
index 1c88159..5e74280 100644
--- a/src/tree/vieter_tree_node.h
+++ b/src/tree/vieter_tree_node.h
@@ -5,11 +5,9 @@
 #include <stdint.h>
 #include <stdbool.h>
 
-#define VIETER_NODE_BLACK ((uint8_t)1)
-
-typedef enum vieter_node_flag {
-    vieter_node_black = ((uint8_t)1) << 0
-} vieter_node_flag;
+typedef enum vieter_tree_node_flag {
+    vieter_tree_node_black = ((uint8_t)1) << 0
+} vieter_tree_node_flag;
 
 typedef struct vieter_tree_node {
     uint64_t key;
@@ -26,20 +24,24 @@ void vieter_tree_node_free(vieter_tree_node *node);
 
 vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key, void *data);
 
+/*
+ * Return the node representing the requested value.
+ */
 vieter_tree_error vieter_tree_node_search_node(vieter_tree_node **out, vieter_tree_node *root, uint64_t key);
+
 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);
 
-void vieter_tree_node_replace(vieter_tree_node *to_replace, vieter_tree_node *replacement);
+void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace, vieter_tree_node *replacement);
 
-void vieter_tree_node_set(vieter_tree_node *node, vieter_node_flag flag);
-void vieter_tree_node_unset(vieter_tree_node *node, vieter_node_flag flag);
-bool vieter_tree_node_get(vieter_tree_node *node, vieter_node_flag flag);
+void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag);
+void vieter_tree_node_unset(vieter_tree_node *node, vieter_tree_node_flag flag);
+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);
 
-int vieter_tree_node_validate(vieter_tree_node *node, uint64_t passed_black_nodes, uint64_t expected_black_nodes);
+bool vieter_tree_node_validate(vieter_tree_node *node, uint64_t passed_black_nodes, uint64_t expected_black_nodes);
 
 #endif
diff --git a/test/tree/test_tree.c b/test/tree/test_tree.c
index c875137..257ec66 100644
--- a/test/tree/test_tree.c
+++ b/test/tree/test_tree.c
@@ -18,7 +18,7 @@ void test_insert() {
     for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_insert(tree, i, NULL) == vieter_tree_ok);
         TEST_SIZE(tree, i + 1);
-        TEST_CHECK(vieter_tree_validate(tree) == 0);
+        TEST_CHECK(vieter_tree_validate(tree));
     }
 
     void *out;
@@ -37,7 +37,7 @@ void test_remove() {
 
     for (uint64_t i = 0; i < 250; i++) {
         TEST_CHECK(vieter_tree_insert(tree, i, NULL) == vieter_tree_ok);
-        TEST_CHECK(vieter_tree_validate(tree) == 0);
+        TEST_CHECK(vieter_tree_validate(tree));
     }
 
     void *out;
@@ -45,7 +45,7 @@ void test_remove() {
     for (uint64_t i = 0; i < 25; 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) == 0);
+        TEST_CHECK(vieter_tree_validate(tree));
         TEST_CHECK(vieter_tree_search(&out, tree, i) == vieter_tree_not_present);
         TEST_SIZE(tree, 250 - i - 1);
     }

From 193fc4b454135d67d331afcebc7f956a62ab08fb Mon Sep 17 00:00:00 2001
From: Chewing_Bever <roosensjef@gmail.com>
Date: Sun, 29 Jan 2023 20:33:11 +0100
Subject: [PATCH 2/4] refactor(tree): store children as static array instead

---
 src/tree/vieter_tree.c           | 10 +++---
 src/tree/vieter_tree_balancing.c | 56 ++++++++++++++++----------------
 src/tree/vieter_tree_node.c      | 45 +++++++++----------------
 src/tree/vieter_tree_node.h      |  8 +++--
 4 files changed, 54 insertions(+), 65 deletions(-)

diff --git a/src/tree/vieter_tree.c b/src/tree/vieter_tree.c
index bd0200e..92f8831 100644
--- a/src/tree/vieter_tree.c
+++ b/src/tree/vieter_tree.c
@@ -80,13 +80,13 @@ void vieter_tree_free(vieter_tree *tree) {
       stack = realloc(stack, capacity * sizeof(vieter_tree_node *));
     }
 
-    if (node->left != NULL) {
-      stack[size] = node->left;
+    if (node->children[0] != NULL) {
+      stack[size] = node->children[0];
       size++;
     }
 
-    if (node->right != NULL) {
-      stack[size] = node->right;
+    if (node->children[1] != NULL) {
+      stack[size] = node->children[1];
       size++;
     }
 
@@ -113,7 +113,7 @@ bool vieter_tree_validate(vieter_tree *tree) {
       expected_black_nodes++;
     }
 
-    node = node->left;
+    node = node->children[0];
   }
 
   return vieter_tree_node_get(tree->root, vieter_tree_node_black) &&
diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c
index d9fc66e..c04825d 100644
--- a/src/tree/vieter_tree_balancing.c
+++ b/src/tree/vieter_tree_balancing.c
@@ -16,15 +16,15 @@ bool vieter_tree_node_validate(vieter_tree_node *node,
   // black.
   bool correctly_colored_children =
       vieter_tree_node_get(node, vieter_tree_node_black) ||
-      ((node->left == NULL ||
-        vieter_tree_node_get(node->left, vieter_tree_node_black)) &&
-       (node->right == NULL ||
-        vieter_tree_node_get(node->right, vieter_tree_node_black)));
+      ((node->children[0] == NULL ||
+        vieter_tree_node_get(node->children[0], vieter_tree_node_black)) &&
+       (node->children[1] == NULL ||
+        vieter_tree_node_get(node->children[1], vieter_tree_node_black)));
 
   return correctly_colored_children &&
-         vieter_tree_node_validate(node->left, passed_black_nodes,
+         vieter_tree_node_validate(node->children[0], passed_black_nodes,
                                    expected_black_nodes) &&
-         vieter_tree_node_validate(node->right, passed_black_nodes,
+         vieter_tree_node_validate(node->children[1], passed_black_nodes,
                                    expected_black_nodes);
 }
 
@@ -48,10 +48,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
     left = node;
     right = parent;
 
-    children[0] = node->left;
-    children[1] = node->right;
-    children[2] = parent->right;
-    children[3] = grand_parent->right;
+    children[0] = node->children[0];
+    children[1] = node->children[1];
+    children[2] = parent->children[1];
+    children[3] = grand_parent->children[1];
   } else if (node->key < parent->key && parent->key > grand_parent->key) {
     root->key = node->key;
     root->data = node->data;
@@ -62,10 +62,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
     left = node;
     right = parent;
 
-    children[0] = grand_parent->left;
-    children[1] = node->left;
-    children[2] = node->right;
-    children[3] = parent->right;
+    children[0] = grand_parent->children[0];
+    children[1] = node->children[0];
+    children[2] = node->children[1];
+    children[3] = parent->children[1];
   } else if (node->key > parent->key && parent->key < grand_parent->key) {
     root->key = node->key;
     root->data = node->data;
@@ -76,10 +76,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
     left = parent;
     right = node;
 
-    children[0] = parent->left;
-    children[1] = node->left;
-    children[2] = node->right;
-    children[3] = grand_parent->right;
+    children[0] = parent->children[0];
+    children[1] = node->children[0];
+    children[2] = node->children[1];
+    children[3] = grand_parent->children[1];
   } else {
     root->key = parent->key;
     root->data = parent->data;
@@ -90,18 +90,18 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
     left = parent;
     right = node;
 
-    children[0] = grand_parent->left;
-    children[1] = parent->left;
-    children[2] = node->left;
-    children[3] = node->right;
+    children[0] = grand_parent->children[0];
+    children[1] = parent->children[0];
+    children[2] = node->children[0];
+    children[3] = node->children[1];
   }
 
   vieter_tree_node_set(left, vieter_tree_node_black);
   vieter_tree_node_set(right, vieter_tree_node_black);
   vieter_tree_node_unset(root, vieter_tree_node_black);
 
-  left->left = children[0];
-  left->right = children[1];
+  left->children[0] = children[0];
+  left->children[1] = children[1];
 
   if (children[0] != NULL)
     children[0]->parent = left;
@@ -109,8 +109,8 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
   if (children[1] != NULL)
     children[1]->parent = left;
 
-  right->left = children[2];
-  right->right = children[3];
+  right->children[0] = children[2];
+  right->children[1] = children[3];
 
   if (children[2] != NULL)
     children[2]->parent = right;
@@ -118,8 +118,8 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
   if (children[3] != NULL)
     children[3]->parent = right;
 
-  root->left = left;
-  root->right = right;
+  root->children[0] = left;
+  root->children[1] = right;
   left->parent = root;
   right->parent = root;
 
diff --git a/src/tree/vieter_tree_node.c b/src/tree/vieter_tree_node.c
index b7e82a8..ca506b9 100644
--- a/src/tree/vieter_tree_node.c
+++ b/src/tree/vieter_tree_node.c
@@ -13,11 +13,7 @@ void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
     return;
   }
 
-  if (key < parent->key) {
-    parent->left = child;
-  } else {
-    parent->right = child;
-  }
+  parent->children[key > parent->key] = child;
 
   if (child != NULL) {
     child->parent = parent;
@@ -28,14 +24,14 @@ void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace,
                                          vieter_tree_node *replacement) {
   to_replace->key = replacement->key;
   to_replace->data = replacement->data;
-  to_replace->left = replacement->left;
-  to_replace->right = replacement->right;
+  to_replace->children[0] = replacement->children[0];
+  to_replace->children[1] = replacement->children[1];
 
-  if (to_replace->left != NULL)
-    to_replace->left->parent = to_replace;
+  if (to_replace->children[0] != NULL)
+    to_replace->children[0]->parent = to_replace;
 
-  if (to_replace->right != NULL)
-    to_replace->right->parent = to_replace;
+  if (to_replace->children[1] != NULL)
+    to_replace->children[1]->parent = to_replace;
 }
 
 vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key,
@@ -49,12 +45,7 @@ vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key,
     }
 
     parent = node;
-
-    if (key < parent->key) {
-      node = parent->left;
-    } else {
-      node = parent->right;
-    }
+    node = node->children[key > parent->key];
   }
 
   vieter_tree_node *new_node = vieter_tree_node_init();
@@ -79,11 +70,7 @@ vieter_tree_error vieter_tree_node_search_node(vieter_tree_node **out,
       return vieter_tree_ok;
     }
 
-    if (key < node->key) {
-      node = node->left;
-    } else {
-      node = node->right;
-    }
+    node = node->children[key > node->key];
   }
 
   return vieter_tree_not_present;
@@ -114,13 +101,13 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
 
   *out = target->data;
 
-  if (target->left == NULL && target->right == NULL) {
+  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->left == NULL) ^ (target->right == NULL)) {
+  } else if ((target->children[0] == NULL) ^ (target->children[1] == NULL)) {
     vieter_tree_node *child =
-        target->left != NULL ? target->left : target->right;
+        target->children[0] != NULL ? target->children[0] : target->children[1];
 
     if (target->parent != NULL) {
       vieter_tree_node_add_child(target->parent, child->key, child);
@@ -131,10 +118,10 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
     }
 
   } else {
-    vieter_tree_node *replacement = target->right;
+    vieter_tree_node *replacement = target->children[1];
 
-    while (replacement->left != NULL) {
-      replacement = replacement->left;
+    while (replacement->children[0] != NULL) {
+      replacement = replacement->children[0];
     }
 
     // We use replacement->key here because the right child can be NULL, so
@@ -142,7 +129,7 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
     // will still get placed into the right location because of how binary
     // trees work.
     vieter_tree_node_add_child(replacement->parent, replacement->key,
-                               replacement->right);
+                               replacement->children[1]);
 
     target->key = replacement->key;
     target->data = replacement->data;
diff --git a/src/tree/vieter_tree_node.h b/src/tree/vieter_tree_node.h
index 5e74280..21bbb81 100644
--- a/src/tree/vieter_tree_node.h
+++ b/src/tree/vieter_tree_node.h
@@ -6,15 +6,17 @@
 #include <stdbool.h>
 
 typedef enum vieter_tree_node_flag {
-    vieter_tree_node_black = ((uint8_t)1) << 0
+  vieter_tree_node_black = ((uint8_t)1) << 0,
+  // This flag is used to index the children array. If this flag isn't set,
+  // it'll evaluate to false, or index 0 in the children array.
+  vieter_tree_node_right = ((uint8_t)1) << 1
 } vieter_tree_node_flag;
 
 typedef struct vieter_tree_node {
     uint64_t key;
     void *data;
     struct vieter_tree_node *parent;
-    struct vieter_tree_node *left;
-    struct vieter_tree_node *right;
+    struct vieter_tree_node *children[2];
     uint8_t flags;
 } vieter_tree_node;
 

From e5130a638377751197c5f574533efe2325e23591 Mon Sep 17 00:00:00 2001
From: Chewing_Bever <roosensjef@gmail.com>
Date: Sun, 29 Jan 2023 20:57:08 +0100
Subject: [PATCH 3/4] feat(tree): bit flag so children know which child they
 are

---
 src/tree/vieter_tree.c           |   2 +-
 src/tree/vieter_tree_balancing.c | 134 +++++++++++++++++--------------
 src/tree/vieter_tree_node.c      |  24 +++---
 src/tree/vieter_tree_node.h      |   3 +-
 4 files changed, 89 insertions(+), 74 deletions(-)

diff --git a/src/tree/vieter_tree.c b/src/tree/vieter_tree.c
index 92f8831..b066af1 100644
--- a/src/tree/vieter_tree.c
+++ b/src/tree/vieter_tree.c
@@ -16,7 +16,7 @@ vieter_tree_error vieter_tree_insert(vieter_tree *tree, uint64_t key,
     tree->root = vieter_tree_node_init();
     tree->root->key = key;
     tree->root->data = data;
-    vieter_tree_node_set(tree->root, vieter_tree_node_black);
+    vieter_tree_node_set(tree->root, vieter_tree_node_black, true);
     tree->size = 1;
 
     return vieter_tree_ok;
diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c
index c04825d..6a58179 100644
--- a/src/tree/vieter_tree_balancing.c
+++ b/src/tree/vieter_tree_balancing.c
@@ -1,4 +1,5 @@
 #include "vieter_tree_balancing.h"
+#include <string.h>
 
 bool vieter_tree_node_validate(vieter_tree_node *node,
                                uint64_t passed_black_nodes,
@@ -38,85 +39,96 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
   uint64_t key_root = root->key;
   void *data_root = root->data;
 
-  if (node->key < parent->key && parent->key < grand_parent->key) {
-    root->key = parent->key;
-    root->data = parent->data;
-
-    parent->key = key_root;
-    parent->data = data_root;
-
-    left = node;
-    right = parent;
-
-    children[0] = node->children[0];
-    children[1] = node->children[1];
-    children[2] = parent->children[1];
-    children[3] = grand_parent->children[1];
-  } else if (node->key < parent->key && parent->key > grand_parent->key) {
-    root->key = node->key;
-    root->data = node->data;
-
-    node->key = key_root;
-    node->data = data_root;
-
-    left = node;
-    right = parent;
-
-    children[0] = grand_parent->children[0];
-    children[1] = node->children[0];
-    children[2] = node->children[1];
-    children[3] = parent->children[1];
-  } else if (node->key > parent->key && parent->key < grand_parent->key) {
-    root->key = node->key;
-    root->data = node->data;
-
-    node->key = key_root;
-    node->data = data_root;
-
+  if (vieter_tree_node_get(node, vieter_tree_node_right)) {
     left = parent;
     right = node;
 
-    children[0] = parent->children[0];
-    children[1] = node->children[0];
-    children[2] = node->children[1];
-    children[3] = grand_parent->children[1];
+    if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
+      root->key = parent->key;
+      root->data = parent->data;
+
+      parent->key = key_root;
+      parent->data = data_root;
+
+      children[0] = grand_parent->children[0];
+      children[1] = parent->children[0];
+      children[2] = node->children[0];
+      children[3] = node->children[1];
+    } else {
+      root->key = node->key;
+      root->data = node->data;
+
+      node->key = key_root;
+      node->data = data_root;
+
+      children[0] = parent->children[0];
+      children[1] = node->children[0];
+      children[2] = node->children[1];
+      children[3] = grand_parent->children[1];
+    }
   } else {
-    root->key = parent->key;
-    root->data = parent->data;
+    left = node;
+    right = parent;
 
-    parent->key = key_root;
-    parent->data = data_root;
+    if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
+      root->key = node->key;
+      root->data = node->data;
 
-    left = parent;
-    right = node;
+      node->key = key_root;
+      node->data = data_root;
 
-    children[0] = grand_parent->children[0];
-    children[1] = parent->children[0];
-    children[2] = node->children[0];
-    children[3] = node->children[1];
+      children[0] = grand_parent->children[0];
+      children[1] = node->children[0];
+      children[2] = node->children[1];
+      children[3] = parent->children[1];
+    } else {
+      root->key = parent->key;
+      root->data = parent->data;
+
+      parent->key = key_root;
+      parent->data = data_root;
+
+      children[0] = node->children[0];
+      children[1] = node->children[1];
+      children[2] = parent->children[1];
+      children[3] = grand_parent->children[1];
+    }
   }
 
-  vieter_tree_node_set(left, vieter_tree_node_black);
-  vieter_tree_node_set(right, vieter_tree_node_black);
-  vieter_tree_node_unset(root, vieter_tree_node_black);
+  vieter_tree_node_set(
+      root, vieter_tree_node_right,
+      vieter_tree_node_get(grand_parent, vieter_tree_node_right));
 
-  left->children[0] = children[0];
-  left->children[1] = children[1];
+  vieter_tree_node_set(left, vieter_tree_node_black, true);
+  vieter_tree_node_set(left, vieter_tree_node_right, false);
+  vieter_tree_node_set(right, vieter_tree_node_black, true);
+  vieter_tree_node_set(left, vieter_tree_node_right, true);
 
-  if (children[0] != NULL)
+  vieter_tree_node_set(root, vieter_tree_node_black, false);
+
+  memcpy(&left->children, children, 2 * sizeof(vieter_tree_node *));
+
+  if (children[0] != NULL) {
     children[0]->parent = left;
+    vieter_tree_node_set(children[0], vieter_tree_node_right, false);
+  }
 
-  if (children[1] != NULL)
+  if (children[1] != NULL) {
     children[1]->parent = left;
+    vieter_tree_node_set(children[0], vieter_tree_node_right, true);
+  }
 
-  right->children[0] = children[2];
-  right->children[1] = children[3];
+  memcpy(&right->children, children + 2, 2 * sizeof(vieter_tree_node *));
 
-  if (children[2] != NULL)
+  if (children[2] != NULL) {
     children[2]->parent = right;
+    vieter_tree_node_set(children[2], vieter_tree_node_right, false);
+  }
 
-  if (children[3] != NULL)
+  if (children[3] != NULL) {
     children[3]->parent = right;
+    vieter_tree_node_set(children[2], vieter_tree_node_right, true);
+  }
 
   root->children[0] = left;
   root->children[1] = right;
@@ -135,7 +147,7 @@ void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
 
   // The root must always be black
   if (node->parent == NULL) {
-    vieter_tree_node_set(node, vieter_tree_node_black);
+    vieter_tree_node_set(node, vieter_tree_node_black, true);
   }
 }
 
diff --git a/src/tree/vieter_tree_node.c b/src/tree/vieter_tree_node.c
index ca506b9..853281a 100644
--- a/src/tree/vieter_tree_node.c
+++ b/src/tree/vieter_tree_node.c
@@ -1,6 +1,8 @@
 #include "vieter_tree_node.h"
 #include "vieter_tree_balancing.h"
 
+#include <string.h>
+
 vieter_tree_node *vieter_tree_node_init() {
   return calloc(1, sizeof(vieter_tree_node));
 }
@@ -13,10 +15,12 @@ void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
     return;
   }
 
-  parent->children[key > parent->key] = child;
+  bool side = key > parent->key;
+  parent->children[side] = child;
 
   if (child != NULL) {
     child->parent = parent;
+    vieter_tree_node_set(child, vieter_tree_node_right, side);
   }
 }
 
@@ -24,8 +28,8 @@ void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace,
                                          vieter_tree_node *replacement) {
   to_replace->key = replacement->key;
   to_replace->data = replacement->data;
-  to_replace->children[0] = replacement->children[0];
-  to_replace->children[1] = replacement->children[1];
+  memcpy(to_replace->children, replacement->children,
+         2 * sizeof(vieter_tree_node *));
 
   if (to_replace->children[0] != NULL)
     to_replace->children[0]->parent = to_replace;
@@ -140,13 +144,13 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
   return vieter_tree_ok;
 }
 
-void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag) {
-  node->flags |= flag;
-}
-
-void vieter_tree_node_unset(vieter_tree_node *node,
-                            vieter_tree_node_flag flag) {
-  node->flags &= ~flag;
+void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag,
+                          bool set) {
+  if (set) {
+    node->flags |= flag;
+  } else {
+    node->flags &= ~flag;
+  }
 }
 
 bool vieter_tree_node_get(vieter_tree_node *node, vieter_tree_node_flag flag) {
diff --git a/src/tree/vieter_tree_node.h b/src/tree/vieter_tree_node.h
index 21bbb81..4012d1d 100644
--- a/src/tree/vieter_tree_node.h
+++ b/src/tree/vieter_tree_node.h
@@ -37,8 +37,7 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root, ui
 
 void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace, vieter_tree_node *replacement);
 
-void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag);
-void vieter_tree_node_unset(vieter_tree_node *node, vieter_tree_node_flag flag);
+void vieter_tree_node_set(vieter_tree_node *node, vieter_tree_node_flag flag, bool value);
 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,

From b74a4e932651466259ba02d34b114d2037864aca Mon Sep 17 00:00:00 2001
From: Chewing_Bever <roosensjef@gmail.com>
Date: Sun, 29 Jan 2023 21:08:17 +0100
Subject: [PATCH 4/4] refactor(tree): bit more cleanup

---
 src/tree/vieter_tree_balancing.c | 81 +++++++++++---------------------
 src/tree/vieter_tree_node.c      | 21 +++++----
 src/tree/vieter_tree_node.h      |  1 +
 3 files changed, 41 insertions(+), 62 deletions(-)

diff --git a/src/tree/vieter_tree_balancing.c b/src/tree/vieter_tree_balancing.c
index 6a58179..b160563 100644
--- a/src/tree/vieter_tree_balancing.c
+++ b/src/tree/vieter_tree_balancing.c
@@ -33,15 +33,15 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
   vieter_tree_node *parent = node->parent;
   vieter_tree_node *grand_parent = parent->parent;
   vieter_tree_node *root = grand_parent;
-  vieter_tree_node *left, *right;
-  vieter_tree_node *children[4];
+  vieter_tree_node *children[2];
+  vieter_tree_node *grand_children[4];
 
   uint64_t key_root = root->key;
   void *data_root = root->data;
 
   if (vieter_tree_node_get(node, vieter_tree_node_right)) {
-    left = parent;
-    right = node;
+    children[0] = parent;
+    children[1] = node;
 
     if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
       root->key = parent->key;
@@ -50,10 +50,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       parent->key = key_root;
       parent->data = data_root;
 
-      children[0] = grand_parent->children[0];
-      children[1] = parent->children[0];
-      children[2] = node->children[0];
-      children[3] = node->children[1];
+      grand_children[0] = grand_parent->children[0];
+      grand_children[1] = parent->children[0];
+      grand_children[2] = node->children[0];
+      grand_children[3] = node->children[1];
     } else {
       root->key = node->key;
       root->data = node->data;
@@ -61,14 +61,14 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       node->key = key_root;
       node->data = data_root;
 
-      children[0] = parent->children[0];
-      children[1] = node->children[0];
-      children[2] = node->children[1];
-      children[3] = grand_parent->children[1];
+      grand_children[0] = parent->children[0];
+      grand_children[1] = node->children[0];
+      grand_children[2] = node->children[1];
+      grand_children[3] = grand_parent->children[1];
     }
   } else {
-    left = node;
-    right = parent;
+    children[0] = node;
+    children[1] = parent;
 
     if (vieter_tree_node_get(parent, vieter_tree_node_right)) {
       root->key = node->key;
@@ -77,10 +77,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       node->key = key_root;
       node->data = data_root;
 
-      children[0] = grand_parent->children[0];
-      children[1] = node->children[0];
-      children[2] = node->children[1];
-      children[3] = parent->children[1];
+      grand_children[0] = grand_parent->children[0];
+      grand_children[1] = node->children[0];
+      grand_children[2] = node->children[1];
+      grand_children[3] = parent->children[1];
     } else {
       root->key = parent->key;
       root->data = parent->data;
@@ -88,10 +88,10 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       parent->key = key_root;
       parent->data = data_root;
 
-      children[0] = node->children[0];
-      children[1] = node->children[1];
-      children[2] = parent->children[1];
-      children[3] = grand_parent->children[1];
+      grand_children[0] = node->children[0];
+      grand_children[1] = node->children[1];
+      grand_children[2] = parent->children[1];
+      grand_children[3] = grand_parent->children[1];
     }
   }
 
@@ -99,41 +99,14 @@ vieter_tree_node *vieter_tree_node_balance(vieter_tree_node *node) {
       root, vieter_tree_node_right,
       vieter_tree_node_get(grand_parent, vieter_tree_node_right));
 
-  vieter_tree_node_set(left, vieter_tree_node_black, true);
-  vieter_tree_node_set(left, vieter_tree_node_right, false);
-  vieter_tree_node_set(right, vieter_tree_node_black, true);
-  vieter_tree_node_set(left, vieter_tree_node_right, true);
+  vieter_tree_node_set(children[0], vieter_tree_node_black, true);
+  vieter_tree_node_set(children[1], vieter_tree_node_black, true);
 
   vieter_tree_node_set(root, vieter_tree_node_black, false);
 
-  memcpy(&left->children, children, 2 * sizeof(vieter_tree_node *));
-
-  if (children[0] != NULL) {
-    children[0]->parent = left;
-    vieter_tree_node_set(children[0], vieter_tree_node_right, false);
-  }
-
-  if (children[1] != NULL) {
-    children[1]->parent = left;
-    vieter_tree_node_set(children[0], vieter_tree_node_right, true);
-  }
-
-  memcpy(&right->children, children + 2, 2 * sizeof(vieter_tree_node *));
-
-  if (children[2] != NULL) {
-    children[2]->parent = right;
-    vieter_tree_node_set(children[2], vieter_tree_node_right, false);
-  }
-
-  if (children[3] != NULL) {
-    children[3]->parent = right;
-    vieter_tree_node_set(children[2], vieter_tree_node_right, true);
-  }
-
-  root->children[0] = left;
-  root->children[1] = right;
-  left->parent = root;
-  right->parent = root;
+  vieter_tree_node_set_children(children[0], grand_children);
+  vieter_tree_node_set_children(children[1], grand_children + 2);
+  vieter_tree_node_set_children(root, children);
 
   return root;
 }
diff --git a/src/tree/vieter_tree_node.c b/src/tree/vieter_tree_node.c
index 853281a..2f45b82 100644
--- a/src/tree/vieter_tree_node.c
+++ b/src/tree/vieter_tree_node.c
@@ -24,18 +24,23 @@ void vieter_tree_node_add_child(vieter_tree_node *parent, uint64_t key,
   }
 }
 
+void vieter_tree_node_set_children(vieter_tree_node *parent,
+                                   vieter_tree_node **children) {
+  memcpy(parent->children, children, 2 * sizeof(vieter_tree_node *));
+
+  for (uint8_t i = 0; i < 2; i++) {
+    if (parent->children[i] != NULL) {
+      parent->children[i]->parent = parent;
+      vieter_tree_node_set(parent->children[i], vieter_tree_node_right, i);
+    }
+  }
+}
+
 void vieter_tree_node_replace_with_child(vieter_tree_node *to_replace,
                                          vieter_tree_node *replacement) {
   to_replace->key = replacement->key;
   to_replace->data = replacement->data;
-  memcpy(to_replace->children, replacement->children,
-         2 * sizeof(vieter_tree_node *));
-
-  if (to_replace->children[0] != NULL)
-    to_replace->children[0]->parent = to_replace;
-
-  if (to_replace->children[1] != NULL)
-    to_replace->children[1]->parent = to_replace;
+  vieter_tree_node_set_children(to_replace, replacement->children);
 }
 
 vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key,
diff --git a/src/tree/vieter_tree_node.h b/src/tree/vieter_tree_node.h
index 4012d1d..1ebb15a 100644
--- a/src/tree/vieter_tree_node.h
+++ b/src/tree/vieter_tree_node.h
@@ -42,6 +42,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);
 
 bool vieter_tree_node_validate(vieter_tree_node *node, uint64_t passed_black_nodes, uint64_t expected_black_nodes);