// SPDX-License-Identifier: LGPL-2.1 OR BSD-2-Clause
/*
 * Copyright (c) 2025-2026 Meta Platforms, Inc. and affiliates.
 * Copyright (c) 2025-2026 Emil Tsalapatis <emil@etsalapatis.com>
 */

#include <libarena/common.h>

#include <libarena/asan.h>
#include <libarena/rbtree.h>

int rb_integrity_check(struct rbtree __arena *rbtree);
void rbnode_print(size_t depth, struct rbnode __arena *rbn);
static int rbnode_replace(struct rbtree __arena *rbtree,
			  struct rbnode __arena *existing,
			  struct rbnode __arena *replacement);

struct rbtree __arena *rb_create(enum rbtree_alloc alloc,
				 enum rbtree_insert_mode insert)
{
	struct rbtree __arena *rbtree;

	rbtree = arena_malloc(sizeof(*rbtree));
	if (unlikely(!rbtree))
		return NULL;

	/*
	 * RB_UPDATE overwrites existing values in the nodes, but RB_NOALLOC
	 * trees manage the tree nodes directly (including holding pointers
	 * to them). Disallow mixing the two modes to avoid dealing with
	 * unintuitive semantics.
	 */
	if (alloc == RB_NOALLOC && insert == RB_UPDATE) {
		arena_stderr("WARNING: Cannot combine RB_NOALLOC and RB_UPDATE");
		arena_free(rbtree);
		return NULL;
	}

	rbtree->alloc = alloc;
	rbtree->insert = insert;
	rbtree->root = NULL;

	return rbtree;
}

__weak
int rb_destroy(struct rbtree __arena *rbtree)
{
	int ret = 0;

	arena_subprog_init();

	if (unlikely(!rbtree))
		return -EINVAL;

	if (rbtree->alloc == RB_NOALLOC) {
		/*
		 * We cannot do anything about RB_NOALLOC nodes. The whole
		 * point of RB_NOALLOC is that the nodes are directly owned
		 * by the caller that allocates and inserts them. We could
		 * unilaterally grab all nodes and free them anyway, but that
		 * would almost certainly cause UAF as the callers keep accessing
		 * the now freed nodes. Throw an error instead.
		 */
		if (rbtree->root) {
			arena_stderr("WARNING: Destroying RB_NOALLOC tree with > 0 nodes");
			return -EBUSY;
		}

		goto out;
	}

	while (rbtree->root && can_loop) {
		ret = rb_remove(rbtree, rbtree->root->key);
		if (ret)
			break;
	}

out:
	arena_free(rbtree);
	return ret;
}

static inline int rbnode_dir(struct rbnode __arena *node)
{
	/* Arbitrarily choose a direction for the root. */
	if (unlikely(!node->parent))
		return 0;

	return (node->parent->left == node) ? 0 : 1;
}

/*
 * The __noinline is to prevent inlining from bloating the add
 * remove calls, in turn causing register splits and increasing
 * stack usage above what is permitted.
 */
__noinline
int rbnode_rotate(struct rbtree __arena *rbtree,
		  struct rbnode __arena *node, int dir)
{
	struct rbnode __arena *tmp, *parent;
	int parentdir;

	parent = node->parent;
	if (parent)
		parentdir = rbnode_dir(node);

	/* If we're doing a root change, are we the root? */
	if (unlikely(!parent && rbtree->root != node))
		return -EINVAL;

	/*
	 * Does the node we're turning into the root into exist?
	 * Note that the new root is on the opposite side of the
	 * rotation's direction.
	 */
	tmp = node->child[1 - dir];
	if (unlikely(!tmp))
		return -EINVAL;

	/* Steal the closest child of the new root. */
	node->child[1 - dir] = tmp->child[dir];
	if (node->child[1 - dir])
		node->child[1 - dir]->parent = node;

	/* Put the node below the new root.*/
	tmp->child[dir] = node;
	node->parent = tmp;

