treap.c

/*
 * Copyright (c) 2014 by Alex Reece <[email protected]>. All rights reserved.
 */

#include <stdbool.h>

#include "debug.h"
#include "treap.h"
#include "util.h"

static void *
n_from_tn(treap_t *tree, treap_node_t *tnode)
{
	uintptr_t tn_uint = (uintptr_t) tnode;
	uintptr_t n_uint = tn_uint - tree->t_offset;
	ASSERT3U(tn_uint, >, tree->t_offset);
	return ((void *) n_uint);
}

static treap_node_t *
tn_from_n(treap_t *tree, void *node)
{
	uintptr_t n_uint = (uintptr_t) node;
	uintptr_t tn_uint = n_uint + tree->t_offset;
	return ((treap_node_t *) tn_uint);
}

static void
treap_node_invalidate(treap_node_t *tnode)
{
	tnode->tn_priority = -1;
	tnode->tn_left = (treap_node_t *)-1;
	tnode->tn_right = (treap_node_t *)-1;
}

static bool UNUSED
check_node_invariants(UNUSED treap_t *tree, treap_node_t *ctn)
{
	if (ctn->tn_right != NULL) {
		ASSERT3S(ctn->tn_priority, >, ctn->tn_right->tn_priority);
		ASSERT3S(tree->t_compare(n_from_tn(tree, ctn),
		    n_from_tn(tree, ctn->tn_right)), <, 0);
	}
	if (ctn->tn_left != NULL) {
		ASSERT3S(ctn->tn_priority, >, ctn->tn_left->tn_priority);
		ASSERT3S(tree->t_compare(n_from_tn(tree, ctn->tn_left),
		    n_from_tn(tree, ctn)), <, 0);
	}
	return (true);
}

void
treap_init(treap_t *tree, int (*compare)(const void *, const void *),
    size_t offset)
{
	/* Assert link is word aligned. */
	ASSERT3U(offset & 7, ==, 0);
	ASSERT3P(tree, !=, NULL);
	ASSERT3P(compare, !=, NULL);

	tree->t_compare = compare;
	tree->t_offset = offset;
	tree->t_root = NULL;
}

treap_node_t *
_treap_find(treap_t *tree, void *node, treap_node_t *ctn)
{
	int compare_result;

	if (ctn == NULL) {
		return (NULL);
	}

	ASSERT(check_node_invariants(tree, ctn));

	compare_result = tree->t_compare(node, n_from_tn(tree, ctn));

	if (compare_result == 0) {
		return (ctn);
	} else if (compare_result < 0) {
		return (_treap_find(tree, node, ctn->tn_left));
	} else {
		return (_treap_find(tree, node, ctn->tn_right));
	}
}

void *
treap_find(treap_t *tree, void *node)
{
	treap_node_t *tn;
	ASSERT3P(tree->t_compare, !=, NULL);
	tn = _treap_find(tree, node, tree->t_root);
	if (tn == NULL) {
		return (NULL);
	} else {
		return (n_from_tn(tree, tn));
	}
}

treap_node_t *
rotate_left(treap_node_t *ctn)
{
	treap_node_t *new_child = ctn;
	treap_node_t *new_parent = new_child->tn_right;
	treap_node_t *temp = new_parent->tn_left;

	ASSERT3S(new_child->tn_priority, <, new_parent->tn_priority);

	new_parent->tn_left = new_child;
	new_child->tn_right = temp;

	return (new_parent);
}

treap_node_t *
rotate_right(treap_node_t *ctn)
{
	treap_node_t *new_child = ctn;
	treap_node_t *new_parent = new_child->tn_left;
	treap_node_t *temp = new_parent->tn_right;

	ASSERT3S(new_child->tn_priority, <, new_parent->tn_priority);

	new_parent->tn_right = new_child;
	new_child->tn_left = temp;

	return (new_parent);
}

