mango/sandbox/btree/btree.c

#include <socks/btree.h>
#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))

#define IS_LEFT_CHILD(p, c) ((p) && (c) && ((p)->b_left == (c)))
#define IS_RIGHT_CHILD(p, c) ((p) && (c) && ((p)->b_right == (c)))

#define HAS_LEFT_CHILD(x) ((x) && ((x)->b_left))
#define HAS_RIGHT_CHILD(x) ((x) && ((x)->b_right))

#define HAS_NO_CHILDREN(x) ((x) && (!(x)->b_left) && (!(x)->b_right))
#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)

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;
}

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)
{
	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) {
		y->b_left->b_parent = x;
	}

	x->b_right = y->b_left;

	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;
}

static void update_height_to_root(btree_node_t *x)
{
	while (x) {
		update_height(x);
		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) {
		x->b_right->b_parent = y;
	}

	y->b_left = x->b_right;

	if (!p) {
		tree->b_root = x;
	} 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;
}

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);

	while (y) {
		update_height(y);
		y = y->b_parent;
	}
}

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);

	while (y) {
		update_height(y);
		y = y->b_parent;
	}
}

static void insert_fixup(btree_t *tree, btree_node_t *w)
{
	int nr_rotations = 0;
	btree_node_t *z = NULL, *y = NULL, *x = NULL;

	z = w;
	while (z) {
		if (bf(z) >= -1 && bf(z) <= 1) {
			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)) {
			    rotate_right(tree, z);
			    update_height_to_root(z);
			} else {
			    rotate_double_right(tree, z);
			}
		} else {
			if (IS_LEFT_CHILD(y, x)) {
			    rotate_double_left(tree, z);
			} else {
			    rotate_left(tree, z);
			    update_height_to_root(z);
			}
		}
		nr_rotations++;

next_ancestor:
		x = y;
		y = z;
		z = z->b_parent;
	}
}

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) {
		printf("bf(z)=%d\n", bf(z));
		printf("bf(z->b_left)=%d\n", bf(z->b_left));

		if (bf(z) > 1) {
			if (bf(z->b_right) >= 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_left) <= 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;
		node->b_height = 1;
		return;
	}

	btree_node_t *cur = tree->b_root;
	while (1) {
		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 {
			nextp = &cur->b_left;
		}

		if (*nextp) {
			cur = *nextp;
		} else {
			node->b_parent = cur;
			*nextp = node;
			break;
		}
	}

	node->b_height = 0;

	cur = node;
	while (cur) {
		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

	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(%llu)\n", node->b_key);
	btree_node_t *w = node->b_parent;
	btree_node_t *p = node->b_parent;
	node->b_parent = NULL;

	if (!p) {
		tree->b_root = NULL;
	} else if (IS_LEFT_CHILD(p, node)) {
		p->b_left = NULL;
	} else {
		p->b_right = NULL;
	}

	while (p) {
		update_height(p);
		p = p->b_parent;
	}

	return w;
}

static btree_node_t *replace_node_with_one_subtree(btree_t *tree, btree_node_t *node)
{
	debug_msg("remove_node_with_one_subtree()\n");
	btree_node_t *p = node->b_parent;
	btree_node_t *z = NULL;

	if (HAS_LEFT_CHILD(node)) {
		z = node->b_left;
	} else {
		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)) {
		p->b_right = z;
	}

	z->b_parent = p;

	node->b_parent = NULL;
	node->b_left = node->b_right = NULL;

	while (z) {
		update_height(z);
		z = z->b_parent;
	}

	return w;
}

static btree_node_t *replace_node_with_two_subtrees(btree_t *tree, btree_node_t *z)
{
	debug_msg("remove_node_with_two_subtrees()\n");

	/* x will replace z */
	btree_node_t *x = z->b_left;

	while (x->b_right) {
		x = x->b_right;
	}

	/* y is the node that will replace x (if x has a left child) */
	btree_node_t *y = x->b_left;

	/* w is the starting point for the height update and fixup */
	btree_node_t *w = x;
	if (w->b_parent != z) {
		w = w->b_parent;
	}

	if (y) {
		w = y;
	}

	if (IS_LEFT_CHILD(x->b_parent, x)) {
		x->b_parent->b_left = y;
	} else if (IS_RIGHT_CHILD(x->b_parent, x)) {
		x->b_parent->b_right = y;
	}

	if (y) {
	    y->b_parent = x->b_parent;
	}

	if (IS_LEFT_CHILD(z->b_parent, z)) {
		z->b_parent->b_left = x;
	} else if (IS_RIGHT_CHILD(z->b_parent, z)) {
		z->b_parent->b_right = x;
	}

	x->b_parent = z->b_parent;
	x->b_left = z->b_left;
	x->b_right = z->b_right;

	if (x->b_left) {
		x->b_left->b_parent = x;
	}

	if (x->b_right) {
		x->b_right->b_parent = x;
	}

	if (!x->b_parent) {
		tree->b_root = x;
	}

	btree_node_t *cur = w;
	while (cur) {
		update_height(cur);
		cur = cur->b_parent;
	}

	printf("w=%llu, h:%u\n", w->b_key, w->b_height);
	return w;
}

void btree_delete(btree_t *tree, btree_node_t *node)
{
	if (!node->b_owner) {
		return;
	}

	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)) {
		w = replace_node_with_one_subtree(tree, node);
	} else if (HAS_TWO_CHILDREN(node)) {
		w = replace_node_with_two_subtrees(tree, node);
	}

	btree_print(tree->b_root, 0);

	if (w) {
		delete_fixup(tree, w);
	}

	node->b_left = node->b_right = node->b_parent = NULL;
}