From 5b51ca823baf56ee87c161b9ed9dbc6a112023d0 Mon Sep 17 00:00:00 2001
From: "Liam R. Howlett" <Liam.Howlett@Oracle.com>
Date: Wed, 31 Jul 2019 12:38:16 -0400
Subject: [PATCH] mas_find implementation

Fixes for prev/next node as well

Signed-off-by: Liam R. Howlett <Liam.Howlett@Oracle.com>
---
 lib/maple_tree.c      | 536 +++++++++++++++++++++++++-----------------
 lib/test_maple_tree.c |  39 +++
 2 files changed, 359 insertions(+), 216 deletions(-)

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 5abe5e5eec30..cc1b5b2b8081 100644
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -16,6 +16,31 @@
 #define ma_mnode_ptr(x) ((struct maple_node*)(x))
 #define ma_enode_ptr(x) ((struct maple_enode*)(x))
 
+/**
+ * mas_for_each() - Iterate over a range of the maple tree.
+ * @mas: Maple Tree operation state (maple_state)
+ * @entry: Entry retrieved from the tree
+ * @max: maximum index to retrieve from the tree
+ *
+ * Note: may return the zero entry.
+ *
+ */
+#define mas_for_each(mas, entry, max) \
+	while ((entry = mas_find(mas, max)) != NULL)
+
+/**
+ * mt_for_each - Iterate over a range of the tree.
+ *
+ *
+ *
+ * Note: Will not return the zero entry.
+ *
+ *
+ */
+#define mt_for_each(tree, entry, index, last, max) \
+
+
+
 static struct kmem_cache *maple_node_cache;
 
 unsigned long mt_max[] = {
@@ -138,6 +163,11 @@ static inline bool mt_is_reserved(const void *entry)
 	return ((unsigned long)entry < 4096) && xa_is_internal(entry);
 }
 