	tmp->parent = parent;
	if (parent)
		parent->child[parentdir] = tmp;
	else
		rbtree->root = tmp;

	return 0;
}

static
struct rbnode __arena *rbnode_find(struct rbnode __arena *subtree, u64 key)
{
	struct rbnode __arena *node = subtree;
	int dir;

	if (!subtree)
		return NULL;

	while (can_loop) {
		if (node->key == key)
			break;

		dir = (key < node->key) ? 0 : 1;

		if (!node->child[dir])
			break;

		node = node->child[dir];
	}

	return node;
}

static
struct rbnode __arena *rbnode_least_upper_bound(struct rbnode __arena *subtree, uint64_t key)
{
	struct rbnode __arena *node = subtree;
	int dir;

	if (!subtree)
		return NULL;

	while (can_loop) {
		dir = (key <= node->key) ? 0 : 1;

		if (!node->child[dir])
			break;

		node = node->child[dir];
	}

	return node;
}

__weak
int rb_find(struct rbtree __arena *rbtree, u64 key, u64 *value)
{
	struct rbnode __arena *node;

	if (unlikely(!rbtree))
		return -EINVAL;

	if (unlikely(!value))
		return -EINVAL;

	node = rbnode_find(rbtree->root, key);
	if (!node || node->key != key)
		return -ENOENT;

	*value = node->value;

	return 0;
}

__weak
struct rbnode __arena *rb_node_alloc(u64 key, u64 value)
{
	struct rbnode __arena *rbnode = NULL;

	rbnode = (struct rbnode __arena *)arena_malloc(sizeof(*rbnode));
	if (!rbnode)
		return NULL;

	/*
	 * WARNING: The order of assignments is weird on purpose.
	 * See comment in rb_insert_node() for more context.
	 * TL;DR: Prevent consecutive 0 assignments from being
	 * promoted into an unverifiable memset by the compiler.
	 */

	rbnode->key = key;
	rbnode->parent = NULL;
	rbnode->value = value;
	rbnode->left = NULL;
	rbnode->is_red = true;
	rbnode->right = NULL;

	return rbnode;
}

__weak
void rb_node_free(struct rbnode __arena *rbnode)
{
	arena_free(rbnode);
}

static
int rb_node_insert(struct rbtree __arena *rbtree,
		   struct rbnode __arena *node)
{
	struct rbnode __arena *grandparent, *parent = rbtree->root;
	u64 key = node->key;
	struct rbnode __arena *uncle;
	int dir;
	int ret;

	if (unlikely(!rbtree))
		return -EINVAL;

	if (!parent) {
		rbtree->root = node;
		return 0;
	}

	if (rbtree->insert != RB_DUPLICATE)
		parent = rbnode_find(parent, key);
	else
		parent = rbnode_least_upper_bound(parent, key);

	if (key == parent->key && rbtree->insert != RB_DUPLICATE) {
		if (rbtree->insert == RB_UPDATE) {
			/*
			 * Replace the old node with the new one.
			 * Free up the old node.
			 */
			ret = rbnode_replace(rbtree, parent, node);
			if (ret)
				return ret;

			if (rbtree->alloc == RB_ALLOC)
				rb_node_free(parent);

			return 0;
		}

		/* Otherwise it's RB_DEFAULT. */
		return -EALREADY;
	}

	node->parent = parent;
	/* Also works if key == parent->key. */
	if (key <= parent->key)
		parent->left = node;
	else
		parent->right = node;

	while (can_loop) {
		parent = node->parent;
		if (!parent)
			return 0;

		if (!parent->is_red)
			return 0;

		grandparent = parent->parent;
		if (!grandparent) {
			parent->is_red = false;
			return 0;
		}

		dir = rbnode_dir(parent);
		uncle = grandparent->child[1 - dir];

		if (!uncle || !uncle->is_red) {
			if (node == parent->child[1 - dir]) {
				rbnode_rotate(rbtree, parent, dir);
				node = parent;
				parent = grandparent->child[dir];
			}

			rbnode_rotate(rbtree, grandparent, 1 - dir);
			parent->is_red = false;
			grandparent->is_red = true;

			return 0;
		}

		/* Uncle is red. */

		parent->is_red = false;
		uncle->is_red = false;
		grandparent->is_red = true;

		node = grandparent;
	}

