maple_tree: Rework mas_next_nentry

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

diff --git a/lib/maple_tree.c b/lib/maple_tree.c
index 38b7bdee6124c8e6548ed37d715fa1160aa3c37c..da299ebde6164c8e040c4afe8ba8de11ecd24735 100644 (file)
--- a/lib/maple_tree.c
+++ b/lib/maple_tree.c
@@ -3738,34 +3738,56 @@ static inline bool mas_next_nentry(struct ma_state *mas, unsigned long max,
                unsigned long *range_start)
 {
        enum maple_type type = mte_node_type(mas->node);
-       unsigned long pivot = mas->min;
+       struct maple_node *node = mas_mn(mas);
+       unsigned long pivot;
        unsigned long r_start;
        unsigned char count, offset = mas->offset;
-       unsigned long *pivots = ma_pivots(mas_mn(mas), type);
-       void **slots = ma_slots(mas_mn(mas), type);
+       unsigned long *pivots = ma_pivots(node, type);
+       void **slots;
 
-       count = mt_slots[type];
        r_start = mas_safe_min(mas, pivots, offset);
-       while (offset < count) {
-               pivot = _mas_safe_pivot(mas, pivots, offset, type);
-               if (!pivot && offset)
+       if (r_start > max) {
+               mas->index = max;
+               goto no_entry;
+       }
+
+       if (type == maple_arange_64)
+               count = ma_meta_end(node, type);
+       else
+               count = mt_pivots[type];
+
+
+       if (!ma_is_leaf(type)) {
+               if (count < offset)
                        goto no_entry;
 
-               if (r_start > max) {
-                       mas->index = max;
+               pivot = _mas_safe_pivot(mas, pivots, offset, type);
+               goto found;
+       }
+
+       slots = ma_slots(node, type);
+       while (offset < count) {
+               pivot = pivots[offset];
+               if (!pivot)
                        goto no_entry;
-               }
 
                if (mas_slot(mas, slots, offset))
                        goto found;
 
-               /* Ran over the limit, this is was the last slot to try */
-               if (pivot >= max)
-                       goto no_entry;
-
                r_start = pivot + 1;
+               if (r_start > max) {
+                       mas->index = max;
+                       goto no_entry;
+               }
                offset++;
        }
+       pivot = _mas_safe_pivot(mas, pivots, offset, type);
+       if (!pivot)
+               goto no_entry;
+
+       if (mas_slot(mas, slots, offset))
+               goto found;
+
 
 no_entry:
        *range_start = r_start;