1 files changed, 65 insertions, 5 deletions

diff --git a/source/tree.c b/source/tree.c
index dcae485..964899f 100644
--- a/source/tree.c
+++ b/source/tree.c
@@ -17,14 +17,14 @@ void kpl_tree_free(kpl_tree *tree) {
 }
 
 static void *kpl_tree_find_helper(kpl_tree_node *root_node, const void *find_node,
-    kpl_tree_node_compare_fn compare_fn) {
+    kpl_tree_node_compare_fn node_compare_fn) {
     if (!root_node)
         return NULL;
-    int32_t compare = compare_fn(root_node, find_node);
+    int32_t compare = node_compare_fn(root_node, find_node);
     if (compare < 0)
-        return kpl_tree_find_helper(root_node->tree_left, find_node, compare_fn);
+        return kpl_tree_find_helper(root_node->tree_left, find_node, node_compare_fn);
     if (compare > 0)
-        return kpl_tree_find_helper(root_node->tree_right, find_node, compare_fn);
+        return kpl_tree_find_helper(root_node->tree_right, find_node, node_compare_fn);
     return root_node;
 }
 
@@ -32,6 +32,66 @@ void *kpl_tree_find(kpl_tree *tree, const void *find_node) {
     return kpl_tree_find_helper(tree->root_node, find_node, tree->node_compare_fn);
 }
 
+static int32_t kpl_tree_node_height_get(const kpl_tree_node *node) {
+    return node ? node->tree_height : 0;
+}
+
+static void kpl_tree_node_height_update(kpl_tree_node *node) {
+    int32_t left_height = kpl_tree_node_height_get(node->tree_left);
+    int32_t right_height = kpl_tree_node_height_get(node->tree_right);
+    node->tree_height = MAX(left_height, right_height) + 1;
+}
+
+static kpl_tree_node *kpl_tree_node_rotate_right(kpl_tree_node *right_node) {
+   kpl_tree_node *right_left = right_node->tree_left, *right_left_right = right_left->tree_right;
+   right_left->tree_right = right_node;
+   right_node->tree_left = right_left_right;
+   kpl_tree_node_height_update(right_node);
+   kpl_tree_node_height_update(right_left);
+   return right_left;
+}
+
+static kpl_tree_node *kpl_tree_node_rotate_left(kpl_tree_node *left_node) {
+    kpl_tree_node *left_right = left_node->tree_right, *left_right_left = left_right->tree_left;
+    left_right->tree_left = left_node;
+    left_node->tree_right = left_right_left;
+    kpl_tree_node_height_update(left_node);
+    kpl_tree_node_height_update(left_right);
+    return left_right;
+}
+
+static kpl_tree_node *kpl_tree_insert_helper(kpl_tree_node *restrict root_node,
+    kpl_tree_node *restrict insert_node, kpl_tree_node_compare_fn node_compare_fn) {
+    if (!root_node)
+        return insert_node;
+    int32_t compare = node_compare_fn(root_node, insert_node);
+    if (compare < 0)
+        root_node->tree_left = kpl_tree_insert_helper(root_node->tree_left, insert_node, node_compare_fn);
+    else if (compare > 0)
+        root_node->tree_right = kpl_tree_insert_helper(root_node->tree_right, insert_node, node_compare_fn);
+    else {
+        insert_node->tree_duplicates = root_node->tree_duplicates;
+        root_node->tree_duplicates = insert_node;
+        return root_node;
+    }
+    kpl_tree_node_height_update(root_node);
+    int32_t balance = kpl_tree_node_height_get(root_node->tree_left) -
+        kpl_tree_node_height_get(root_node->tree_right);
+    if (balance > 1 && compare < 0)
+        return kpl_tree_node_rotate_right(root_node);
+    if (balance < -1 && compare > 0)
+        return kpl_tree_node_rotate_left(root_node);
+    if (balance > 1 && compare > 0) {
+        root_node->tree_left = kpl_tree_node_rotate_left(root_node->tree_left);
+        return kpl_tree_node_rotate_right(root_node);
+    }
+    if (balance < -1 && compare < 0) {
+        root_node->tree_right = kpl_tree_node_rotate_right(root_node->tree_right);
+        return kpl_tree_node_rotate_left(root_node);
+    }
+    return root_node;
+}
+
 void kpl_tree_insert(kpl_tree *tree, void *insert_node) {
-    // TODO
+    tree->root_node = kpl_tree_insert_helper(tree->root_node, insert_node, tree->node_compare_fn);
 }