	return 0;
}

int rb_insert_node(struct rbtree __arena *rbtree,
		   struct rbnode __arena *node)
{
	if (unlikely(!rbtree))
		return -EINVAL;

	if (unlikely(rbtree->alloc == RB_ALLOC))
		return -EINVAL;

	node->left = NULL;

	/*
	 * Workaround to break an optimization that causes
	 * verification failures on some compilers. Assignments
	 * of the kind
	 *
	 * *(r0 + 0) = 0;
	 * *(r0 + 8) = 0;
	 * *(r0 + 16) = 0;
	 *
	 * get promoted into a memset, and that in turn is not
	 * handled properly for arena memory by LLVM 21 and GCC 15.
	 * Add a barrier for now to prevent the assignments from being fused.
	 */
	barrier();

	node->parent = NULL;
	node->right = NULL;
	
	node->is_red = true;

	return rb_node_insert(rbtree, node);
}

__weak
int rb_insert(struct rbtree __arena *rbtree, u64 key, u64 value)
{
	struct rbnode __arena *node;
	int ret;

	if (unlikely(!rbtree))
		return -EINVAL;

	if (unlikely(rbtree->alloc != RB_ALLOC))
		return -EINVAL;

	node = rb_node_alloc(key, value);
	if (!node)
		return -ENOMEM;

	ret = rb_node_insert(rbtree, node);
	if (ret) {
		rb_node_free(node);
		return ret;
	}

	return 0;
}

static inline struct rbnode __arena *rbnode_least(struct rbnode __arena *subtree)
{
	while (subtree->left && can_loop)
		subtree = subtree->left;

	return subtree;
}

__weak int rb_least(struct rbtree __arena *rbtree, u64 *key, u64 *value)
{
	struct rbnode __arena *least;

	if (unlikely(!rbtree))
		return -EINVAL;

	if (!rbtree->root)
		return -ENOENT;

	least = rbnode_least(rbtree->root);
	if (key)
		*key = least->key;
	if (value)
		*value = least->value;

	return 0;
}


/*
 * If we are referencing ourselves, a and b have a parent-child relation,
 * and we should be pointing at the other node instead.
 */
static inline void rbnode_fixup_pointers(struct rbnode __arena *a,
					 struct rbnode __arena *b)
{
#define fixup(n1, n2, member) do { if (n1->member == n1) n1->member = n2; } while (0)
	fixup(a, b, left);
	fixup(a, b, right);
	fixup(a, b, parent);
#undef fixup
}

static inline void rbnode_swap_values(struct rbnode __arena *a,
				      struct rbnode __arena *b)
{
#define swap(n1, n2, tmp) do { (tmp) = (n1); (n1) = (n2); (n2) = (tmp); } while (0)
	struct rbnode __arena *tmpnode;
	u64 tmp;

	/* Swap the pointers. */
	swap(a->is_red, b->is_red, tmp);

	swap(a->left, b->left, tmpnode);
	swap(a->right, b->right, tmpnode);
	swap(a->parent, b->parent, tmpnode);
#undef swap

	/* Account for the nodes being parent and child. */
	rbnode_fixup_pointers(b, a);
	rbnode_fixup_pointers(a, b);
}

static inline void rbnode_adjust_neighbors(struct rbtree __arena *rbtree,
					   struct rbnode __arena *node, int dir)
{
	if (node->left)
		node->left->parent = node;
	if (node->right)
		node->right->parent = node;

	if (node->parent) {
		node->parent->child[dir] = node;
		return;
	}

	rbtree->root = node;
}

/*
 * Directly replace an existing node with a replacement. The replacement node
 * should not already be in the tree.
 */
