From: Liam R. Howlett <Liam.Howlett@oracle.com>
Date: Wed, 10 Nov 2021 03:17:05 +0000 (-0500)
Subject: maple_tree: Add mt_lookup() with no state support for faster lookups
X-Git-Url: https://www.infradead.org/git/?a=commitdiff_plain;h=9472f1f66cf2df4b34e2c9a009bc4da646e2a92e;p=users%2Fjedix%2Flinux-maple.git

maple_tree: Add mt_lookup() with no state support for faster lookups

Signed-off-by: Liam R. Howlett <Liam.Howlett@oracle.com>
---

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 8f4d044f73f4..ae2c2b58fa49 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -3708,6 +3708,86 @@ done:
 		mas->min = min;
 	}
 }
+
+/*
+ * mas_lookup_walk() - Internal quick lookup that does not keep maple state up
+ * to date.
+ *
+ * @mas: The maple state.
+ *
+ * Return: The entry for @mas->index or %NULL on dead node.
+ */
+static inline void *mas_lookup_walk(struct ma_state *mas)
+{
+	unsigned long *pivots;
+	unsigned char offset;
+	struct maple_node *node;
+	struct maple_enode *next;
+	enum maple_type type;
+	void __rcu **slots;
+	unsigned char end;
+	unsigned long max;
+
+	next = mas->node;
+	max = ULONG_MAX;
+	do {
+		offset = 0;
+		node = mte_to_node(next);
+		type = mte_node_type(next);
+		pivots = ma_pivots(node, type);
+		end = ma_data_end(node, type, pivots, max);
+		if (unlikely(ma_dead_node(node)))
+			goto dead_node;
+
+		while((offset < end) &&
+		      (pivots[offset] < mas->index))
+			offset++;
+
+		if (likely(offset > mt_pivots[type]) && pivots[offset])
+			max = pivots[offset];
+
+		slots = ma_slots(node, type);
+		next = mt_slot(mas->tree, slots, offset);
+		if (unlikely(ma_dead_node(node)))
+			goto dead_node;
+	} while (!ma_is_leaf(type));
+
+	return (void *) next;
+
+dead_node:
+	mas_reset(mas);
+	return NULL;
+}
+
+/*
+ * mt_lookup() - Fast value lookup.
+ *
+ */
+void *mt_lookup(struct maple_tree *mt, unsigned long index)
+{
+	MA_STATE(mas, mt, index, index);
+	void *entry;
+
+retry:
+	entry = mas_start(&mas);
+	if (unlikely(mas_is_none(&mas)))
+		return NULL;
+
+	if (unlikely(mas_is_ptr(&mas))) {
+		if (!index)
+			return entry;
+
+		return NULL;
+	}
+
+	entry = mas_lookup_walk(&mas);
+	if (!entry && unlikely(mas_is_start(&mas)))
+		goto retry;
+
+	return entry;
+}
+EXPORT_SYMBOL_GPL(mt_lookup);
+
 /*
  * mas_descend_walk(): Locates a value and sets the mas->node and slot
  * accordingly.  range_min and range_max are set to the range which the entry is
diff --git a/lib/test_maple_tree.c b/lib/test_maple_tree.c
index 1a54902b4eca..0a9d128b9d85 100644
--- a/lib/test_maple_tree.c
+++ b/lib/test_maple_tree.c
@@ -97,6 +97,11 @@ static noinline void check_load(struct maple_tree *mt, unsigned long index,
 {
 	void *ret = mtree_test_load(mt, index);
 
+	if (ret != mt_lookup(mt, index)) {
+		pr_err("lookup of %lu %p expected %p\n", index, mt_lookup(mt, index), ret);
+	}
+	MT_BUG_ON(mt, ret != mt_lookup(mt, index));
+
 	if (ret != ptr)
 		pr_err("Load %lu returned %p expect %p\n", index, ret, ptr);
 	MT_BUG_ON(mt, ret != ptr);