AVL tree is now fixed after node is deleted

sandbox/btree/btree.c

@@ -16,6 +16,32 @@
 #define HAS_ONE_CHILD(x) ((HAS_LEFT_CHILD(x) && !HAS_RIGHT_CHILD(x)) || (!HAS_LEFT_CHILD(x) && HAS_RIGHT_CHILD(x)))
 #define HAS_TWO_CHILDREN(x) (HAS_LEFT_CHILD(x) && HAS_RIGHT_CHILD(x))
 
+#define HEIGHT(x) ((x) ? (x)->b_height : 0)
+
+static inline void update_height(btree_node_t *x)
+{
+	x->b_height = MAX(HEIGHT(x->b_left), HEIGHT((x->b_right))) + 1;
+}
+
+static inline int bf(btree_node_t *x)
+{
+	int bf = 0;
+
+	if (!x) {
+		return bf;
+	}
+
+	if (x->b_right) {
+		bf += x->b_right->b_height;
+	}
+
+	if (x->b_left) {
+		bf -= x->b_left->b_height;
+	}
+
+	return bf;
+}
+
 static void rotate_left(btree_t *tree, btree_node_t *x)
 {
 	btree_node_t *y = x->b_right;
@@ -39,14 +65,11 @@ static void rotate_left(btree_t *tree, btree_node_t *x)
 	y->b_left = x;
 	y->b_parent = p;
 
-	x->b_bfactor = x->b_bfactor - 1 - MAX(y->b_bfactor, 0);
-
-	p = y;
+	p = x;
 
 	while (p) {
-	    p->b_bfactor = p->b_bfactor - 1 + MIN(y->b_bfactor, 0);
-	    x = p;
-	    p = p->b_parent;
+		update_height(p);
+		p = p->b_parent;
 	}
 }
 
@@ -73,12 +96,11 @@ static void rotate_right(btree_t *tree, btree_node_t *y)
 	x->b_right = y;
 	x->b_parent = p;
 
-	y->b_bfactor = y->b_bfactor + 1 - MIN(x->b_bfactor, 0);
+	p = y;
 
-	p = x;
 	while (p) {
-	    p->b_bfactor = x->b_bfactor + 1 + MAX(y->b_bfactor, 0);
-	    y = p;
+		update_height(p);
+		p = p->b_parent;
 	}
 }
 
@@ -89,6 +111,16 @@ static void rotate_double_left(btree_t *tree, btree_node_t *z)
 
 	rotate_right(tree, x);
 	rotate_left(tree, z);
+
+	update_height(z);
+	update_height(x);
+	update_height(y);
+
+	btree_node_t *p = y->b_parent;
+	while (p) {
+		p->b_height--;
+		p = p->b_parent;
+	}
 }
 
 static void rotate_double_right(btree_t *tree, btree_node_t *z)
@@ -98,15 +130,25 @@ static void rotate_double_right(btree_t *tree, btree_node_t *z)
 
 	rotate_left(tree, x);
 	rotate_right(tree, z);
+
+	update_height(z);
+	update_height(x);
+	update_height(y);
+
+	btree_node_t *p = y->b_parent;
+	while (p) {
+		p->b_height--;
+		p = p->b_parent;
+	}
 }
 