static int rbnode_replace(struct rbtree __arena *rbtree,
			  struct rbnode __arena *existing,
			  struct rbnode __arena *replacement)
{
	int dir = 0;

	if (unlikely(replacement->parent || replacement->left || replacement->right))
		return -EINVAL;

	if (existing->parent)
		dir = rbnode_dir(existing);

	replacement->is_red = existing->is_red;
	replacement->left = existing->left;
	replacement->right = existing->right;
	replacement->parent = existing->parent;

	/* Fix up the new node's neighbors. */
	rbnode_adjust_neighbors(rbtree, replacement, dir);

	return 0;
}

/*
 * Switch two nodes in the tree in place. This is useful during node deletion.
 * This is more involved than switching the values of the two nodes because we
 * must update all tree pointers.
 */
static void rbnode_switch(struct rbtree __arena *rbtree,
			  struct rbnode __arena *a,
			  struct rbnode __arena *b)
{
	int adir = 0, bdir = 0;

	/*
	 * Store the direction in the parent because we will not
	 * be able to recompute it once we start swapping values.
	 */
	if (a->parent)
		adir = rbnode_dir(a);

	if (b->parent)
		bdir = rbnode_dir(b);

	rbnode_swap_values(a, b);

	/*
	 * Fix up the pointers from the children/parent to the
	 * new nodes.
	 */
	rbnode_adjust_neighbors(rbtree, a, bdir);
	rbnode_adjust_neighbors(rbtree, b, adir);
}

static inline int rbnode_remove_node_single_child(struct rbtree __arena *rbtree,
						  struct rbnode __arena *node,
						  bool free)
{
	struct rbnode __arena *child;
	int dir;

	if (unlikely(node->is_red)) {
		arena_stderr("Node unexpectedly red\n");
		return -EINVAL;
	}

	child = node->left ? node->left : node->right;
	if (unlikely(!child->is_red)) {
		arena_stderr("Only child is black\n");
		return -EINVAL;
	}

	/*
	 * Since it's the immediate child, we can just
	 * remove the parent.
	 */
	child->parent = node->parent;

	if (node->parent) {
		dir = rbnode_dir(node);
		node->parent->child[dir] = child;
	} else {
		rbtree->root = child;
	}

	/* Color the child black. */
	child->is_red = false;

	/* Only free if called from rb_remove. */
	if (free)
		rb_node_free(node);

	return 0;
}

static inline bool rbnode_has_red_children(struct rbnode __arena *node)
{
	if (node->left && node->left->is_red)
		return true;

	return node->right && node->right->is_red;
}

static
int rb_node_remove(struct rbtree __arena *rbtree,
		   struct rbnode __arena *node)
{
	struct rbnode __arena *parent, *sibling, *close_nephew, *distant_nephew;
	bool free = (rbtree->alloc == RB_ALLOC);
	struct rbnode __arena *replace, *initial;
	bool is_red;
	int dir;

	/* Both children present, replace with next largest key. */
	if (node->left && node->right) {
		/*
		 * Swap the node itself instead of just the
		 * key/value pair to account for nodes embedded
		 * in other structs.
		 */

		replace = rbnode_least(node->right);
		rbnode_switch(rbtree, replace, node);

		/*
		 * FALLTHROUGH: We moved the node we are removing to
		 * the leftmost position of the subtree. We can now
		 * remove it as if it was always where we moved it to.
		 */
	}

	initial = node;

	/* Only one child present, replace with child and paint it black. */
	if (!node->left != !node->right)
		return rbnode_remove_node_single_child(rbtree, node, free);

	/* (!node->left && !node->right) */

	parent = node->parent;
	if (!parent) {
		/* Check that we're _actually_ the root. */
		if (rbtree->root == node)
			rbtree->root = NULL;
		else
			arena_stderr("WARNING: Attempting to remove detached node from rbtree\n");

		if (free)
			rb_node_free(node);
		return 0;
	}

