From: Liam R. Howlett <Liam.Howlett@oracle.com>
Date: Tue, 25 Apr 2023 20:10:31 +0000 (-0400)
Subject: maple_tree: Add ma_peek() interface
X-Git-Url: https://www.infradead.org/git/?a=commitdiff_plain;h=60646824e8992553554f9beede8025b41f55b51f;p=users%2Fjedix%2Flinux-maple.git

maple_tree: Add ma_peek() interface

Add a thin lookup which uses minimal state.

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

diff --git a/include/linux/maple_tree.h b/include/linux/maple_tree.h
index bddced4ec7f2..dcb3ec36568f 100644
--- a/include/linux/maple_tree.h
+++ b/include/linux/maple_tree.h
@@ -465,6 +465,7 @@ void maple_tree_init(void);
 void mas_destroy(struct ma_state *mas);
 int mas_expected_entries(struct ma_state *mas, unsigned long nr_entries);
 
+void *ma_peek(struct maple_tree *mt, unsigned long addr);
 void *mas_prev(struct ma_state *mas, unsigned long min);
 void *mas_prev_range(struct ma_state *mas, unsigned long max);
 void *mas_next(struct ma_state *mas, unsigned long max);
diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 04b73499baff..9a01cd7dd449 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -5025,6 +5025,62 @@ done:
 	return found;
 }
 
+/* ma_peek() - Cheaper version of walking to a specific address with minimal
+ * state.
+ * @mt: The maple tree
+ * @addr: The address
+ *
+ * Note: Only works holding the write lock.
+ */
+void *ma_peek(struct maple_tree *mt, unsigned long addr)
+{
+	struct maple_enode *enode;
+	struct maple_node *node;
+	unsigned long max;
+	unsigned long *pivots;
+	unsigned char offset, end;
+	enum maple_type type;
+	void __rcu **slots;
+
+	enode = rcu_dereference_protected(mt->ma_root, mt_locked(mt));
+
+	if (unlikely(!enode))
+		return NULL;
+
+	if (unlikely(!xa_is_node(enode))) {
+		if (addr == 0)
+			return (void *)enode;
+
+		return NULL;
+	}
+
+	enode = mte_safe_root(enode);
+	max = ULONG_MAX;
+	do {
+		offset = 0;
+		node = mte_to_node(enode);
+		type = mte_node_type(enode);
+		pivots = ma_pivots(node, type);
+
+		if (pivots[0] >= addr)
+			goto descend;
+
+		end = ma_data_end(node, type, pivots, max);
+		do {
+			offset++;
+		} while ((offset < end) && (pivots[offset] < addr));
+
+descend:
+		if (pivots[offset])
+			max = pivots[offset];
+		slots = ma_slots(node, type);
+		enode = mt_slot_locked(mt, slots, offset);
+	} while (!ma_is_leaf(type));
+
+	return (void *) enode;
+}
+EXPORT_SYMBOL_GPL(ma_peek);
+
 /**
  * mas_walk() - Search for @mas->index in the tree.
  * @mas: The maple state.
diff --git a/lib/test_maple_tree.c b/lib/test_maple_tree.c
index adbf59542951..0f2874b95980 100644
--- a/lib/test_maple_tree.c
+++ b/lib/test_maple_tree.c
@@ -3500,6 +3500,33 @@ static noinline void __init check_slot_iterators(struct maple_tree *mt)
 	mt_set_non_kernel(0);
 }
 
+static noinline void __init check_peek(struct maple_tree *mt)
+{
+	MA_STATE(mas, mt, 0, 0);
+	unsigned long i, index = 40;
+	void *tmp;
+
+	mt_set_non_kernel(99999);
+
+	mas_lock(&mas);
+	for (i = 0; i <= index; i++) {
+		unsigned long end = 5;
+		if (i > 18 && i < 35)
+			end = 9;
+		mas_set_range(&mas, i*10, i*10 + end);
+		mas_store_gfp(&mas, xa_mk_value(i), GFP_KERNEL);
+	}
+
+	mas_set_range(&mas, 18*10, 18*10 + 2);
+	mas_store_gfp(&mas, xa_mk_value(18), GFP_KERNEL);
+	mas_set(&mas, 0);
+	mas_for_each(&mas, tmp, ULONG_MAX)
+		MAS_BUG_ON(&mas, tmp != ma_peek(mt, mas.index));
+
+	mas_unlock(&mas);
+	mt_set_non_kernel(0);
+}
+
 static DEFINE_MTREE(tree);
 static int __init maple_tree_seed(void)
 {
@@ -3769,6 +3796,10 @@ static int __init maple_tree_seed(void)
 	check_slot_iterators(&tree);
 	mtree_destroy(&tree);
 
+	mt_init_flags(&tree, MT_FLAGS_ALLOC_RANGE);
+	check_peek(&tree);
+	mtree_destroy(&tree);
+
 #if defined(BENCH)
 skip:
 #endif