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