sandbox: AVL tree deletion bugfixes and testing

sandbox/btree/btree.c

@@ -56,7 +56,7 @@ 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);
 
@@ -109,7 +109,7 @@ static void rotate_right(btree_t *tree, btree_node_t *y)
 	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;
@@ -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,9 +263,10 @@ 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;
 		}
-		
+
 		if (node->b_key >= cur->b_key) {
 			nextp = &cur->b_right;
 		} else {
@@ -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,8 +410,9 @@ 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)) {
 		w = remove_node_with_no_children(tree, node);
 	} else if (HAS_ONE_CHILD(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;