feat(tree): first attempt at balanced inserts

src/tree/balancing.c

@@ -1,3 +1,95 @@
 #include "balancing.h"
 
-void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {}
+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 *grand_grand_parent = grand_parent->parent;
+  vieter_tree_node *root, *left, *right;
+  vieter_tree_node *children[4];
+
+  if (node->key < parent->key && parent->key < grand_parent->key) {
+    root = parent;
+    left = node;
+    right = grand_parent;
+
+    children[0] = node->left;
+    children[1] = node->right;
+    children[2] = parent->right;
+    children[3] = grand_parent->right;
+  } else if (node->key < parent->key && parent->key > grand_parent->key) {
+    root = node;
+    left = grand_parent;
+    right = parent;
+
+    children[0] = grand_parent->left;
+    children[1] = node->left;
+    children[2] = node->right;
+    children[3] = parent->right;
+  } else if (node->key > parent->key && parent->key < grand_parent->key) {
+    root = node;
+    left = parent;
+    right = grand_parent;
+
+    children[0] = parent->left;
+    children[1] = node->left;
+    children[2] = node->right;
+    children[3] = grand_parent->right;
+  } else {
+    root = parent;
+    left = grand_parent;
+    right = node;
+
+    children[0] = grand_parent->left;
+    children[1] = parent->left;
+    children[2] = node->left;
+    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);
+
+  left->left = children[0];
+  left->right = children[1];
+
+  if (children[0] != NULL)
+    children[0]->parent = left;
+
+  if (children[1] != NULL)
+    children[1]->parent = left;
+
+  right->left = children[2];
+  right->right = children[3];
+
+  if (children[2] != NULL)
+    children[2]->parent = right;
+
+  if (children[3] != NULL)
+    children[3]->parent = right;
+
+  root->left = left;
+  root->right = right;
+  left->parent = root;
+  right->parent = root;
+
+  if (grand_grand_parent != NULL) {
+    vieter_tree_node_add_child(grand_grand_parent, root->key, root);
+  } else {
+    root->parent = NULL;
+  }
+
+  return root;
+}
+
+void vieter_tree_node_balance_after_insert(vieter_tree_node *node) {
+  while (!vieter_tree_node_get(node, vieter_node_black) &&
+         node->parent != NULL &&
+         !vieter_tree_node_get(node->parent, vieter_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);
+  }
+}

src/tree/node.c

@@ -1,4 +1,5 @@
 #include "node.h"
+#include "balancing.h"
 
 vieter_tree_node *vieter_tree_node_init() {
   return calloc(1, sizeof(vieter_tree_node));
@@ -55,6 +56,7 @@ vieter_tree_error vieter_tree_node_insert(vieter_tree_node *root, uint64_t key,
   new_node->data = data;
 
   vieter_tree_node_add_child(parent, key, new_node);
+  vieter_tree_node_balance_after_insert(new_node);
 
   return vieter_tree_ok;
 }
@@ -95,20 +97,6 @@ vieter_tree_error vieter_tree_node_search(void **out, vieter_tree_node *root,
   return vieter_tree_ok;
 }
 
-vieter_tree_node *vieter_tree_node_successor(vieter_tree_node *node) {
-  if ((node->left != NULL) ^ (node->right != NULL)) {
-    return node->left != NULL ? node->left : node->right;
-  }
-
-  node = node->right;
-
-  while (node->left != NULL) {
-    node = node->left;
-  }
-
-  return node;
-}
-
 vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root,
                                           uint64_t key) {
   vieter_tree_node *target;

src/tree/node.h

@@ -37,4 +37,7 @@ 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_add_child(vieter_tree_node *parent, uint64_t key,
+                                vieter_tree_node *child);
+
 #endif

src/tree/tree.c

@@ -16,6 +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);
     tree->size = 1;
 
     return vieter_tree_ok;