+static inline bool mt_is_advanced(const void *entry)
+{
+	return xa_is_internal(entry) && (entry <= XA_RETRY_ENTRY);
+}
+
 static inline void mas_set_err(struct ma_state *mas, long err)
 {
 	mas->node = MA_ERROR(err);
@@ -418,6 +448,9 @@ static inline unsigned long ma_get_pivot(const struct maple_enode *mn,
 {
 	return _ma_get_pivot(mn, slot, mt_node_type(mn));
 }
+
+/* ma_get_safe_pivot() - Return the pivot or the mas->max.
+ */
 static inline unsigned long ma_get_safe_pivot(const struct ma_state *mas,
 					 unsigned char slot)
 {
@@ -688,6 +721,38 @@ static inline void ma_encoded_parent(struct ma_state *mas)
 	mas_update_limits(mas, slot, gptype);
 	mas->node = p_enode;
 }
+/* ma_get_prev_pivot() - Return the previous pivot.
+ *
+ * Mainly for extracting the previous pivot in the case of slot = 0.
+ * Walk up the tree until the minimum changes, then get the previous pivot.
+ *
+ */
+static inline unsigned long ma_get_prev_pivot(const struct ma_state *mas,
+					 unsigned char slot)
+{
+	unsigned char p_slot = MAPLE_NODE_SLOTS; // parent slot.
+
+	if (slot > 0)
+		return ma_get_safe_pivot(mas, slot - 1);
+
+	if (mas->min == 0)
+		return 0;
+
+	MA_STATE(prev_piv, mas->tree, 0, 0);
+
+	prev_piv.node = mas->node;
+	prev_piv.min = mas->min;
+
+	while (prev_piv.min == mas->min) {
+		p_slot = mt_parent_slot(prev_piv.node);
+		ma_encoded_parent(&prev_piv);
+		if (p_slot)
+			break;
+	}
+	p_slot--;
+	return ma_get_pivot(prev_piv.node, slot);
+}
+
 static inline struct maple_node *ma_next_alloc(struct ma_state *ms)
 {
 	int cnt;
@@ -1976,6 +2041,210 @@ static inline void ma_root_expand(struct ma_state *ms, void *entry)
 	rcu_assign_pointer(ms->tree->ma_root, mt_mk_root(ms->node));
 }
 
+static inline int mas_dead_node(struct ma_state *mas, unsigned long index);
+
+/*
+ * Find the prev non-null entry at the same level in the tree.  The prev value
+ * will be mas->node[ma_get_slot(mas)] or MAS_NONE.
+ */
+static inline void mas_prev_node (struct ma_state *mas, unsigned long min)
+{
+	int level;
+	unsigned char slot;
+	unsigned long start_piv;
+
+restart_prev_node:
+	level = 0;
+	slot = ma_get_slot(mas);
+	start_piv = ma_get_safe_pivot(mas, slot);
+	if (ma_is_root(mas->node))
+		goto no_entry;
+
+	while (1) {
+		ma_encoded_parent(mas);
+		level++;
+
+		if (mas_dead_node(mas, start_piv))
+			goto restart_prev_node;
+
+		if (!slot)
+			goto ascend;
+
+		slot--;
+		do {
+			struct maple_enode *mn;
+			unsigned long last_pivot;
+			unsigned long pivot = ma_get_safe_pivot(mas, slot);
+
+			if (pivot < min)
+				goto no_entry;
+
+			if (slot != 0 && pivot == 0)
+				break;
+
+			mn = ma_get_rcu_slot(mas->node, slot);
+			if (!mn)
+				continue;
+
+			if (level == 1) {
+				ma_set_slot(mas, slot);
+				mas->node = mn;
+				if (mas_dead_node(mas, start_piv))
+					goto restart_prev_node;
+				return;
+			}
+
+			level--;
+			mas->node = mn;
+			slot = ma_data_end(mas, mt_node_type(mn), &last_pivot);
+		} while(slot-- > 0);
+
+ascend:
+		if (ma_is_root(mas->node))
+			goto no_entry;
+
+		slot = mt_parent_slot(mas->node);
+	}
+
+no_entry:
+	mas->node = MAS_NONE;
+	return;
+}
+
+/*
+ * Find the next non-null entry at the same level in the tree.  The next value
+ * will be mas->node[ma_get_slot(mas)] or MAS_NONE.
+ *
+ */
+static inline unsigned long
+mas_next_node (struct ma_state *mas, unsigned long max)
+{
+	int level;
+	unsigned long start_piv;
+
+restart_next_node:
+	level = 0;
+	while (1) {
+		unsigned char count, slot;
+		struct maple_enode *mn;
+		unsigned long prev_piv;
+
+		mn = mas->node;
+		slot = ma_get_slot(mas);
+		start_piv = ma_get_safe_pivot(mas, slot);
+		level++;
+		if (!ma_is_root(mas->node))
+			ma_encoded_parent(mas);
+
+		if (mas_dead_node(mas, start_piv))
+			goto restart_next_node;
+
+		count = mt_slot_count(mas->node);
+
+		prev_piv = ma_get_safe_pivot(mas, slot);
+
+		while (++slot < count) {
+			unsigned long pivot = ma_get_safe_pivot(mas, slot);
+
+			if (prev_piv > max)
+				goto no_entry;
+
+			if (slot != 0 && pivot == 0)
+				break;
+
+			mn = ma_get_rcu_slot(mas->node, slot);
+			if (!mn) {
+				prev_piv = pivot;
+				continue;
+			}
+
+			mas->min = prev_piv + 1;
+			mas->max = pivot;
+
+			if (level == 1) {
+				ma_set_slot(mas, slot);
+				mas->node = mn;
+				if (mas_dead_node(mas, start_piv))
+					goto restart_next_node;
+				return pivot;
+			}
+
+			level--;
+			mas->node = mn;
+			slot = -1;
+			count = mt_slot_count(mas->node);
+		}
+
+		if (ma_is_root(mas->node))
+			goto no_entry;
+	}
+
+no_entry:
+	mas->node = MAS_NONE;
+	return mas->max;
+
+}
+
+// Next node entry or return 0 on none.
+static inline bool
+mas_next_nentry (struct ma_state *mas, unsigned long max, unsigned long *piv)
+{
+	unsigned long pivot = mas->min;
+	unsigned char slot = ma_get_slot(mas);
+	unsigned char count = mt_slot_count(mas->node);
+	void *entry;
+
+	while (slot < count) {
+		pivot = ma_get_safe_pivot(mas, slot);
+
+		if (slot != 0 && pivot == 0)
+			goto no_entry;
+
+		*piv = pivot;
+
+		/* Valid pivot */;
+
+		entry = ma_get_rcu_slot(mas->node, slot);
+		if (entry)
+			break;
+
+		/* Ran over the limit, this is was the last slot to try */
+		if (pivot >= max)
+			goto no_entry;
+
+		slot++;
+	}
+	ma_set_slot(mas, slot);
+	return true;
+
+no_entry:
+	return false;
+}
+static inline unsigned long
+mas_next_entry (struct ma_state *mas, unsigned long max)
+{
+	unsigned long pivot = max;
+
+	if (mas->node == MAS_NONE)
+		return max;
+
+	if (!mt_is_leaf(mas->node))
+		pivot = mas_first_entry(mas, max);
+	else {
+		while (mas->node != MAS_NONE) {
+			unsigned char p_slot;
+			if (mas_next_nentry(mas, max, &pivot))
+				break;
+
+			p_slot = mt_parent_slot(mas->node);
+			ma_set_slot(mas, p_slot);
+			mas_next_node(mas, max);
+			if (mas->node != MAS_NONE)
+				ma_set_slot(mas, 0);
+		}
+	}
+	return pivot;
+}
 /* Private
  * Combine nulls with the same pivot value
  *
@@ -2322,7 +2591,6 @@ static inline bool __mas_walk(struct ma_state *mas)
 	unsigned long pivot_cnt, max, min, i;
 	bool ret = false;
 
-	mas->node = mas_start(mas);
 	min = mas->min;
 	max = mas->max;
 
@@ -2391,201 +2659,71 @@ static inline int mas_dead_node(struct ma_state *mas, unsigned long index)
 	_mas_walk(mas);
 	return 1;
 }
-/*
- * Find the prev non-null entry at the same level in the tree.  The prev value
- * will be mas->node[ma_get_slot(mas)] or MAS_NONE.
- * 
- */
-static inline void mas_prev_node (struct ma_state *mas, unsigned long min)
-{
-	int level;
-	unsigned char slot;
-	unsigned long start_piv;
 
-restart_prev_node:
-	level = 0;
-	slot = ma_get_slot(mas);
-	start_piv = ma_get_safe_pivot(mas, slot);
-	if (ma_is_root(mas->node))
-		goto no_entry;
-
-	while (1) {
-		ma_encoded_parent(mas);
-		level++;
-
-		if (mas_dead_node(mas, start_piv))
-			goto restart_prev_node;
-
-		if (!slot)
-			goto ascend;
-
-		slot--;
-		do {
-			struct maple_enode *mn;
-			unsigned long last_pivot;
-			unsigned long pivot = ma_get_safe_pivot(mas, slot);
-
-			if (pivot < min)
-				goto no_entry;
-
-			if (slot != 0 && pivot == 0)
-				break;
-
-			mn = ma_get_rcu_slot(mas->node, slot);
-			if (!mn)
-				continue;
-
-			if (level == 1) {
-				ma_set_slot(mas, slot);
-				mas->node = mn;
-				if (mas_dead_node(mas, start_piv))
-					goto restart_prev_node;
-				return;
-			}
-
-			level--;
-			mas->node = mn;
-			slot = ma_data_end(mas, mt_node_type(mn), &last_pivot);
-		} while(slot-- > 0);
-
-ascend:
-		if (ma_is_root(mas->node))
-			goto no_entry;
-
-		slot = mt_parent_slot(mas->node);
-	}
-
-no_entry:
-	mas->node = MAS_NONE;
-	return;
-}
-
-/*
- * Find the next non-null entry at the same level in the tree.  The next value
- * will be mas->node[ma_get_slot(mas)] or MAS_NONE.
+/* mas_find:
+ * Finds first entry that overlaps the range.
+ *   If mas->node is root, finds mas->index and returns the value.  If the
+ *   value is null, returns the next non-null range.
+ *   If mas->node is not root, finds the next range after the range that
+ *   contains mas->index.
+ *
+ * If an entry exists, last and index are updated accordingly.
  *
+ * returns entry or null
  */
-static inline unsigned long
-mas_next_node (struct ma_state *mas, unsigned long max)
+void *mas_find(struct ma_state *mas, unsigned long max)
 {
-	int level;
-	unsigned long start_piv;
+	void *entry = NULL;
+	unsigned long last_piv;
+	unsigned char slot = 0;
 
-restart_next_node:
-	level = 0;
-	while (1) {
-		unsigned char count, slot;
-		struct maple_enode *mn;
-		unsigned long prev_piv;
+	if (!mas->node)
+		return MAS_NONE;
 
-		mn = mas->node;
+	if (mas->node == MAS_START) {
+		_mas_walk(mas);
 		slot = ma_get_slot(mas);
-		start_piv = ma_get_safe_pivot(mas, slot);
-		level++;
-		if (!ma_is_root(mas->node))
-			ma_encoded_parent(mas);
-
-		if (mas_dead_node(mas, start_piv))
-			goto restart_next_node;
-
-		count = mt_slot_count(mas->node);
-
-		prev_piv = ma_get_safe_pivot(mas, slot);
-
-		while (++slot < count) {
-			unsigned long pivot = ma_get_safe_pivot(mas, slot);
-
-			if (prev_piv > max)
-				goto no_entry;
-
-			if (slot != 0 && pivot == 0)
-				break;
-
-			mn = ma_get_rcu_slot(mas->node, slot);
-			if (!mn) {
-				prev_piv = pivot;
-				continue;
-			}
-
-			mas->min = prev_piv + 1;
-			mas->max = pivot;
-
-			if (level == 1) {
-				ma_set_slot(mas, slot);
-				mas->node = mn;
-				if (mas_dead_node(mas, start_piv))
-					goto restart_next_node;
-				return pivot;
-			}
+		last_piv = ma_get_safe_pivot(mas, slot);
+		if (slot != MAPLE_NODE_SLOTS)
+			entry = ma_get_rcu_slot(mas->node, slot);
 
-			level--;
-			mas->node = mn;
-			slot = -1;
-			count = mt_slot_count(mas->node);
+		if (entry) {
+			mas->index = last_piv;
+			if (last_piv)
+				mas->index--;
+			goto done;
 		}
-
-		if (ma_is_root(mas->node))
-			goto no_entry;
 	}
 
-no_entry:
-	mas->node = MAS_NONE;
-	return mas->max;
-
-}
-
-// Next node entry or return 0 on none.
-static inline unsigned long
-mas_next_nentry (struct ma_state *mas, unsigned long max)
-{
-	unsigned long pivot = 0;
-	unsigned char slot = ma_get_slot(mas);
-	unsigned char count = mt_slot_count(mas->node);
-	void *entry;
+	last_piv = mas_next_entry(mas, max);
+	if (mas->node == MAS_NONE) {
+		goto not_found;
+	}
 
-	while (slot++ < count) {
-		pivot = mas->max;
-		if (slot < count - 1)
-			pivot = ma_get_pivot(mas->node, slot);
+	slot = ma_get_slot(mas);
+	entry = ma_get_rcu_slot(mas->node, slot);
+	if (!entry)
+		goto not_found;
 
-		if (pivot > max)
-			goto no_entry;
+	mas->index = ma_get_prev_pivot(mas, slot) + 1;
+done:
+	mas->last = last_piv;
+	ma_set_slot(mas, slot+1);
+	return entry;
 
-		if (slot != 0 && pivot == 0)
-			goto no_entry;
+not_found:
+	mas->node = MAS_NONE;
+	return NULL;
 
-		entry = ma_get_rcu_slot(mas->node, slot);
-		if (entry)
-			break;
-	}
-	ma_set_slot(mas, slot);
-	return pivot;
 
-no_entry:
-	return 0;
 }
-static inline unsigned long
-mas_next_entry (struct ma_state *mas, unsigned long max)
-{
-	unsigned long pivot = max;
-
-	if (mas->node == MAS_NONE)
-		return max;
 
-	if (!mt_is_leaf(mas->node))
-		pivot = mas_first_entry(mas, max);
-	else {
-		while (mas->node != MAS_NONE) {
-			unsigned char p_slot;
-			pivot = mas_next_nentry(mas, max);
-			if (pivot)
-				break;
-			p_slot = mt_parent_slot(mas->node);
-			ma_set_slot(mas, p_slot);
-			mas_next_node(mas, max);
-		}
-	}
-	return pivot;
+void *mt_find_after(struct ma_state *mas, unsigned long max)
+{
+	void *entry = mas_find(mas, max);
+	if (xa_is_zero(entry))
+		return NULL;
+	return entry;
 }
 static inline void ma_inactive_insert(struct ma_state *mas, void *entry);
 static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
@@ -2600,26 +2738,18 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 
 	
 
-	printk("\n%s:\n", __func__);
 	slot_cnt = 1 + ma_get_slot(mas);
 	ma_set_slot(mas, mt_parent_slot(mas->node));
-	printk("%p %u Replace %ld-%ld\n", mas->node, slot_cnt, mas->index, mas->last);
 	while (mas->node != MAS_NONE) {
 		last = mas->node;
 		mas_prev_node(mas, 0);
-		printk("%p\n", mas->node);
 		l_node_cnt++;
 	}
 	mas->node = last;
-	printk("%s: left node count = %ld\n", __func__, l_node_cnt);
-	printk("mas->node = %p\n", mas->node);
-	
 	_mas_walk(&r_mas);
 
-	printk("Starting right at %p[%u]\n", r_mas.node, ma_get_slot(&r_mas));
 	if (ma_get_slot(&r_mas) != MAPLE_NODE_SLOTS) {
 		unsigned long piv;
-		printk("Right node %p[%u]\n", r_node, r_slot);
 		r_slot_cnt += ma_data_end(&r_mas, mt_node_type(r_mas.node), &piv);
 		if (piv != ULONG_MAX) {
 			unsigned char p_slot = mt_parent_slot(r_mas.node);
@@ -2631,7 +2761,6 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 			while (r_mas.node != MAS_NONE) {
 				last = r_mas.node;
 				mas_next_node(&r_mas, ULONG_MAX);
-				printk("%p\n", r_mas.node);
 				r_node_cnt++;
 			}
 			r_mas.node = last;
@@ -2639,9 +2768,6 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 		if (r_mas.max < mas->last)
 			slot_cnt++;
 	}
-	printk("%s: right node count = %ld\n", __func__, r_node_cnt);
-	printk("%s: right slots %u\n", __func__, slot_cnt);
-
 
 	if (slot_cnt > mt_slot_count(mas->node))
 		nodes++;
@@ -2650,18 +2776,14 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 
 	node_cnt = 1; // Root node.
 	while (nodes) {
-		printk("Adding %ld\n", nodes);
 		node_cnt += nodes;
 		nodes /= 8;
 	}
 
-	printk("%s: Requesting %ld nodes\n", __func__, node_cnt + 1);
-
 	mas_node_cnt(mas, node_cnt + 1);
 	if (mas_is_err(mas))
 		return 0;
 
-	printk("\n\n%s: Here we go! %lu-%lu -> %p\n", __func__, mas->index, mas->last, new_entry);
 	// Create a new tree.
 	DEFINE_MTREE(new_tree);
 	MA_STATE(new_mas, &new_tree, 0, 0);
@@ -2679,7 +2801,6 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 		while (slot < mt_slot_count(mas->node)) {
 			if (new_mas.index == mas->index)
 				break;
-			printk("On %p[%u]\n", mas->node, slot);
 			new_mas.last = ma_get_safe_pivot(mas, slot);
 			if (!new_mas.last && slot)
 				break;
@@ -2689,7 +2810,6 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 
 			entry = ma_get_rcu_slot(mas->node, slot);
 			if (entry) {
-				printk("Add %lu-%lu -> %p\n", new_mas.index, new_mas.last, entry);
 				ma_inactive_insert(&new_mas, entry);
 				if (mas_is_err(&new_mas))
 					BUG_ON(1);
@@ -2703,17 +2823,13 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 			slot++;
 		}
 		ma_set_slot(mas, mt_parent_slot(mas->node));
-		printk("Going to next from %p's parent[%u]\n", mas->node, mt_parent_slot(mas->node));
 		mas_next_node(mas, mas->index);
-		printk("Returned %p\n", mas->node);
 		slot = 0;
 	} while (mas->node != MAS_NONE);
 
-	printk("Done left\n");
 	// Insert the new value.
 	new_mas.index = mas->index;
 	new_mas.last = mas->last;
-	printk("Add %lu-%lu -> %p\n", new_mas.index, new_mas.last, new_entry);
 	ma_inactive_insert(&new_mas, new_entry);
 	if (mas_is_err(&new_mas))
 		BUG_ON(1);
@@ -2728,7 +2844,6 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 	 */
 	if (!r_node) // No entry beyond new_mas.last
 		goto skip_right;
-	printk("Now the Right\n");
 
 	new_mas.index = new_mas.last + 1;
 	new_mas.last = ma_get_pivot(r_node, r_slot);
@@ -2736,27 +2851,20 @@ static inline int mas_replace_tree(struct ma_state *mas, void *new_entry)
 	ma_set_slot(mas, r_slot);
 	mas->node = r_node;
 	mas->min = new_mas.index;
-	printk("going from %p[%u] to the end\n", r_node, r_slot);
-	printk("Inserting %lu - %lu -> ?\n", new_mas.index, new_mas.last);
 	do {
 		new_mas.index = mas->min;
 		new_mas.last = mas->min;
 		slot = ma_get_slot(mas);
-		printk("do %p[%u]\n", mas->node, slot);
 		while (slot++ < mt_slot_count(mas->node)) {
 			new_mas.last = ma_get_safe_pivot(mas, slot);
 			if (!new_mas.last)
 				break;
-			printk("pivot %lu\n", new_mas.last);
 
 			entry = ma_get_rcu_slot(mas->node, slot);
 			if (entry) {
-				printk("inserting %lu-%lu->%p\n", new_mas.index, new_mas.last, entry);
 				ma_inactive_insert(&new_mas, entry);
-				if (mas_is_err(&new_mas)) {
-					printk("Error %p\n", new_mas.node);
+				if (mas_is_err(&new_mas))
 					BUG_ON(1);
-				}
 			}
 
 			new_mas.index = new_mas.last + 1;
@@ -2772,7 +2880,6 @@ skip_right:
 	_mt_replace(mas, false, false);
 	mas->node = 0;
 	mas->alloc = new_mas.alloc;
-	printk("Destroy %p\n", r_node);
 	ma_destroy_walk(r_node);
 
 	// Copy right side
@@ -3158,7 +3265,7 @@ int mtree_store_range(struct maple_tree *mt, unsigned long first,
 {
 	MA_STATE(mas, mt, first, last);
 
-	if (WARN_ON_ONCE(mt_is_reserved(entry)))
+	if (WARN_ON_ONCE(mt_is_advanced(entry)))
 		return -EINVAL;
 
 	if (first > last)
@@ -3171,12 +3278,9 @@ retry:
 		goto retry;
 
 	mtree_unlock(mas.tree);
-	if (mas_is_err(&mas)) {
-		printk("Error %p\n", mas.node);
+	if (mas_is_err(&mas))
 		return xa_err(mas.node);
-	}
 
-	printk("Return 0\n");
 	return 0;
 }
 EXPORT_SYMBOL(mtree_store_range);
@@ -3192,7 +3296,7 @@ int mtree_insert_range(struct maple_tree *mt, unsigned long first,
 {
 	MA_STATE(ms, mt, first, last);
 
-	if (WARN_ON_ONCE(mt_is_reserved(entry)))
+	if (WARN_ON_ONCE(mt_is_advanced(entry)))
 		return -EINVAL;
 
 	if (first > last)
diff --git a/lib/test_maple_tree.c b/lib/test_maple_tree.c
index 703051da9d07..60e749d8403b 100644
--- a/lib/test_maple_tree.c
+++ b/lib/test_maple_tree.c
@@ -359,6 +359,42 @@ static noinline void check_mid_split(struct maple_tree *mt)
 	check_insert(mt, 0, (void*) 0);
 	check_lb_not_empty(mt);
 }
+
+static noinline void check_find(struct maple_tree *mt)
+{
+	unsigned long val = 1;
+	unsigned long count = 20;
+	void *entry;
+	MA_STATE(mas, mt, 0, 0);
+
+	// Insert 0.
+	MT_BUG_ON(mt, mtree_insert_index(mt, 0, GFP_KERNEL));
+
+	for (int i = 0; i <= count; i++) {
+		if (val != 64)
+			MT_BUG_ON(mt, mtree_insert_index(mt, val, GFP_KERNEL));
+		else
+			MT_BUG_ON(mt, mtree_insert(mt, val,
+				XA_ZERO_ENTRY, GFP_KERNEL));
+
+		val <<= 2;
+	}
+
+	mas.node = MAS_START;
+	val = 0;
+	while ( (entry = mas_find(&mas, 268435456)) != NULL) {
+		if (val != 64)
+			MT_BUG_ON(mt, xa_mk_value(val) != entry);
+		else
+			MT_BUG_ON(mt, entry != XA_ZERO_ENTRY);
+
+		val <<= 2;
+		// For zero check.
+		if (!val)
+			val = 1;
+	}
+
+}
 static noinline void check_alloc_rev_range(struct maple_tree *mt)
 {
 	// cat /proc/self/maps|awk '{print $1}'| awk -F "-" '{printf "0x%s, 0x%s, ", $1, $2}'
@@ -816,6 +852,9 @@ static int maple_tree_seed(void)
 	check_lower_bound_split(&tree);
 	check_upper_bound_split(&tree);
 	check_mid_split(&tree);
+
+	check_find(&tree);
+
 	rcu_barrier();
 
 	printk("maple_tree: %u of %u tests passed\n", tests_passed, tests_run);
-- 
2.50.1