feat(tree): first attempt at balanced inserts

Makefile

@@ -27,7 +27,7 @@ INC_FLAGS := $(addprefix -I,$(INC_DIRS))
 #  object file is also recompiled if only a header is changed.
 # -MP: generate a dummy target for every header file (according to the  docs it
 #  prevents some errors when removing header files)
-CFLAGS ?= -MMD -MP -Wall -Wextra
+CFLAGS ?= -MMD -MP -Wall -Wextra -g
 CFLAGS += $(INC_FLAGS)
 
 .PHONY: all

src/tree/balancing.c

@@ -0,0 +1,95 @@
+#include "balancing.h"
+
+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/balancing.h

@@ -0,0 +1,8 @@
+#ifndef VIETER_TREE_BALANCING
+#define VIETER_TREE_BALANCING
+
+#include "node.h"
+
+void vieter_tree_node_balance_after_insert(vieter_tree_node *node);
+
+#endif

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;
@@ -158,3 +146,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) {
+  node->flags |= flag;
+}
+
+void vieter_tree_node_unset(vieter_tree_node *node, vieter_node_flag flag) {
+  node->flags &= ~flag;
+}
+
+bool vieter_tree_node_get(vieter_tree_node *node, vieter_node_flag flag) {
+  return (node->flags & flag) != 0;
+}

src/tree/node.h

@@ -3,6 +3,13 @@
 
 #include "vieter_tree.h"
 #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 struct vieter_tree_node {
     uint64_t key;
@@ -10,6 +17,7 @@ typedef struct vieter_tree_node {
     struct vieter_tree_node *parent;
     struct vieter_tree_node *left;
     struct vieter_tree_node *right;
+    uint8_t flags;
 } vieter_tree_node;
 
 vieter_tree_node *vieter_tree_node_init();
@@ -25,4 +33,11 @@ vieter_tree_error vieter_tree_node_remove(void **out, vieter_tree_node *root, ui
 
 void vieter_tree_node_replace(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_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;