2 files changed, 36 insertions, 5 deletions

diff --git a/lib/test_xarray.c b/lib/test_xarray.c
index 6f09c845187e..4676c0a1eeca 100644
--- a/lib/test_xarray.c
+++ b/lib/test_xarray.c
@@ -633,6 +633,15 @@ static noinline void check_xa_alloc(void)
                                 GFP_KERNEL) != -ENOSPC);
         XA_BUG_ON(&xa0, id != 0xffffffffU);
         xa_destroy(&xa0);
+
+        id = 10;
+        XA_BUG_ON(&xa0, xa_alloc(&xa0, &id, 5, xa_mk_index(id),
+                                GFP_KERNEL) != -ENOSPC);
+        XA_BUG_ON(&xa0, xa_store_index(&xa0, 3, GFP_KERNEL) != 0);
+        XA_BUG_ON(&xa0, xa_alloc(&xa0, &id, 5, xa_mk_index(id),
+                                GFP_KERNEL) != -ENOSPC);
+        xa_erase_index(&xa0, 3);
+        XA_BUG_ON(&xa0, !xa_empty(&xa0));
 }
 
 static noinline void __check_store_iter(struct xarray *xa, unsigned long start,
@@ -822,10 +831,34 @@ static noinline void check_find_2(struct xarray *xa)
         xa_destroy(xa);
 }
 
+static noinline void check_find_3(struct xarray *xa)
+{
+        XA_STATE(xas, xa, 0);
+        unsigned long i, j, k;
+        void *entry;
+
+        for (i = 0; i < 100; i++) {
+                for (j = 0; j < 100; j++) {
+                        for (k = 0; k < 100; k++) {
+                                xas_set(&xas, j);
+                                xas_for_each_marked(&xas, entry, k, XA_MARK_0)
+                                        ;
+                                if (j > k)
+                                        XA_BUG_ON(xa,
+                                                xas.xa_node != XAS_RESTART);
+                        }
+                }
+                xa_store_index(xa, i, GFP_KERNEL);
+                xa_set_mark(xa, i, XA_MARK_0);
+        }
+        xa_destroy(xa);
+}
+
 static noinline void check_find(struct xarray *xa)
 {
         check_find_1(xa);
         check_find_2(xa);
+        check_find_3(xa);
         check_multi_find(xa);
         check_multi_find_2(xa);
 }
diff --git a/lib/xarray.c b/lib/xarray.c
index bbacca576593..5f3f9311de89 100644
--- a/lib/xarray.c
+++ b/lib/xarray.c
@@ -1131,7 +1131,7 @@ void *xas_find_marked(struct xa_state *xas, unsigned long max, xa_mark_t mark)
                 entry = xa_head(xas->xa);
                 xas->xa_node = NULL;
                 if (xas->xa_index > max_index(entry))
-                        goto bounds;
+                        goto out;
                 if (!xa_is_node(entry)) {
                         if (xa_marked(xas->xa, mark))
                                 return entry;
@@ -1180,11 +1180,9 @@ void *xas_find_marked(struct xa_state *xas, unsigned long max, xa_mark_t mark)
         }
 
 out:
-        if (!max)
+        if (xas->xa_index > max)
                 goto max;
-bounds:
-        xas->xa_node = XAS_BOUNDS;
-        return NULL;
+        return set_bounds(xas);
 max:
         xas->xa_node = XAS_RESTART;
         return NULL;