Re: [PATCH 1/6] rbtree: add rb_find_add_cached() to rbtree.h

[Peter Zijlstra <peterz@infradead.org>]

On Thu, Dec 12, 2024 at 10:46:18AM -0600, Roger L. Beckermeyer III wrote:
> Adds rb_find_add_cached() as a helper function for use with
> red-black trees. Used in btrfs to reduce boilerplate code.
> 
> Suggested-by: Josef Bacik <josef@toxicpanda.com>
> Signed-off-by: Roger L. Beckermeyer III <beckerlee3@gmail.com>

Acked-by: Peter Zijlstra (Intel) <peterz@infradead.org>

> ---
>  include/linux/rbtree.h | 37 +++++++++++++++++++++++++++++++++++++
>  1 file changed, 37 insertions(+)
> 
> diff --git a/include/linux/rbtree.h b/include/linux/rbtree.h
> index 7c173aa64e1e..0d4444c0cfb3 100644
> --- a/include/linux/rbtree.h
> +++ b/include/linux/rbtree.h
> @@ -210,6 +210,43 @@ rb_add(struct rb_node *node, struct rb_root *tree,
>  	rb_insert_color(node, tree);
>  }
>  
> +/**
> + * rb_find_add_cached() - find equivalent @node in @tree, or add @node
> + * @node: node to look-for / insert
> + * @tree: tree to search / modify
> + * @cmp: operator defining the node order
> + *
> + * Returns the rb_node matching @node, or NULL when no match is found and @node
> + * is inserted.
> + */
> +static __always_inline struct rb_node *
> +rb_find_add_cached(struct rb_node *node, struct rb_root_cached *tree,
> +	    int (*cmp)(struct rb_node *, const struct rb_node *))
> +{
> +	bool leftmost = true;
> +	struct rb_node **link = &tree->rb_root.rb_node;
> +	struct rb_node *parent = NULL;
> +	int c;
> +
> +	while (*link) {
> +		parent = *link;
> +		c = cmp(node, parent);
> +
> +		if (c < 0) {
> +			link = &parent->rb_left;
> +		} else if (c > 0) {
> +			link = &parent->rb_right;
> +			leftmost = false;
> +		} else {
> +			return parent;
> +		}
> +	}
> +
> +	rb_link_node(node, parent, link);
> +	rb_insert_color_cached(node, tree, leftmost);
> +	return NULL;
> +}
> +
>  /**
>   * rb_find_add() - find equivalent @node in @tree, or add @node
>   * @node: node to look-for / insert
> -- 
> 2.45.2
>