-static void insert_fixup(btree_t *tree, btree_node_t *w)
+static void fix_tree(btree_t *tree, btree_node_t *w)
 {
 	btree_node_t *z = NULL, *y = NULL, *x = NULL;
 
 	z = w;
 	while (z) {
-		if (z->b_bfactor >= -1 && z->b_bfactor <= 1) {
+		if (bf(z) >= -1 && bf(z) <= 1) {
 			goto next_ancestor;
 		}
 
@@ -138,6 +180,7 @@ void btree_insert(btree_t *tree, btree_node_t *node)
 	if (!tree->b_root) {
 		tree->b_root = node;
 		node->b_parent = NULL;
+		node->b_height = 1;
 		return;
 	}
 
@@ -145,17 +188,12 @@ void btree_insert(btree_t *tree, btree_node_t *node)
 	while (1) {
 		btree_node_t **nextp = NULL;
 
-		int bfactor_diff = 0;
 		if (node->b_key >= cur->b_key) {
 			nextp = &cur->b_right;
-			bfactor_diff = 1;
 		} else {
 			nextp = &cur->b_left;
-			bfactor_diff = -1;
 		}
 
-		cur->b_bfactor = cur->b_bfactor + bfactor_diff;
-
 		if (*nextp) {
 			cur = *nextp;
 		} else {
@@ -165,13 +203,21 @@ void btree_insert(btree_t *tree, btree_node_t *node)
 		}
 	}
 
-	node->b_bfactor = 0;
-	insert_fixup(tree, node);
+	node->b_height = 0;
+
+	cur = node;
+	while (cur) {
+		cur->b_height++;
+		cur = cur->b_parent;
+	}
+
+	fix_tree(tree, node);
 }
 
-static void remove_node_with_no_children(btree_t *tree, btree_node_t *node)
+static btree_node_t *remove_node_with_no_children(btree_t *tree, btree_node_t *node)
 {
 	printf("remove_node_with_no_children()\n");
+	btree_node_t *w = node->b_parent;
 	btree_node_t *p = node->b_parent;
 	node->b_parent = NULL;
 
@@ -182,9 +228,16 @@ static void remove_node_with_no_children(btree_t *tree, btree_node_t *node)
 	} else {
 		p->b_right = NULL;
 	}
+
+	while (p) {
+		update_height(p);
+		p = p->b_parent;
+	}
+
+	return w;
 }
 
-static void replace_node_with_one_subtree(btree_t *tree, btree_node_t *node)
+static btree_node_t *replace_node_with_one_subtree(btree_t *tree, btree_node_t *node)
 {
 	printf("remove_node_with_one_subtree()\n");
 	btree_node_t *p = node->b_parent;
@@ -196,6 +249,7 @@ static void replace_node_with_one_subtree(btree_t *tree, btree_node_t *node)
 		z = node->b_right;
 	}
 
+	btree_node_t *w = z;
 	if (IS_LEFT_CHILD(p, node)) {
 		p->b_left = z;
 	} else if (IS_RIGHT_CHILD(p, node)) {
@@ -205,10 +259,17 @@ static void replace_node_with_one_subtree(btree_t *tree, btree_node_t *node)
 	z->b_parent = p;
 
 	node->b_parent = NULL;
-	node->b_left = node->b_parent = NULL;
+	node->b_left = node->b_right = NULL;
+
+	while (z) {
+		update_height(z);
+		z = z->b_parent;
+	}
+
+	return w;
 }
 
-static void replace_node_with_two_subtrees(btree_t *tree, btree_node_t *node)
+static btree_node_t *replace_node_with_two_subtrees(btree_t *tree, btree_node_t *node)
 {
 	printf("remove_node_with_two_subtrees()\n");
 	btree_node_t *cur = node->b_left;
@@ -218,13 +279,12 @@ static void replace_node_with_two_subtrees(btree_t *tree, btree_node_t *node)
 	}
 
 	btree_node_t *p = cur->b_parent;
+	btree_node_t *k = p;
 
 	if (IS_LEFT_CHILD(p, cur)) {
 		p->b_left = NULL;
-		p->b_bfactor++;
 	} else {
 		p->b_right = NULL;
-		p->b_bfactor--;
 	}
 
 	p = node->b_parent;
@@ -239,30 +299,41 @@ static void replace_node_with_two_subtrees(btree_t *tree, btree_node_t *node)
 
 	cur->b_left = node->b_left;
 	cur->b_right = node->b_right;
-	cur->b_bfactor = node->b_bfactor;
+	cur->b_height = node->b_height;
 	cur->b_parent = p;
+	btree_node_t *w = cur;
 
 	if (!p) {
 		tree->b_root = cur;
 	}
+
+	if (tree->b_root->b_left) {
+		tree->b_root->b_left->b_parent = tree->b_root;
+	}
+
+	if (tree->b_root->b_right) {
+		tree->b_root->b_right->b_parent = tree->b_root;
+	}
+
+	while (k) {
+		update_height(k);
+		k = k->b_parent;
+	}
+
+	return w;
 }
 
 void btree_delete(btree_t *tree, btree_node_t *node)
 {
-	btree_node_t *x = NULL;
-
+	btree_node_t *w = NULL;
+	
 	if (HAS_NO_CHILDREN(node)) {
-		remove_node_with_no_children(tree, node);
-		return;
+		w = remove_node_with_no_children(tree, node);
+	} else if (HAS_ONE_CHILD(node)) {
+		w = replace_node_with_one_subtree(tree, node);
+	} else if (HAS_TWO_CHILDREN(node)) {
+		w = replace_node_with_two_subtrees(tree, node);
 	}
 
-	if (HAS_ONE_CHILD(node)) {
-		replace_node_with_one_subtree(tree, node);
-		return;
-	}
-
-	if (HAS_TWO_CHILDREN(node)) {
-		replace_node_with_two_subtrees(tree, node);
-		return;
-	}
+	fix_tree(tree, w);
 }

sandbox/btree/include/socks/btree.h

@@ -7,7 +7,7 @@ typedef uint64_t btree_key_t;
 
 typedef struct btree_node {
 	struct btree_node *b_parent, *b_left, *b_right;
-	signed char b_bfactor;
+	unsigned short b_height;
 	btree_key_t b_key;
 } btree_node_t;