From b2944f4cb8a2ec138270a2cc7f819d002f69fa68 Mon Sep 17 00:00:00 2001
From: Brendan Hansen <brendan.f.hansen@gmail.com>
Date: Tue, 10 May 2022 19:53:45 -0500
Subject: [PATCH] added avl_tree implementation; missing delete()

---
 core/container/avl_tree.onyx | 95 +++++++++++++++++++++++++++++++++++-
 core/std.onyx                |  1 +
 tests/avl_test               |  4 ++
 tests/avl_test.onyx          | 22 +++++++++
 4 files changed, 121 insertions(+), 1 deletion(-)
 create mode 100644 tests/avl_test
 create mode 100644 tests/avl_test.onyx

diff --git a/core/container/avl_tree.onyx b/core/container/avl_tree.onyx
index 73d35f9d..554ea44e 100644
--- a/core/container/avl_tree.onyx
+++ b/core/container/avl_tree.onyx
@@ -1,14 +1,107 @@
 package core.avl_tree
 
+#local math :: package core.math
+
 AVL_Tree :: struct (T: type_expr) {
     data: T;
     left, right: ^AVL_Tree(T);
+    height: i32;
+}
 
-    insert :: (tree: ^AVL_Tree, data: tree.T) {
+insert :: (ptree: ^^AVL_Tree($T), data: T) {
+    tree := *ptree;
+    if tree == null {
         node := new(typeof *tree);
         node.data = data;
         node.left = null;
         node.right = null;
+        node.height = 0;
+        *ptree = node;
+        return;
+    }
+
+    if data < tree.data {
+        insert(^tree.left, data);
+    } else {
+        insert(^tree.right, data);
+    }
+
+    tree.height = math.max(get_height(tree.left), get_height(tree.right)) + 1;
+
+    bf := get_height(tree.left) - get_height(tree.right);
+    if bf < -1 {
+        child_bf := get_height(tree.right.left) - get_height(tree.right.right);
+        if child_bf < 0 {
+            rotate_left(ptree);
+        } else {
+            rotate_right(^tree.right);
+            rotate_left(ptree);
+        }
+
+    } elseif bf > 1 {
+        child_bf := get_height(tree.left.left) - get_height(tree.left.right);
+        if child_bf < 0 {
+            rotate_right(ptree);
+        } else {
+            rotate_left(^tree.left);
+            rotate_right(ptree);
+        }
     }
 }
 
+delete :: (tree: ^AVL_Tree, data: tree.T) {
+
+}
+
+contains :: (tree: ^AVL_Tree, data: tree.T) -> bool {
+    if tree == null do return false;
+
+    if tree.data == data do return true;
+
+    if data < tree.data do return contains(tree.left, data);
+    else                do return contains(tree.right, data);
+}
+
+print :: (tree: ^AVL_Tree) {
+    use package core {print as std_print, printf}
+
+    if tree == null {
+        std_print("_ ");
+        return;
+    }
+
+    printf("{}[{}] ", tree.data, tree.height);
+    if tree.left != null || tree.right != null {
+        std_print("( ");
+        print(tree.left);
+        print(tree.right);
+        std_print(") ");
+    }
+}
+
+#local get_height :: (tree: ^AVL_Tree) -> i32 {
+    if tree == null do return -1;
+    return tree.height;
+}
+
+#local rotate_left :: (tree: ^^AVL_Tree) {
+    A := *tree;
+    B := A.right;
+    A.right = B.left;
+    B.left = A;
+    *tree = B;
+
+    A.height = math.max(get_height(A.left), get_height(A.right)) + 1;
+    B.height = math.max(get_height(B.left), get_height(B.right)) + 1;
+}
+
+#local rotate_right :: (tree: ^^AVL_Tree) {
+    A := *tree;
+    B := A.left;
+    A.left = B.right;
+    B.right = A;
+    *tree = B;
+
+    A.height = math.max(get_height(A.left), get_height(A.right)) + 1;
+    B.height = math.max(get_height(B.left), get_height(B.right)) + 1;
+}
\ No newline at end of file
diff --git a/core/std.onyx b/core/std.onyx
index 28a22830..8c888b77 100644
--- a/core/std.onyx
+++ b/core/std.onyx
@@ -5,6 +5,7 @@ package core
 #load "./memory"
 
 #load "./container/array"
+#load "./container/avl_tree"
 #load "./container/map"
 #load "./container/list"
 #load "./container/iter"
diff --git a/tests/avl_test b/tests/avl_test
new file mode 100644
index 00000000..4a0e271d
--- /dev/null
+++ b/tests/avl_test
@@ -0,0 +1,4 @@
+8[3] ( 5[1] ( 0[0] _ ) 20[2] ( 15[1] ( 10[0] _ ) 25[0] ) ) 
+Contains 7? false
+Contains 15? true
+Contains 0? true
diff --git a/tests/avl_test.onyx b/tests/avl_test.onyx
new file mode 100644
index 00000000..a727e494
--- /dev/null
+++ b/tests/avl_test.onyx
@@ -0,0 +1,22 @@
+#load "core/std"
+
+use package core
+
+main :: () {
+    tree: ^avl_tree.AVL_Tree(i32);
+
+    avl_tree.insert(^tree, 20);
+    avl_tree.insert(^tree, 5);
+    avl_tree.insert(^tree, 8);
+    avl_tree.insert(^tree, 25);
+    avl_tree.insert(^tree, 0);
+    avl_tree.insert(^tree, 15);
+    avl_tree.insert(^tree, 10);
+
+    avl_tree.print(tree);
+    print("\n");
+
+    printf("Contains {}? {}\n", 7, avl_tree.contains(tree, 7));
+    printf("Contains {}? {}\n", 15, avl_tree.contains(tree, 15));
+    printf("Contains {}? {}\n", 0, avl_tree.contains(tree, 0));
+}
\ No newline at end of file
-- 
2.25.1