mas_find implementation

Fixes for prev/next node as well

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

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 5abe5e5eec30dc646823be14622707ad609b22af..cc1b5b2b80819f3e1850134cf11fa480ffdf5a81 100644 (file)
--- 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 703051da9d07d26bb92a32682a775bdb12bde26f..60e749d8403bb1c6bf3dab8617d84f360b7aca6c 100644 (file)
--- 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);