tree impl

6 files changed, 92 insertions, 9 deletions

diff --git a/include/kpl/base.h b/include/kpl/base.h
index 144fd51..9420a44 100644
--- a/include/kpl/base.h
+++ b/include/kpl/base.h
@@ -11,6 +11,7 @@
 #include <stdbool.h>
 #include <string.h>
 #include <ctype.h>
+#include <limits.h>
 #include <pthread.h>
 #include <semaphore.h>
 #include <fcntl.h>
@@ -19,6 +20,7 @@
 #include <sys/ioctl.h>
 #include <sys/mman.h>
 #include <sys/uio.h>
+#include <sys/param.h>
 #include <linux/fs.h>
 #include <linux/io_uring.h>
 // KPL includes
diff --git a/include/kpl/tree.h b/include/kpl/tree.h
index 8e52524..90de3a8 100644
--- a/include/kpl/tree.h
+++ b/include/kpl/tree.h
@@ -4,9 +4,9 @@
 
 #pragma once
 
-// tree_duplicate is a linked stack
+// tree_duplicates is a linked stack
 #define KPL_TREE_HEADER(STRUCT_NAME) \
-    struct STRUCT_NAME *tree_left, *tree_right, *tree_duplicate; \
+    struct STRUCT_NAME *tree_left, *tree_right, *tree_duplicates; \
     int32_t tree_height
 
 /*
diff --git a/include/kpl_test.h b/include/kpl_test.h
index a656761..70a99b2 100644
--- a/include/kpl_test.h
+++ b/include/kpl_test.h
@@ -38,7 +38,7 @@ void _kpl_test_add(kpl_test *_test); // do not call outside KPL_TEST macro
 #define KPL_TEST(NAME) _KPL_TEST_FN_DEFINITION(NAME); \
 static void __attribute__((constructor)) _ ## NAME(void) { \
     static kpl_test _test = { \
-        .tree_left = NULL, .tree_right = NULL, .tree_duplicate = NULL, \
+        .tree_left = NULL, .tree_right = NULL, .tree_duplicates = NULL, \
         .tree_height = KPL_TREE_TREE_HEIGHT_INIT, \
         .error_line_number = KPL_TEST_ERROR_LINE_NUMBER_INIT, \
         .list_next = NULL, \
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);
 }
diff --git a/test/test.c b/test/test.c
index 6cede92..02a8730 100644
--- a/test/test.c
+++ b/test/test.c
@@ -67,7 +67,22 @@ static int32_t kpl_test_run_all(void) {
 int32_t main(int32_t argc, char *argv[]) {
     if (argc < 2)
         return kpl_test_run_all();
-    // TODO run by name
+    kpl_test find = {};
+    for (int32_t arg_index = 1; arg_index < argc; arg_index++) {
+        find.name = argv[arg_index];
+        kpl_test *found = kpl_tree_find(&test_tree, &find);
+        if (!found) {
+            printf(KPL_COLOR2(BOLD, RED) "TEST: " KPL_COLOR(RESET) KPL_COLOR2(BOLD, WHITE) "%s " KPL_COLOR(RESET)
+                KPL_COLOR2(BOLD, RED) "NOT FOUND\n" KPL_COLOR(RESET), find.name);
+            return KPL_EXIT_CODE_ERROR;
+        }
+        kpl_test *duplicates = found->tree_duplicates;
+        while (duplicates) {
+            kpl_test_run_one(duplicates);
+            duplicates = duplicates->tree_duplicates;
+        }
+        kpl_test_run_one(found);
+    }
     return kpl_test_report_error();
 }
 
diff --git a/test/tree.c b/test/tree.c
new file mode 100644
index 0000000..caaa23a
--- /dev/null
+++ b/test/tree.c
@@ -0,0 +1,6 @@
+
+#include "kpl_test.h"
+
+KPL_TEST(c_string_trees) {
+    // TODO
+}