treap_node_t *
_treap_add(treap_t *tree, void *node, treap_node_t *ctn)
{
	int compare_result;

	if (ctn == NULL) {
		treap_node_t *tn = tn_from_n(tree, node);
		tn->tn_left = NULL;
		tn->tn_right = NULL;
		tn->tn_priority = rand();
		ASSERT3S(tn->tn_priority, >=, 0);
		return (tn);
	}

	ASSERT(check_node_invariants(tree, ctn));

	compare_result = tree->t_compare(node, n_from_tn(tree, ctn));
	ASSERT3S(compare_result, !=, 0);

	if (compare_result < 0) {
		ctn->tn_left = _treap_add(tree, node, ctn->tn_left);
		if (ctn->tn_left->tn_priority > ctn->tn_priority) {
			ctn = rotate_right(ctn);
		}
	} else {
		ctn->tn_right = _treap_add(tree, node, ctn->tn_right);
		if (ctn->tn_right->tn_priority > ctn->tn_priority) {
			ctn = rotate_left(ctn);
		}
	}

	ASSERT(check_node_invariants(tree, ctn));
	return (ctn);
}

void
treap_add(treap_t *tree, void *node)
{
	ASSERT3P(tree->t_compare, !=, NULL);
	treap_node_invalidate(tn_from_n(tree, node));
	tree->t_root = _treap_add(tree, node, tree->t_root);
	ASSERT3P(tree->t_root, !=, NULL);
}

static treap_node_t *
_treap_remove_node(treap_t *tree, treap_node_t *ctn)
{
	ASSERT3S(ctn->tn_priority, ==, 0);

	if (ctn->tn_left == NULL && ctn->tn_right == NULL) {
		treap_node_invalidate(ctn);
		return (NULL);
	}

	/*
	 * If there is no right node OR the left node is bigger than the right
	 * node, rotate the left node over the current node.
	 */
	if (ctn->tn_right == NULL || (ctn->tn_left != NULL &&
	    ctn->tn_left->tn_priority > ctn->tn_right->tn_priority)) {
		ctn = rotate_right(ctn);
		ctn->tn_right = _treap_remove_node(tree, ctn->tn_right);
	} else {
		ctn = rotate_left(ctn);
		ctn->tn_left = _treap_remove_node(tree, ctn->tn_left);
	}

	if (ctn != NULL) {
		ASSERT(check_node_invariants(tree, ctn));
	}
	return (ctn);
}

static treap_node_t *
_treap_remove_find(treap_t *tree, void *node, treap_node_t *ctn)
{
	int compare_result;

	ASSERT3P(ctn, !=, NULL);
	ASSERT(check_node_invariants(tree, ctn));

	compare_result = tree->t_compare(node, n_from_tn(tree, ctn));
	if (compare_result == 0) {
		ctn->tn_priority = 0;
		ctn = _treap_remove_node(tree, ctn);
	} else if (compare_result < 0) {
		ctn->tn_left = _treap_remove_find(tree, node, ctn->tn_left);
	} else {
		ctn->tn_right = _treap_remove_find(tree, node, ctn->tn_right);
	}

	if (ctn != NULL) {
		ASSERT(check_node_invariants(tree, ctn));
	}
	return (ctn);
}

void
treap_remove(treap_t *tree, void *node)
{
	ASSERT3P(tree->t_compare, !=, NULL);
	ASSERT3P(tree->t_root, !=, NULL);
	ASSERT(check_node_invariants(tree, tree->t_root));
	tree->t_root = _treap_remove_find(tree, node, tree->t_root);
}

void
_treap_walk(treap_t *tree, void (*walkfunc)(void *arg, void *node), void *arg,
    treap_node_t *ctn)
{
	if (ctn == NULL) {
		return;
	}

	_treap_walk(tree, walkfunc, arg, ctn->tn_left);
	walkfunc(arg, n_from_tn(tree, ctn));
	_treap_walk(tree, walkfunc, arg, ctn->tn_right);
}

void
treap_walk(treap_t *tree, void (*walkfunc)(void *arg, void *node), void *arg)
{
	ASSERT3P(tree->t_compare, !=, NULL);
	_treap_walk(tree, walkfunc, arg, tree->t_root);
}

void
_treap_clear(treap_t *tree, void (*clearfunc)(void *arg, void *node), void *arg,
    treap_node_t *ctn)
{
	void *node;

	if (ctn == NULL) {
		return;
	}

	_treap_clear(tree, clearfunc, arg, ctn->tn_left);
	_treap_clear(tree, clearfunc, arg, ctn->tn_right);

	node = n_from_tn(tree, ctn);
	treap_node_invalidate(ctn);
	clearfunc(arg, node);
}

void
treap_clear(treap_t *tree, void (*clearfunc)(void *arg, void *node), void *arg)
{
	ASSERT3P(tree->t_compare, !=, NULL);
	_treap_clear(tree, clearfunc, arg, tree->t_root);
	tree->t_root = NULL;
}