sandbox: AVL tree deletion bugfixes and testing

sandbox/base/main.c

@@ -3,13 +3,14 @@
 #include <stdlib.h>
 #include <inttypes.h>
 #include <time.h>
+#include <assert.h>
 #include <sys/mman.h>
 #include <socks/types.h>
 #include <socks/btree.h>
 #include <socks/memblock.h>
 #include <socks/vm.h>
 
-#define NR_BTREE_NODES 16383
+#define NR_BTREE_NODES 13
 
 /* we're working with 512MiB of simulated system RAM */
 #define MEMORY_SIZE_MB 512
@@ -166,6 +167,10 @@ void btree_print(btree_node_t *node, int depth)
  */
 static int btree_avl_validate(btree_node_t *x)
 {
+	if (!x) {
+		return 0;
+	}
+
 	if (!x->b_left && !x->b_right) {
 		return 1;
 	}
@@ -210,38 +215,42 @@ static int btree_test(void)
 	btree_node_t *nodes = calloc(NR_BTREE_NODES, sizeof *nodes);
 
 	for (int i = 0; i < NR_BTREE_NODES; i++) {
-		nodes[i].b_key = (rand() % 512000) + 1;
+		nodes[i].b_key = (rand() % 128) + 1;
 		printf(" - node %d: %llu\n", i, nodes[i].b_key);
 	}
 
+	int validation_result = 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");
+
+		validation_result = btree_avl_validate(tree.b_root);
+		assert(validation_result >= 1);
 	}
 
 	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];
-
-	for (int i = 0; i < nr_to_delete; i++) {
-		int node_index = to_delete[i];
+	printf("AVL tree height: %d\n", result);
 
+	for (int i = 0; i < NR_BTREE_NODES; i++) {
 		printf("#######################\n");
-		printf("deleting node %llu\n", nodes[node_index].b_key);
+		printf("deleting node #%d: %llu\n", i, nodes[i].b_key);
 		printf("#######################\n");
 
-		btree_delete(&tree, &nodes[node_index]);
+		if (nodes[i].b_key == 89) {
+			printf("brk\n");
+		}
+
+		btree_delete(&tree, &nodes[i]);
 		btree_print(tree.b_root, 0);
+
+		validation_result = btree_avl_validate(tree.b_root);
+		assert(validation_result >= 0);
 	}
-	#endif
 
 	free(nodes);
 	return 0;

sandbox/btree/btree.c

@@ -175,7 +175,7 @@ static void rotate_double_right(btree_t *tree, btree_node_t *z)
 	}
 }
 
-static void fix_tree(btree_t *tree, btree_node_t *w)
+static void insert_fixup(btree_t *tree, btree_node_t *w)
 {
 	int nr_rotations = 0;
 	btree_node_t *z = NULL, *y = NULL, *x = NULL;
@@ -186,6 +186,9 @@ static void fix_tree(btree_t *tree, btree_node_t *w)
 			goto next_ancestor;
 		}
 
+		assert(x && y && z);
+		assert(x == y->b_left || x == y->b_right);
+		assert(y == z->b_left || y == z->b_right);
 
 		if (IS_LEFT_CHILD(z, y)) {
 			if (IS_LEFT_CHILD(y, x)) {
@@ -209,12 +212,45 @@ next_ancestor:
 		y = z;
 		z = z->b_parent;
 	}
+}
 
-	assert(nr_rotations <= 1);
+static void delete_fixup(btree_t *tree, btree_node_t *w)
+{
+	printf("delete_fixup(%llu)\n", w->b_key);
+	btree_node_t *z = w;
+
+	int nr_rotations = 0;
+
+	while (z) {
+		if (bf(z) > 1) {
+			if (bf(z->b_left) <= 0) {
+				rotate_left(tree, z);
+				update_height_to_root(z);
+			} else {
+				rotate_double_left(tree, z);
+			}
+		} else if (bf(z) < -1) {
+			if (bf(z->b_right) > 0) {
+				rotate_right(tree, z);
+				update_height_to_root(z);
+			} else {
+				rotate_double_right(tree, z);
+			}
+		}
+
+		z = z->b_parent;
+		nr_rotations++;
+
+		if (nr_rotations > 10) {
+			assert(0 && "too many rotations");
+		}
+	}
 }
 
 void btree_insert(btree_t *tree, btree_node_t *node)
 {
+	node->b_owner = tree;
+
 	if (!tree->b_root) {
 		tree->b_root = node;
 		node->b_parent = NULL;
@@ -227,6 +263,7 @@ void btree_insert(btree_t *tree, btree_node_t *node)
 		btree_node_t **nextp = NULL;
 
 		if (node->b_key == cur->b_key) {
+			node->b_owner = NULL;
 			return;
 		}
 
@@ -258,12 +295,13 @@ void btree_insert(btree_t *tree, btree_node_t *node)
 	btree_print(tree->b_root, 0);
 	printf("################################################\n");
 	#endif
-	fix_tree(tree, node);
+
+	insert_fixup(tree, node);
 }
 
 static btree_node_t *remove_node_with_no_children(btree_t *tree, btree_node_t *node)
 {
-	debug_msg("remove_node_with_no_children()\n");
+	debug_msg("remove_node_with_no_children(%llu)\n", node->b_key);
 	btree_node_t *w = node->b_parent;
 	btree_node_t *p = node->b_parent;
 	node->b_parent = NULL;
@@ -362,9 +400,9 @@ static btree_node_t *replace_node_with_two_subtrees(btree_t *tree, btree_node_t
 		tree->b_root->b_right->b_parent = tree->b_root;
 	}
 
-	while (k) {
-		update_height(k);
-		k = k->b_parent;
+	while (cur) {
+		update_height(cur);
+		cur = cur->b_parent;
 	}
 
 	return w;
@@ -372,6 +410,7 @@ static btree_node_t *replace_node_with_two_subtrees(btree_t *tree, btree_node_t
 
 void btree_delete(btree_t *tree, btree_node_t *node)
 {
+	assert(node->b_owner == tree);
 	btree_node_t *w = NULL;
 
 	if (HAS_NO_CHILDREN(node)) {
@@ -382,5 +421,9 @@ void btree_delete(btree_t *tree, btree_node_t *node)
 		w = replace_node_with_two_subtrees(tree, node);
 	}
 
-	fix_tree(tree, w);
+	if (w) {
+		delete_fixup(tree, w);
+	}
+
+	node->b_left = node->b_right = node->b_parent = NULL;
 }

sandbox/btree/include/socks/btree.h

@@ -5,13 +5,16 @@
 
 typedef uint64_t btree_key_t;
 
+struct btree;
+
 typedef struct btree_node {
 	struct btree_node *b_parent, *b_left, *b_right;
 	unsigned short b_height;
 	btree_key_t b_key;
+	struct btree *b_owner;
 } btree_node_t;
 
-typedef struct btree_t {
+typedef struct btree {
 	struct btree_node *b_root;
 } btree_t;