sandbox: added some AVL tree insertion stress testing

sandbox/base/main.c

@@ -9,7 +9,7 @@
 #include <socks/memblock.h>
 #include <socks/vm.h>
 
-#define NR_BTREE_NODES 3
+#define NR_BTREE_NODES 4096
 
 /* we're working with 512MiB of simulated system RAM */
 #define MEMORY_SIZE_MB 512
@@ -127,7 +127,7 @@ static int memory_test(void)
 	return 0;
 }
 
-static void btree_print(btree_node_t *node, int depth)
+void btree_print(btree_node_t *node, int depth)
 {
 	if (!node) {
 		return;
@@ -155,22 +155,24 @@ static void btree_print(btree_node_t *node, int depth)
 static int btree_test(void)
 {
 	btree_t tree = {};
-	btree_node_t nodes[] = {
+	btree_node_t *nodes = calloc(NR_BTREE_NODES, sizeof *nodes);
-		{ .b_key = 1 },
-		{ .b_key = 2 },
-		{ .b_key = 3 },
-		{ .b_key = 4 },
-		{ .b_key = 5 },
-		{ .b_key = 6 },
-		{ .b_key = 7 },
-	};
 
-	int nr_nodes = sizeof nodes / sizeof nodes[0];
+	for (int i = 0; i < NR_BTREE_NODES; i++) {
-	for (int i = 0; i < nr_nodes; i++) {
+		nodes[i].b_key = (rand() % 512000) + 1;
-		btree_insert(&tree, &nodes[i]);
+		printf(" - node %d: %llu\n", i, nodes[i].b_key);
 	}
-	btree_print(tree.b_root, 0);
 
+	for (int i = 0; i < NR_BTREE_NODES; i++) {
+		printf("#######################\n");
+		printf("inserting node #%d: %llu\n", i, nodes[i].b_key);
+
+		btree_insert(&tree, &nodes[i]);
+		printf("#######################\n");
+
+		btree_print(tree.b_root, 0);
+	}
+
+	#if 0
 	int to_delete[] = { 3, 1, 0 };
 	int nr_to_delete = sizeof to_delete / sizeof to_delete[0];
 
@@ -184,7 +186,9 @@ static int btree_test(void)
 		btree_delete(&tree, &nodes[node_index]);
 		btree_print(tree.b_root, 0);
 	}
+	#endif
 
+	free(nodes);
 	return 0;
 }
 

sandbox/btree/btree.c

@@ -2,6 +2,15 @@
 #include <stddef.h>
 #include <stdlib.h>
 #include <stdio.h>
+#include <assert.h>
+
+#define BTREE_VERBOSE
+
+#ifdef BTREE_VERBOSE
+#define debug_msg(...) printf(__VA_ARGS__)
+#else
+#define debug_msg(...)
+#endif
 
 #define MAX(a, b) ((a) > (b) ? (a) : (b))
 #define MIN(a, b) ((a) < (b) ? (a) : (b))
@@ -18,6 +27,8 @@
 
 #define HEIGHT(x) ((x) ? (x)->b_height : 0)
 
+extern void btree_print(btree_node_t *, int);
+
 static inline void update_height(btree_node_t *x)
 {
 	x->b_height = MAX(HEIGHT(x->b_left), HEIGHT((x->b_right))) + 1;
@@ -44,7 +55,18 @@ static inline int bf(btree_node_t *x)
 
 static void rotate_left(btree_t *tree, btree_node_t *x)
 {
+	assert(x != NULL);
+	
 	btree_node_t *y = x->b_right;
+	assert(y != NULL);
+
+	assert(y == x->b_left || y == x->b_right);
+	if (x->b_parent) {
+		assert(x == x->b_parent->b_left || x == x->b_parent->b_right);
+	}
+
+	debug_msg("  rotate_left(x: %llu, y: %llu)\n", x->b_key, y->b_key);
+
 	btree_node_t *p = x->b_parent;
 
 	if (y->b_left) {
@@ -56,26 +78,40 @@ static void rotate_left(btree_t *tree, btree_node_t *x)
 	if (!p) {
 		tree->b_root = y;
 	} else if (x == p->b_left) {
+		debug_msg("    %llu->left was %llu, now %llu\n", p->b_key, p->b_left->b_key, y->b_key);
 		p->b_left = y;
 	} else {
+		debug_msg("    %llu->right was %llu, now %llu\n", p->b_key, p->b_right->b_key, y->b_key);
 		p->b_right = y;
 	}
 
 	x->b_parent = y;
 	y->b_left = x;
 	y->b_parent = p;
+}
 
-	p = x;
+static void update_height_to_root(btree_node_t *x)
-
+{
-	while (p) {
+	while (x) {
-		update_height(p);
+		update_height(x);
-		p = p->b_parent;
+		x = x->b_parent;
 	}
 }
 
 static void rotate_right(btree_t *tree, btree_node_t *y)
 {
+	assert(y);
+
 	btree_node_t *x = y->b_left;
+	assert(x);
+
+	assert(x == y->b_left || x == y->b_right);
+	if (y->b_parent) {
+		assert(y == y->b_parent->b_left || y == y->b_parent->b_right);
+	}
+	
+	debug_msg("  rotate_right(y: %llu, x: %llu)\n", y->b_key, x->b_key);
+
 	btree_node_t *p = y->b_parent;
 
 	if (x->b_right) {
@@ -86,40 +122,36 @@ static void rotate_right(btree_t *tree, btree_node_t *y)
 
 	if (!p) {
 		tree->b_root = x;
-	} else if (x == p->b_left) {
+	} else if (y == p->b_left) {
+		debug_msg("    %llu->left was %llu, now %llu\n", p->b_key, p->b_left->b_key, x->b_key);
 		p->b_left = x;
 	} else {
+		debug_msg("    %llu->right was %llu, now %llu\n", p->b_key, p->b_right->b_key, x->b_key);
 		p->b_right = x;
 	}
 
 	y->b_parent = x;
 	x->b_right = y;
 	x->b_parent = p;
-
-	p = y;
-
-	while (p) {
-		update_height(p);
-		p = p->b_parent;
-	}
 }
 
 static void rotate_double_left(btree_t *tree, btree_node_t *z)
 {
 	btree_node_t *x = z->b_right;
 	btree_node_t *y = x->b_left;
+	debug_msg("rotate_double_left(z: %llu, x: %llu, y: %llu) {\n", z->b_key, x->b_key, y->b_key);
 
 	rotate_right(tree, x);
 	rotate_left(tree, z);
 
+	debug_msg("}\n");
+
 	update_height(z);
 	update_height(x);
-	update_height(y);
 
-	btree_node_t *p = y->b_parent;
+	while (y) {
-	while (p) {
+		update_height(y);
-		p->b_height--;
+		y = y->b_parent;
-		p = p->b_parent;
 	}
 }
 
@@ -127,18 +159,19 @@ static void rotate_double_right(btree_t *tree, btree_node_t *z)
 {
 	btree_node_t *x = z->b_left;
 	btree_node_t *y = x->b_right;
+	debug_msg("rotate_double_right(z: %llu, x: %llu, y: %llu) {\n", z->b_key, x->b_key, y->b_key);
 
 	rotate_left(tree, x);
 	rotate_right(tree, z);
 
+	debug_msg("}\n");
+
 	update_height(z);
 	update_height(x);
-	update_height(y);
 
-	btree_node_t *p = y->b_parent;
+	while (y) {
-	while (p) {
+		update_height(y);
-		p->b_height--;
+		y = y->b_parent;
-		p = p->b_parent;
 	}
 }
 
@@ -156,6 +189,7 @@ static void fix_tree(btree_t *tree, btree_node_t *w)
 		if (IS_LEFT_CHILD(z, y)) {
 			if (IS_LEFT_CHILD(y, x)) {
 			    rotate_right(tree, z);
+			    update_height_to_root(z);
 			} else {
 			    rotate_double_right(tree, z);
 			}
@@ -164,6 +198,7 @@ static void fix_tree(btree_t *tree, btree_node_t *w)
 			    rotate_double_left(tree, z);
 			} else {
 			    rotate_left(tree, z);
+			    update_height_to_root(z);
 			}
 		}
 
@@ -188,6 +223,10 @@ void btree_insert(btree_t *tree, btree_node_t *node)
 	while (1) {
 		btree_node_t **nextp = NULL;
 
+		if (node->b_key == cur->b_key) {
+			return;
+		}
+		
 		if (node->b_key >= cur->b_key) {
 			nextp = &cur->b_right;
 		} else {
@@ -207,16 +246,21 @@ void btree_insert(btree_t *tree, btree_node_t *node)
 
 	cur = node;
 	while (cur) {
-		cur->b_height++;
+		update_height(cur);
 		cur = cur->b_parent;
 	}
 
+	#if 0
+	printf("##### TREE POST-INSERT, PRE-FIXUP ##############\n");
+	btree_print(tree->b_root, 0);
+	printf("################################################\n");
+	#endif
 	fix_tree(tree, node);
 }
 
 static btree_node_t *remove_node_with_no_children(btree_t *tree, btree_node_t *node)
 {
-	printf("remove_node_with_no_children()\n");
+	debug_msg("remove_node_with_no_children()\n");
 	btree_node_t *w = node->b_parent;
 	btree_node_t *p = node->b_parent;
 	node->b_parent = NULL;
@@ -239,7 +283,7 @@ static btree_node_t *remove_node_with_no_children(btree_t *tree, btree_node_t *n
 
 static btree_node_t *replace_node_with_one_subtree(btree_t *tree, btree_node_t *node)
 {
-	printf("remove_node_with_one_subtree()\n");
+	debug_msg("remove_node_with_one_subtree()\n");
 	btree_node_t *p = node->b_parent;
 	btree_node_t *z = NULL;
 
@@ -271,7 +315,7 @@ static btree_node_t *replace_node_with_one_subtree(btree_t *tree, btree_node_t *
 
 static btree_node_t *replace_node_with_two_subtrees(btree_t *tree, btree_node_t *node)
 {
-	printf("remove_node_with_two_subtrees()\n");
+	debug_msg("remove_node_with_two_subtrees()\n");
 	btree_node_t *cur = node->b_left;
 
 	while (cur->b_right) {