XArray: Fix xa_erase of 2-byte aligned entries

xas_store() was interpreting the entry it found in the array as a node
entry if the bottom two bits had value 2.  That's only true if either
the entry is in the root node or in a non-leaf node.

Signed-off-by: Matthew Wilcox <willy@infradead.org>

2 files changed, 16 insertions, 2 deletions

diff --git a/lib/test_xarray.c b/lib/test_xarray.c
index 52f8ecff8c0c..bc202d468a6b 100644
--- a/lib/test_xarray.c
+++ b/lib/test_xarray.c
@@ -1355,6 +1355,20 @@ static void check_align_1(struct xarray *xa, char *name)
         xa_destroy(xa);
 }
 
+static void check_align_2(struct xarray *xa, char *name)
+{
+        int i;
+
+        XA_BUG_ON(xa, !xa_empty(xa));
+
+        for (i = 0; i < 8; i++) {
+                XA_BUG_ON(xa, xa_store(xa, 0, name + i, GFP_KERNEL) != NULL);
+                xa_erase(xa, 0);
+        }
+
+        XA_BUG_ON(xa, !xa_empty(xa));
+}
+
 static noinline void check_align(struct xarray *xa)
 {
         char name[] = "Motorola 68000";
@@ -1363,7 +1377,7 @@ static noinline void check_align(struct xarray *xa)
         check_align_1(xa, name + 1);
         check_align_1(xa, name + 2);
         check_align_1(xa, name + 3);
-//      check_align_2(xa, name);
+        check_align_2(xa, name);
 }
 
 static LIST_HEAD(shadow_nodes);
diff --git a/lib/xarray.c b/lib/xarray.c
index 3f10198f00b7..2cc3798672f7 100644
--- a/lib/xarray.c
+++ b/lib/xarray.c
@@ -800,7 +800,7 @@ void *xas_store(struct xa_state *xas, void *entry)
                  * entry is set to NULL.
                  */
                 rcu_assign_pointer(*slot, entry);
-                if (xa_is_node(next))
+                if (xa_is_node(next) && (!node || node->shift))
                         xas_free_nodes(xas, xa_to_node(next));
                 if (!node)
                         break;