rbtree: add rb_find_add_cached() to rbtree.h

Adds rb_find_add_cached() as a helper function for use with red-black
trees. Used in btrfs to reduce boilerplate code.

And since it's a new helper, the cmp() function will require both
parameter to be const rb_node pointers.

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>
Reviewed-by: Qu Wenruo <wqu@suse.com>
Signed-off-by: Qu Wenruo <wqu@suse.com>
Reviewed-by: David Sterba <dsterba@suse.com>
Signed-off-by: David Sterba <dsterba@suse.com>

diff --git a/include/linux/rbtree.h b/include/linux/rbtree.h
index 7c173aa64e1e3f26a38f94ffa5c5031d9aa99c53..8d2ba3749866f500a492d267e5e13556f6aa3f55 100644 (file)
--- 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)(const struct rb_node *new, const struct rb_node *exist))
+{
+       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