	dir = rbnode_dir(node);
	parent->child[dir] = NULL;
	is_red = node->is_red;

	if (free)
		rb_node_free(node);

	/* If we removed a red node, we did not unbalance the tree.*/
	if (is_red)
		return 0;

	sibling = parent->child[1 - dir];
	if (unlikely(!sibling)) {
		arena_stderr("rbtree: removed black node has no sibling\n");
		return -EINVAL;
	}

	/*
	 * We removed a black node, causing a change in path
	 * weight. Start rebalancing. The invariant is that
	 * all paths going through the node are shortened
	 * by one, and the current node is black.
	 */
	while (can_loop) {

		/* Balancing reached the root, there can be no imbalance. */
		if (!parent)
			return 0;

		/*
		 * We already determined the dir, either above or
		 * at the end of the loop.
		 */

		/*
		 * If we have no sibling, the tree was
		 * already unbalanced.
		 */
		sibling = parent->child[1 - dir];
		if (unlikely(!sibling)) {
			arena_stderr("rbtree: removed black node has no sibling\n");
			return -EINVAL;
		}

		/* Sibling is red, turn it into the grandparent. */
		if (sibling->is_red) {
			/*
			 * Sibling is red. Transform the tree to turn
			 * the sibling into the parent's position, and
			 * repaint them. This does not balance the tree
			 * but makes it so we know the sibling is black
			 * and so can use the transformations to balance.
			 */
			rbnode_rotate(rbtree, parent, dir);
			parent->is_red = true;
			sibling->is_red = false;

			/* Our new sibling is now the close nephew. */
			sibling = parent->child[1 - dir];
			/* If sibling has any red siblings, break out. */
			if (rbnode_has_red_children(sibling))
				break;

			/* We can repaint the sibling and parent, we're done. */
			sibling->is_red = true;
			parent->is_red = false;

			return 0;
		}

		/* Sibling guaranteed to be black. If it has red children, break out. */
		if (rbnode_has_red_children(sibling))
			break;

		/*
		 * Both sibling and children are black. If parent is red, swap
		 * colors with the sibling. Otherwise
		 */
		if (parent->is_red) {
			parent->is_red = false;
			sibling->is_red = true;
			return 0;
		}

		/*
		 * Parent, sibling, and all its children are black. Repaint the sibling.
		 * This shortens the paths through it, so pop up a level in the
		 * tree and repeat the balancing.
		 */
		sibling->is_red = true;
		node = parent;
		parent = node->parent;
		dir = rbnode_dir(node);
	}

	if (node != initial) {
		dir = rbnode_dir(node);
		parent = node->parent;
		sibling = parent->child[1-dir];
	}
	/*
	 * Almost there. We know between the parent, sibling,
	 * and nephews only one or two of the nephews are red. If
	 * it is the close one, rotate it to the sibling position,
	 * paint it black, and paint the previous sibling red.
	 */

	close_nephew = sibling->child[dir];
	distant_nephew = sibling->child[1 - dir];

	/*
	 * If the distant red nephew is not red, rotate
	 * and repaint. We need the distant nephew
	 * to be red. We know the close nephew is red
	 * because at least one of them are, so the
	 * distant one is black if it exists.
	 */
	if (!distant_nephew || !distant_nephew->is_red) {
		rbnode_rotate(rbtree, sibling, 1 - dir);
		sibling->is_red = true;
		close_nephew->is_red = false;
		distant_nephew = sibling;
		sibling = close_nephew;
	}

	/*
	 * We now know it's the distant nephew that's red.
	 * Rotate the sibling into our parent's position
	 * and paint both black.
	 */

	rbnode_rotate(rbtree, parent, dir);
	sibling->is_red = parent->is_red;
	parent->is_red = false;
	distant_nephew->is_red = false;

	return 0;
}

__weak
int rb_remove_node(struct rbtree __arena *rbtree,
		   struct rbnode __arena *node)
{
	if (unlikely(!rbtree))
		return -EINVAL;

