From cd217186a27567d5dd06194dc754f858c51d88f7 Mon Sep 17 00:00:00 2001
From: Max Wash <mbswash@gmail.com>
Date: Sat, 21 Jan 2023 21:42:54 +0000
Subject: [PATCH] sandbox: btree_test() now checks that constructed tree is a
 valid AVL tree

---
 sandbox/base/main.c   | 61 ++++++++++++++++++++++++++++++++++++++++---
 sandbox/btree/btree.c |  5 +++-
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/sandbox/base/main.c b/sandbox/base/main.c
index e1e232a..bc0b65e 100644
--- a/sandbox/base/main.c
+++ b/sandbox/base/main.c
@@ -9,7 +9,7 @@
 #include <socks/memblock.h>
 #include <socks/vm.h>
 
-#define NR_BTREE_NODES 4096
+#define NR_BTREE_NODES 16383
 
 /* we're working with 512MiB of simulated system RAM */
 #define MEMORY_SIZE_MB 512
@@ -142,6 +142,14 @@ void btree_print(btree_node_t *node, int depth)
 	}
 
 	if (node) {
+		if (node->b_parent) {
+			if (node == node->b_parent->b_left) {
+				printf("\\ ");
+			} else {
+				printf("/ ");
+			}
+		}
+
 		printf("%llu (h:%d)\n", node->b_key, node->b_height);
 	} else {
 		printf("\x1b[1;31mNULL\x1b[0m\n");
@@ -152,6 +160,50 @@ void btree_print(btree_node_t *node, int depth)
 	}
 }
 
+/* returns the height of the subtree rooted at node x, or -1 if one of these conditions is true:
+ *   - the calculated height of subtree x does not match the stored height value.
+ *   - the subtree is not a valid AVL tree.
+ */
+static int btree_avl_validate(btree_node_t *x)
+{
+	if (!x->b_left && !x->b_right) {
+		return 1;
+	}
+	
+	int left = 0, right = 0;
+
+	if (x->b_left) {
+		left = btree_avl_validate(x->b_left);
+	}
+
+	if (x->b_right) {
+		right = btree_avl_validate(x->b_right);
+	}
+
+	if (left == -1 || right == -1) {
+		return -1;
+	}
+
+	int diff = right - left;
+	if (diff > 1 || diff < -1) {
+		return -1;
+	}
+
+	int height = 0;
+
+	if (left > right) {
+		height = left + 1;
+	} else {
+		height = right + 1;
+	}
+
+	if (height != x->b_height) {
+		return -1;
+	}
+	
+	return height;
+}
+
 static int btree_test(void)
 {
 	btree_t tree = {};
@@ -168,10 +220,13 @@ static int btree_test(void)
 
 		btree_insert(&tree, &nodes[i]);
 		printf("#######################\n");
-
-		btree_print(tree.b_root, 0);
 	}
 
+	btree_print(tree.b_root, 0);
+
+	int result = btree_avl_validate(tree.b_root);
+	printf("AVL validation result: %d (%s)\n", result, result != -1 ? "pass" : "fail");
+
 	#if 0
 	int to_delete[] = { 3, 1, 0 };
 	int nr_to_delete = sizeof to_delete / sizeof to_delete[0];
diff --git a/sandbox/btree/btree.c b/sandbox/btree/btree.c
index 2e9cde8..d8014ff 100644
--- a/sandbox/btree/btree.c
+++ b/sandbox/btree/btree.c
@@ -177,6 +177,7 @@ static void rotate_double_right(btree_t *tree, btree_node_t *z)
 
 static void fix_tree(btree_t *tree, btree_node_t *w)
 {
+	int nr_rotations = 0;
 	btree_node_t *z = NULL, *y = NULL, *x = NULL;
 
 	z = w;
@@ -201,13 +202,15 @@ static void fix_tree(btree_t *tree, btree_node_t *w)
 			    update_height_to_root(z);
 			}
 		}
-
+		nr_rotations++;
 
 next_ancestor:
 		x = y;
 		y = z;
 		z = z->b_parent;
 	}
+
+	assert(nr_rotations <= 1);
 }
 
 void btree_insert(btree_t *tree, btree_node_t *node)