	if (unlikely(rbtree->alloc == RB_ALLOC))
		return -EINVAL;

	return rb_node_remove(rbtree, node);
}

__weak
int rb_remove(struct rbtree __arena *rbtree, u64 key)
{
	struct rbnode __arena *node;

	if (unlikely(!rbtree))
		return -EINVAL;

	if (unlikely(rbtree->alloc != RB_ALLOC))
		return -EINVAL;

	if (!rbtree->root)
		return -ENOENT;

	node = rbnode_find(rbtree->root, key);
	if (!node || node->key != key)
		return -ENOENT;

	return rb_node_remove(rbtree, node);
}

__weak
int rb_pop(struct rbtree __arena *rbtree, u64 *key, u64 *value)
{
	struct rbnode __arena *node;

	if (unlikely(!rbtree))
		return -EINVAL;

	if (!rbtree->root)
		return -ENOENT;

	if (rbtree->alloc != RB_ALLOC)
		return -EINVAL;

	node = rbnode_least(rbtree->root);
	if (unlikely(!node))
		return -ENOENT;

	if (key)
		*key = node->key;
	if (value)
		*value = node->value;

	return rb_node_remove(rbtree, node);
}

inline void rbnode_print(size_t depth, struct rbnode __arena *rbn)
{
	arena_stderr("[DEPTH %d] %p (%s)\n PARENT %p", depth, rbn, rbn->is_red ? "red" : "black", rbn->parent);
	arena_stderr("\tKV (%ld, %ld)\n LEFT %p RIGHT %p]\n", rbn->key, rbn->value, rbn->left, rbn->right);
}

enum rb_print_state {
	RB_NONE_VISITED,
	RB_LEFT_VISITED,
	RB_RIGHT_VISITED,
};

__weak
enum rb_print_state rb_print_next_state(struct rbnode __arena *rbnode,
					enum rb_print_state state, u64 *next)
{
	if (unlikely(!next))
		return RB_NONE_VISITED;

	switch (state) {
	case RB_NONE_VISITED:
		if (rbnode->left) {
			*next = (u64)rbnode->left;
			state = RB_LEFT_VISITED;
			break;
		}

		/* FALLTHROUGH */

	case RB_LEFT_VISITED:
		if (rbnode->right) {
			*next = (u64)rbnode->right;
			state = RB_RIGHT_VISITED;
			break;
		}

		/* FALLTHROUGH */

	default:
		*next = 0;
		state = RB_RIGHT_VISITED;
	}

	return state;
}

__weak
int rb_print_pop_up(struct rbnode __arena **rbnodep, u8 *depthp, enum rb_print_state (*stack)[RB_MAXLVL_PRINT], enum rb_print_state *state)
{
	struct rbnode __arena *rbnode;
	volatile u8 depth;
	int j;

	if (unlikely(!rbnodep || !depthp || !stack || !state))
		return -EINVAL;

	rbnode = *rbnodep;
	depth = *depthp;

	for (j = 0; j < RB_MAXLVL_PRINT && can_loop; j++) {
		if (*state != RB_RIGHT_VISITED)
			break;

		depth -= 1;
		if (depth < 0 || depth >= RB_MAXLVL_PRINT)
			break;

		*state = (*stack)[depth % RB_MAXLVL_PRINT];
		rbnode = rbnode->parent;
	}

	*rbnodep = rbnode;
	*depthp = depth;

	return 0;
}

__weak
int rb_print(struct rbtree __arena *rbtree)
{
	enum rb_print_state stack[RB_MAXLVL_PRINT];
	struct rbnode __arena *rbnode = rbtree->root;
	enum rb_print_state state;
	struct rbnode __arena *next;
	u64 next_addr;
	u8 depth;
	int ret;

	if (unlikely(!rbtree))
		return -EINVAL;

	depth = 0;
	state = RB_NONE_VISITED;

	arena_stderr("=== RB TREE START ===\n");

	if (!rbtree->root)
		goto out;

	/* Even with can_loop, the verifier doesn't like infinite loops. */
	while (can_loop) {
		if (state == RB_NONE_VISITED)
			rbnode_print(depth, rbnode);

		/* Find which child to traverse next. */
		state = rb_print_next_state(rbnode, state, &next_addr);
		next = (struct rbnode __arena *)next_addr;

		/* Child found. Store the node state and go on. */
		if (next) {
			if (depth < 0 || depth >= RB_MAXLVL_PRINT)
				return 0;

			stack[depth++] = state;

			rbnode = next;
			state = RB_NONE_VISITED;

			continue;
		}

		/* Otherwise, go as far up as possible. */
		ret = rb_print_pop_up(&rbnode, &depth, &stack, &state);
		if (ret)
			return -EINVAL;

		if (depth < 0 || depth >= RB_MAXLVL_PRINT) {
			arena_stderr("=== RB TREE END (depth %d\n)===", depth);
			return 0;
		}

	}

out:
	arena_stderr("=== RB TREE END ===\n");

	return 0;
}

__weak
int rb_integrity_check(struct rbtree __arena *rbtree)
{
	enum rb_print_state stack[RB_MAXLVL_PRINT];
	struct rbnode __arena *rbnode = rbtree->root;
	enum rb_print_state state;
	struct rbnode __arena *next;
	u64 next_addr;
	u8 depth;
	int ret;

	if (unlikely(!rbtree))
		return -EINVAL;

	if (!rbtree->root)
		return 0;

	depth = 0;
	state = RB_NONE_VISITED;

	/* Even with can_loop, the verifier doesn't like infinite loops. */
	while (can_loop) {
		if (rbnode->parent && rbnode->parent->left != rbnode
			&& rbnode->parent->right != rbnode) {
			arena_stderr("WARNING: Inconsistent tree. Parent %p has no child %p\n", rbnode->parent, rbnode);
			return -EINVAL;
		}

		if (rbnode->parent == rbnode) {
			arena_stderr("WARNING: Inconsistent tree, node %p is its own parent\n", rbnode);
			return -EINVAL;
		}

		if (rbnode->left == rbnode) {
			arena_stderr("WARNING: Inconsistent tree, node %p is its own left child\n", rbnode);
			return -EINVAL;
		}

		if (rbnode->right == rbnode) {
			arena_stderr("WARNING: Inconsistent tree, node %p is its own right child\n", rbnode);
			return -EINVAL;
		}

		if (rbnode->is_red) {
			if (rbnode->left && rbnode->left->is_red) {
				arena_stderr("WARNING: Inconsistent tree. Parent has %p has red child %p\n", rbnode, rbnode->left);
				return -EINVAL;
			}
			if (rbnode->right && rbnode->right->is_red) {
				arena_stderr("WARNING: Inconsistent tree. Parent has %p has red child %p\n", rbnode, rbnode->right);
				return -EINVAL;
			}
		} else if (rbnode->parent && rbnode->parent->child[1 - rbnode_dir(rbnode)] == NULL) {
			arena_stderr("WARNING: Inconsistent tree. Black node %p has no sibling\n", rbnode);
			return -EINVAL;
		}

		/* Find which child to traverse next. */
		state = rb_print_next_state(rbnode, state, &next_addr);
		next = (struct rbnode __arena *)next_addr;

		/* Child found. Store the node state and go on. */
		if (next) {
			if (depth < 0 || depth >= RB_MAXLVL_PRINT)
				return 0;

			stack[depth++] = state;

			rbnode = next;
			state = RB_NONE_VISITED;

			continue;
		}

		/* Otherwise, go as far up as possible. */
		ret = rb_print_pop_up(&rbnode, &depth, &stack, &state);
		if (ret)
			return -EINVAL;

		if (depth < 0 || depth >= RB_MAXLVL_PRINT) {
			return 0;
		}

	}

	return 0;
}