xarray: Add XArray load operation

The xa_load function brings with it a lot of infrastructure; xa_empty(),
xa_is_err(), and large chunks of the XArray advanced API that are used
to implement xa_load.

As the test-suite demonstrates, it is possible to use the XArray functions
on a radix tree.  The radix tree functions depend on the GFP flags being
stored in the root of the tree, so it's not possible to use the radix
tree functions on an XArray.

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

14 files changed, 661 insertions, 45 deletions

diff --git a/include/linux/xarray.h b/include/linux/xarray.h
index 52141dfc5a90..a0df8217068c 100644
--- a/include/linux/xarray.h
+++ b/include/linux/xarray.h
@@ -12,6 +12,8 @@
 #include <linux/bug.h>
 #include <linux/compiler.h>
 #include <linux/kconfig.h>
+#include <linux/kernel.h>
+#include <linux/rcupdate.h>
 #include <linux/spinlock.h>
 #include <linux/types.h>
 
@@ -30,6 +32,10 @@
  *
  * 0-62: Sibling entries
  * 256: Retry entry
+ *
+ * Errors are also represented as internal entries, but use the negative
+ * space (-4094 to -2).  They're never stored in the slots array; only
+ * returned by the normal API.
  */
 
 #define BITS_PER_XA_VALUE       (BITS_PER_LONG - 1)
@@ -156,6 +162,42 @@ static inline bool xa_is_internal(const void *entry)
 }
 
 /**
+ * xa_is_err() - Report whether an XArray operation returned an error
+ * @entry: Result from calling an XArray function
+ *
+ * If an XArray operation cannot complete an operation, it will return
+ * a special value indicating an error.  This function tells you
+ * whether an error occurred; xa_err() tells you which error occurred.
+ *
+ * Context: Any context.
+ * Return: %true if the entry indicates an error.
+ */
+static inline bool xa_is_err(const void *entry)
+{
+        return unlikely(xa_is_internal(entry));
+}
+
+/**
+ * xa_err() - Turn an XArray result into an errno.
+ * @entry: Result from calling an XArray function.
+ *
+ * If an XArray operation cannot complete an operation, it will return
+ * a special pointer value which encodes an errno.  This function extracts
+ * the errno from the pointer value, or returns 0 if the pointer does not
+ * represent an errno.
+ *
+ * Context: Any context.
+ * Return: A negative errno or 0.
+ */
+static inline int xa_err(void *entry)
+{
+        /* xa_to_internal() would not do sign extension. */
+        if (xa_is_err(entry))
+                return (long)entry >> 2;
+        return 0;
+}
+
+/**
  * struct xarray - The anchor of the XArray.
  * @xa_lock: Lock that protects the contents of the XArray.
  *
@@ -209,6 +251,7 @@ struct xarray {
 #define DEFINE_XARRAY(name) DEFINE_XARRAY_FLAGS(name, 0)
 
 void xa_init_flags(struct xarray *, gfp_t flags);
+void *xa_load(struct xarray *, unsigned long index);
 
 /**
  * xa_init() - Initialise an empty XArray.
@@ -223,6 +266,18 @@ static inline void xa_init(struct xarray *xa)
         xa_init_flags(xa, 0);
 }
 
+/**
+ * xa_empty() - Determine if an array has any present entries.
+ * @xa: XArray.
+ *
+ * Context: Any context.
+ * Return: %true if the array contains only NULL pointers.
+ */
+static inline bool xa_empty(const struct xarray *xa)
+{
+        return xa->xa_head == NULL;
+}
+
 #define xa_trylock(xa)          spin_trylock(&(xa)->xa_lock)
 #define xa_lock(xa)             spin_lock(&(xa)->xa_lock)
 #define xa_unlock(xa)           spin_unlock(&(xa)->xa_lock)
@@ -280,6 +335,65 @@ struct xa_node {
         };
 };
 
+void xa_dump(const struct xarray *);
+void xa_dump_node(const struct xa_node *);
+
+#ifdef XA_DEBUG
+#define XA_BUG_ON(xa, x) do {                                   \
+                if (x) {                                        \
+                        xa_dump(xa);                            \
+                        BUG();                                  \
+                }                                               \
+        } while (0)
+#define XA_NODE_BUG_ON(node, x) do {                            \
+                if (x) {                                        \
+                        if (node) xa_dump_node(node);           \
+                        BUG();                                  \
+                }                                               \
+        } while (0)
+#else
+#define XA_BUG_ON(xa, x)        do { } while (0)
+#define XA_NODE_BUG_ON(node, x) do { } while (0)
+#endif
+
+/* Private */
+static inline void *xa_head(const struct xarray *xa)
+{
+        return rcu_dereference_check(xa->xa_head,
+                                                lockdep_is_held(&xa->xa_lock));
+}
+
+/* Private */
+static inline void *xa_head_locked(const struct xarray *xa)
+{
+        return rcu_dereference_protected(xa->xa_head,
+                                                lockdep_is_held(&xa->xa_lock));
+}
+
+/* Private */
+static inline void *xa_entry(const struct xarray *xa,
+                                const struct xa_node *node, unsigned int offset)
+{
+        XA_NODE_BUG_ON(node, offset >= XA_CHUNK_SIZE);
+        return rcu_dereference_check(node->slots[offset],
+                                                lockdep_is_held(&xa->xa_lock));
+}
+
+/* Private */
+static inline void *xa_entry_locked(const struct xarray *xa,
+                                const struct xa_node *node, unsigned int offset)
+{
+        XA_NODE_BUG_ON(node, offset >= XA_CHUNK_SIZE);
+        return rcu_dereference_protected(node->slots[offset],
+                                                lockdep_is_held(&xa->xa_lock));
+}
+
+/* Private */
+static inline struct xa_node *xa_to_node(const void *entry)
+{
+        return (struct xa_node *)((unsigned long)entry - 2);
+}
+
 /* Private */
 static inline bool xa_is_node(const void *entry)
 {
@@ -312,4 +426,226 @@ static inline bool xa_is_sibling(const void *entry)
 
 #define XA_RETRY_ENTRY          xa_mk_internal(256)
 
+/**
+ * xa_is_retry() - Is the entry a retry entry?
+ * @entry: Entry retrieved from the XArray
+ *
+ * Return: %true if the entry is a retry entry.
+ */
+static inline bool xa_is_retry(const void *entry)
+{
+        return unlikely(entry == XA_RETRY_ENTRY);
+}
+
+/**
+ * typedef xa_update_node_t - A callback function from the XArray.
+ * @node: The node which is being processed
+ *
+ * This function is called every time the XArray updates the count of
+ * present and value entries in a node.  It allows advanced users to
+ * maintain the private_list in the node.
+ *
+ * Context: The xa_lock is held and interrupts may be disabled.
+ *          Implementations should not drop the xa_lock, nor re-enable
+ *          interrupts.
+ */
+typedef void (*xa_update_node_t)(struct xa_node *node);
+
+/*
+ * The xa_state is opaque to its users.  It contains various different pieces
+ * of state involved in the current operation on the XArray.  It should be
+ * declared on the stack and passed between the various internal routines.
+ * The various elements in it should not be accessed directly, but only
+ * through the provided accessor functions.  The below documentation is for
+ * the benefit of those working on the code, not for users of the XArray.
+ *
+ * @xa_node usually points to the xa_node containing the slot we're operating
+ * on (and @xa_offset is the offset in the slots array).  If there is a
+ * single entry in the array at index 0, there are no allocated xa_nodes to
+ * point to, and so we store %NULL in @xa_node.  @xa_node is set to
+ * the value %XAS_RESTART if the xa_state is not walked to the correct
+ * position in the tree of nodes for this operation.  If an error occurs
+ * during an operation, it is set to an %XAS_ERROR value.  If we run off the
+ * end of the allocated nodes, it is set to %XAS_BOUNDS.
+ */
+struct xa_state {
+        struct xarray *xa;
+        unsigned long xa_index;
+        unsigned char xa_shift;
+        unsigned char xa_sibs;
+        unsigned char xa_offset;
+        unsigned char xa_pad;           /* Helps gcc generate better code */
+        struct xa_node *xa_node;
+        struct xa_node *xa_alloc;
+        xa_update_node_t xa_update;
+};
+
+/*
+ * We encode errnos in the xas->xa_node.  If an error has happened, we need to
+ * drop the lock to fix it, and once we've done so the xa_state is invalid.
+ */
+#define XA_ERROR(errno) ((struct xa_node *)(((unsigned long)errno << 2) | 2UL))
+#define XAS_BOUNDS      ((struct xa_node *)1UL)
+#define XAS_RESTART     ((struct xa_node *)3UL)
+
+#define __XA_STATE(array, index, shift, sibs)  {        \
+        .xa = array,                                    \
+        .xa_index = index,                              \
+        .xa_shift = shift,                              \
+        .xa_sibs = sibs,                                \
+        .xa_offset = 0,                                 \
+        .xa_pad = 0,                                    \
+        .xa_node = XAS_RESTART,                         \
+        .xa_alloc = NULL,                               \
+        .xa_update = NULL                               \
+}
+
+/**
+ * XA_STATE() - Declare an XArray operation state.
+ * @name: Name of this operation state (usually xas).
+ * @array: Array to operate on.
+ * @index: Initial index of interest.
+ *
+ * Declare and initialise an xa_state on the stack.
+ */
+#define XA_STATE(name, array, index)                            \
+        struct xa_state name = __XA_STATE(array, index, 0, 0)
+
+/**
+ * XA_STATE_ORDER() - Declare an XArray operation state.
+ * @name: Name of this operation state (usually xas).
+ * @array: Array to operate on.
+ * @index: Initial index of interest.
+ * @order: Order of entry.
+ *
+ * Declare and initialise an xa_state on the stack.  This variant of
+ * XA_STATE() allows you to specify the 'order' of the element you
+ * want to operate on.`
+ */
+#define XA_STATE_ORDER(name, array, index, order)               \
+        struct xa_state name = __XA_STATE(array,                \
+                        (index >> order) << order,              \
+                        order - (order % XA_CHUNK_SHIFT),       \
+                        (1U << (order % XA_CHUNK_SHIFT)) - 1)
+
+#define xas_marked(xas, mark)   xa_marked((xas)->xa, (mark))
+#define xas_trylock(xas)        xa_trylock((xas)->xa)
+#define xas_lock(xas)           xa_lock((xas)->xa)
+#define xas_unlock(xas)         xa_unlock((xas)->xa)
+#define xas_lock_bh(xas)        xa_lock_bh((xas)->xa)
+#define xas_unlock_bh(xas)      xa_unlock_bh((xas)->xa)
+#define xas_lock_irq(xas)       xa_lock_irq((xas)->xa)
+#define xas_unlock_irq(xas)     xa_unlock_irq((xas)->xa)
+#define xas_lock_irqsave(xas, flags) \
+                                xa_lock_irqsave((xas)->xa, flags)
+#define xas_unlock_irqrestore(xas, flags) \
+                                xa_unlock_irqrestore((xas)->xa, flags)
+
+/**
+ * xas_error() - Return an errno stored in the xa_state.
+ * @xas: XArray operation state.
+ *
+ * Return: 0 if no error has been noted.  A negative errno if one has.
+ */
+static inline int xas_error(const struct xa_state *xas)
+{
+        return xa_err(xas->xa_node);
+}
+
+/**
+ * xas_set_err() - Note an error in the xa_state.
+ * @xas: XArray operation state.
+ * @err: Negative error number.
+ *
+ * Only call this function with a negative @err; zero or positive errors
+ * will probably not behave the way you think they should.  If you want
+ * to clear the error from an xa_state, use xas_reset().
+ */
+static inline void xas_set_err(struct xa_state *xas, long err)
+{
+        xas->xa_node = XA_ERROR(err);
+}
+
+/**
+ * xas_invalid() - Is the xas in a retry or error state?
+ * @xas: XArray operation state.
+ *
+ * Return: %true if the xas cannot be used for operations.
+ */
+static inline bool xas_invalid(const struct xa_state *xas)
+{
+        return (unsigned long)xas->xa_node & 3;
+}
+
+/**
+ * xas_valid() - Is the xas a valid cursor into the array?
+ * @xas: XArray operation state.
+ *
+ * Return: %true if the xas can be used for operations.
+ */
+static inline bool xas_valid(const struct xa_state *xas)
+{
+        return !xas_invalid(xas);
+}
+
+/**
+ * xas_reset() - Reset an XArray operation state.
+ * @xas: XArray operation state.
+ *
+ * Resets the error or walk state of the @xas so future walks of the
+ * array will start from the root.  Use this if you have dropped the
+ * xarray lock and want to reuse the xa_state.
+ *
+ * Context: Any context.
+ */
+static inline void xas_reset(struct xa_state *xas)
+{
+        xas->xa_node = XAS_RESTART;
+}
+
+/**
+ * xas_retry() - Retry the operation if appropriate.
+ * @xas: XArray operation state.
+ * @entry: Entry from xarray.
+ *
+ * The advanced functions may sometimes return an internal entry, such as
+ * a retry entry or a zero entry.  This function sets up the @xas to restart
+ * the walk from the head of the array if needed.
+ *
+ * Context: Any context.
+ * Return: true if the operation needs to be retried.
+ */
+static inline bool xas_retry(struct xa_state *xas, const void *entry)
+{
+        if (!xa_is_retry(entry))
+                return false;
+        xas_reset(xas);
+        return true;
+}
+
+void *xas_load(struct xa_state *);
+
+/**
+ * xas_reload() - Refetch an entry from the xarray.
+ * @xas: XArray operation state.
+ *
+ * Use this function to check that a previously loaded entry still has
+ * the same value.  This is useful for the lockless pagecache lookup where
+ * we walk the array with only the RCU lock to protect us, lock the page,
+ * then check that the page hasn't moved since we looked it up.
+ *
+ * The caller guarantees that @xas is still valid.  If it may be in an
+ * error or restart state, call xas_load() instead.
+ *
+ * Return: The entry at this location in the xarray.
+ */
+static inline void *xas_reload(struct xa_state *xas)
+{
+        struct xa_node *node = xas->xa_node;
+
+        if (node)
+                return xa_entry(xas->xa, node, xas->xa_offset);
+        return xa_head(xas->xa);
+}
+
 #endif /* _LINUX_XARRAY_H */
diff --git a/lib/Kconfig.debug b/lib/Kconfig.debug
index 4966c4fbe7f7..091155e12422 100644
--- a/lib/Kconfig.debug
+++ b/lib/Kconfig.debug
@@ -1813,6 +1813,9 @@ config TEST_BITFIELD
 config TEST_UUID
         tristate "Test functions located in the uuid module at runtime"
 
+config TEST_XARRAY
+        tristate "Test the XArray code at runtime"
+
 config TEST_OVERFLOW
         tristate "Test check_*_overflow() functions at runtime"
 
diff --git a/lib/Makefile b/lib/Makefile
index 057385f1f145..e6f809556af5 100644
--- a/lib/Makefile
+++ b/lib/Makefile
@@ -68,6 +68,7 @@ obj-$(CONFIG_TEST_PRINTF) += test_printf.o
 obj-$(CONFIG_TEST_BITMAP) += test_bitmap.o
 obj-$(CONFIG_TEST_BITFIELD) += test_bitfield.o
 obj-$(CONFIG_TEST_UUID) += test_uuid.o
+obj-$(CONFIG_TEST_XARRAY) += test_xarray.o
 obj-$(CONFIG_TEST_PARMAN) += test_parman.o
 obj-$(CONFIG_TEST_KMOD) += test_kmod.o
 obj-$(CONFIG_TEST_DEBUG_VIRTUAL) += test_debug_virtual.o
diff --git a/lib/radix-tree.c b/lib/radix-tree.c
index 8a568cea1237..b8e961428484 100644
--- a/lib/radix-tree.c
+++ b/lib/radix-tree.c
@@ -256,49 +256,6 @@ static unsigned long next_index(unsigned long index,
 }
 
 #ifndef __KERNEL__
-static void dump_node(struct radix_tree_node *node, unsigned long index)
-{
-        unsigned long i;
-
-        pr_debug("radix node: %p offset %d indices %lu-%lu parent %p tags %lx %lx %lx shift %d count %d nr_values %d\n",
-                node, node->offset, index, index | node_maxindex(node),
-                node->parent,
-                node->tags[0][0], node->tags[1][0], node->tags[2][0],
-                node->shift, node->count, node->nr_values);
-
-        for (i = 0; i < RADIX_TREE_MAP_SIZE; i++) {
-                unsigned long first = index | (i << node->shift);
-                unsigned long last = first | ((1UL << node->shift) - 1);
-                void *entry = node->slots[i];
-                if (!entry)
-                        continue;
-                if (entry == RADIX_TREE_RETRY) {
-                        pr_debug("radix retry offset %ld indices %lu-%lu parent %p\n",
-                                        i, first, last, node);
-                } else if (!radix_tree_is_internal_node(entry)) {
-                        pr_debug("radix entry %p offset %ld indices %lu-%lu parent %p\n",
-                                        entry, i, first, last, node);
-                } else if (xa_is_sibling(entry)) {
-                        pr_debug("radix sblng %p offset %ld indices %lu-%lu parent %p val %p\n",
-                                        entry, i, first, last, node,
-                                        node->slots[xa_to_sibling(entry)]);
-                } else {
-                        dump_node(entry_to_node(entry), first);
-                }
-        }
-}
-
-/* For debug */
-static void radix_tree_dump(struct radix_tree_root *root)
-{
-        pr_debug("radix root: %p xa_head %p tags %x\n",
-                        root, root->xa_head,
-                        root->xa_flags >> ROOT_TAG_SHIFT);
-        if (!radix_tree_is_internal_node(root->xa_head))
-                return;
-        dump_node(entry_to_node(root->xa_head), 0);
-}
-
 static void dump_ida_node(void *entry, unsigned long index)
 {
         unsigned long i;
diff --git a/lib/test_xarray.c b/lib/test_xarray.c
new file mode 100644
index 000000000000..a7248b87617f
--- /dev/null
+++ b/lib/test_xarray.c
@@ -0,0 +1,87 @@
+// SPDX-License-Identifier: GPL-2.0+
+/*
+ * test_xarray.c: Test the XArray API
+ * Copyright (c) 2017-2018 Microsoft Corporation
+ * Author: Matthew Wilcox <willy@infradead.org>
+ */
+
+#include <linux/xarray.h>
+#include <linux/module.h>
+
+static unsigned int tests_run;
+static unsigned int tests_passed;
+
+#ifndef XA_DEBUG
+# ifdef __KERNEL__
+void xa_dump(const struct xarray *xa) { }
+# endif
+#undef XA_BUG_ON
+#define XA_BUG_ON(xa, x) do {                                   \
+        tests_run++;                                            \
+        if (x) {                                                \
+                printk("BUG at %s:%d\n", __func__, __LINE__);   \
+                xa_dump(xa);                                    \
+                dump_stack();                                   \
+        } else {                                                \
+                tests_passed++;                                 \
+        }                                                       \
+} while (0)
+#endif
+
+static void *xa_store_index(struct xarray *xa, unsigned long index, gfp_t gfp)
+{
+        radix_tree_insert(xa, index, xa_mk_value(index));
+        return NULL;
+}
+
+static void xa_erase_index(struct xarray *xa, unsigned long index)
+{
+        radix_tree_delete(xa, index);
+}
+
+static noinline void check_xa_load(struct xarray *xa)
+{
+        unsigned long i, j;
+
+        for (i = 0; i < 1024; i++) {
+                for (j = 0; j < 1024; j++) {
+                        void *entry = xa_load(xa, j);
+                        if (j < i)
+                                XA_BUG_ON(xa, xa_to_value(entry) != j);
+                        else
+                                XA_BUG_ON(xa, entry);
+                }
+                XA_BUG_ON(xa, xa_store_index(xa, i, GFP_KERNEL) != NULL);
+        }
+
+        for (i = 0; i < 1024; i++) {
+                for (j = 0; j < 1024; j++) {
+                        void *entry = xa_load(xa, j);
+                        if (j >= i)
+                                XA_BUG_ON(xa, xa_to_value(entry) != j);
+                        else
+                                XA_BUG_ON(xa, entry);
+                }
+                xa_erase_index(xa, i);
+        }
+        XA_BUG_ON(xa, !xa_empty(xa));
+}
+
+static RADIX_TREE(array, GFP_KERNEL);
+
+static int xarray_checks(void)
+{
+        check_xa_load(&array);
+
+        printk("XArray: %u of %u tests passed\n", tests_passed, tests_run);
+        return (tests_run == tests_passed) ? 0 : -EINVAL;
+}
+
+static void xarray_exit(void)
+{
+}
+
+module_init(xarray_checks);
+module_exit(xarray_exit);
+MODULE_AUTHOR("Matthew Wilcox <willy@infradead.org>");
+MODULE_LICENSE("GPL");
diff --git a/lib/xarray.c b/lib/xarray.c
index 862f4c64c754..19cfcbc69a68 100644
--- a/lib/xarray.c
+++ b/lib/xarray.c
@@ -24,6 +24,100 @@
  * @entry refers to something stored in a slot in the xarray
  */
 
+/* extracts the offset within this node from the index */
+static unsigned int get_offset(unsigned long index, struct xa_node *node)
+{
+        return (index >> node->shift) & XA_CHUNK_MASK;
+}
+
+/* move the index either forwards (find) or backwards (sibling slot) */
+static void xas_move_index(struct xa_state *xas, unsigned long offset)
+{
+        unsigned int shift = xas->xa_node->shift;
+        xas->xa_index &= ~XA_CHUNK_MASK << shift;
+        xas->xa_index += offset << shift;
+}
+
+static void *set_bounds(struct xa_state *xas)
+{
+        xas->xa_node = XAS_BOUNDS;
+        return NULL;
+}
+
+/*
+ * Starts a walk.  If the @xas is already valid, we assume that it's on
+ * the right path and just return where we've got to.  If we're in an
+ * error state, return NULL.  If the index is outside the current scope
+ * of the xarray, return NULL without changing @xas->xa_node.  Otherwise
+ * set @xas->xa_node to NULL and return the current head of the array.
+ */
+static void *xas_start(struct xa_state *xas)
+{
+        void *entry;
+
+        if (xas_valid(xas))
+                return xas_reload(xas);
+        if (xas_error(xas))
+                return NULL;
+
+        entry = xa_head(xas->xa);
+        if (!xa_is_node(entry)) {
+                if (xas->xa_index)
+                        return set_bounds(xas);
+        } else {
+                if ((xas->xa_index >> xa_to_node(entry)->shift) > XA_CHUNK_MASK)
+                        return set_bounds(xas);
+        }
+
+        xas->xa_node = NULL;
+        return entry;
+}
+
+static void *xas_descend(struct xa_state *xas, struct xa_node *node)
+{
+        unsigned int offset = get_offset(xas->xa_index, node);
+        void *entry = xa_entry(xas->xa, node, offset);
+
+        xas->xa_node = node;
+        if (xa_is_sibling(entry)) {
+                offset = xa_to_sibling(entry);
+                entry = xa_entry(xas->xa, node, offset);
+        }
+
+        xas->xa_offset = offset;
+        return entry;
+}
+
+/**
+ * xas_load() - Load an entry from the XArray (advanced).
+ * @xas: XArray operation state.
+ *
+ * Usually walks the @xas to the appropriate state to load the entry
+ * stored at xa_index.  However, it will do nothing and return %NULL if
+ * @xas is in an error state.  xas_load() will never expand the tree.
+ *
+ * If the xa_state is set up to operate on a multi-index entry, xas_load()
+ * may return %NULL or an internal entry, even if there are entries
+ * present within the range specified by @xas.
+ *
+ * Context: Any context.  The caller should hold the xa_lock or the RCU lock.
+ * Return: Usually an entry in the XArray, but see description for exceptions.
+ */
+void *xas_load(struct xa_state *xas)
+{
+        void *entry = xas_start(xas);
+
+        while (xa_is_node(entry)) {
+                struct xa_node *node = xa_to_node(entry);
+
+                if (xas->xa_shift > node->shift)
+                        break;
+                entry = xas_descend(xas, node);
+        }
+        return entry;
+}
+EXPORT_SYMBOL_GPL(xas_load);
+
 /**
  * xa_init_flags() - Initialise an empty XArray with flags.
  * @xa: XArray.
@@ -42,3 +136,104 @@ void xa_init_flags(struct xarray *xa, gfp_t flags)
         xa->xa_head = NULL;
 }
 EXPORT_SYMBOL(xa_init_flags);
+
+/**
+ * xa_load() - Load an entry from an XArray.
+ * @xa: XArray.
+ * @index: index into array.
+ *
+ * Context: Any context.  Takes and releases the RCU lock.
+ * Return: The entry at @index in @xa.
+ */
+void *xa_load(struct xarray *xa, unsigned long index)
+{
+        XA_STATE(xas, xa, index);
+        void *entry;
+
+        rcu_read_lock();
+        do {
+                entry = xas_load(&xas);
+        } while (xas_retry(&xas, entry));
+        rcu_read_unlock();
+
+        return entry;
+}
+EXPORT_SYMBOL(xa_load);
+
+#ifdef XA_DEBUG
+void xa_dump_node(const struct xa_node *node)
+{
+        unsigned i, j;
+
+        if (!node)
+                return;
+        if ((unsigned long)node & 3) {
+                pr_cont("node %px\n", node);
+                return;
+        }
+
+        pr_cont("node %px %s %d parent %px shift %d count %d values %d "
+                "array %px list %px %px marks",
+                node, node->parent ? "offset" : "max", node->offset,
+                node->parent, node->shift, node->count, node->nr_values,
+                node->array, node->private_list.prev, node->private_list.next);
+        for (i = 0; i < XA_MAX_MARKS; i++)
+                for (j = 0; j < XA_MARK_LONGS; j++)
+                        pr_cont(" %lx", node->marks[i][j]);
+        pr_cont("\n");
+}
+
+void xa_dump_index(unsigned long index, unsigned int shift)
+{
+        if (!shift)
+                pr_info("%lu: ", index);
+        else if (shift >= BITS_PER_LONG)
+                pr_info("0-%lu: ", ~0UL);
+        else
+                pr_info("%lu-%lu: ", index, index | ((1UL << shift) - 1));
+}
+
+void xa_dump_entry(const void *entry, unsigned long index, unsigned long shift)
+{
+        if (!entry)
+                return;
+
+        xa_dump_index(index, shift);
+
+        if (xa_is_node(entry)) {
+                if (shift == 0) {
+                        pr_cont("%px\n", entry);
+                } else {
+                        unsigned long i;
+                        struct xa_node *node = xa_to_node(entry);
+                        xa_dump_node(node);
+                        for (i = 0; i < XA_CHUNK_SIZE; i++)
+                                xa_dump_entry(node->slots[i],
+                                      index + (i << node->shift), node->shift);
+                }
+        } else if (xa_is_value(entry))
+                pr_cont("value %ld (0x%lx) [%px]\n", xa_to_value(entry),
+                                                xa_to_value(entry), entry);
+        else if (!xa_is_internal(entry))
+                pr_cont("%px\n", entry);
+        else if (xa_is_retry(entry))
+                pr_cont("retry (%ld)\n", xa_to_internal(entry));
+        else if (xa_is_sibling(entry))
+                pr_cont("sibling (slot %ld)\n", xa_to_sibling(entry));
+        else
+                pr_cont("UNKNOWN ENTRY (%px)\n", entry);
+}
+
+void xa_dump(const struct xarray *xa)
+{
+        void *entry = xa->xa_head;
+        unsigned int shift = 0;
+
+        pr_info("xarray: %px head %px flags %x marks %d %d %d\n", xa, entry,
+                        xa->xa_flags, radix_tree_tagged(xa, 0),
+                        radix_tree_tagged(xa, 1), radix_tree_tagged(xa, 2));
+        if (xa_is_node(entry))
+                shift = xa_to_node(entry)->shift + XA_CHUNK_SHIFT;
+        xa_dump_entry(entry, 0, shift);
+}
+#endif
diff --git a/tools/include/linux/kernel.h b/tools/include/linux/kernel.h
index 0ad884452c5c..6935ef94e77a 100644
--- a/tools/include/linux/kernel.h
+++ b/tools/include/linux/kernel.h
@@ -70,6 +70,7 @@
 #define BUG_ON(cond) assert(!(cond))
 #endif
 #endif
+#define BUG()   BUG_ON(1)
 
 #if __BYTE_ORDER == __BIG_ENDIAN
 #define cpu_to_le16 bswap_16
diff --git a/tools/testing/radix-tree/.gitignore b/tools/testing/radix-tree/.gitignore
index d4706c0ffceb..3834899b6693 100644
--- a/tools/testing/radix-tree/.gitignore
+++ b/tools/testing/radix-tree/.gitignore
@@ -4,3 +4,4 @@ idr-test
 main
 multiorder
 radix-tree.c
+xarray
diff --git a/tools/testing/radix-tree/Makefile b/tools/testing/radix-tree/Makefile
index c0cf1c79efd5..1379f1d78d0b 100644
--- a/tools/testing/radix-tree/Makefile
+++ b/tools/testing/radix-tree/Makefile
@@ -4,7 +4,7 @@ CFLAGS += -I. -I../../include -g -Og -Wall -D_LGPL_SOURCE -fsanitize=address \
           -fsanitize=undefined
 LDFLAGS += -fsanitize=address -fsanitize=undefined
 LDLIBS+= -lpthread -lurcu
-TARGETS = main idr-test multiorder
+TARGETS = main idr-test multiorder xarray
 CORE_OFILES := xarray.o radix-tree.o idr.o linux.o test.o find_bit.o
 OFILES = main.o $(CORE_OFILES) regression1.o regression2.o regression3.o \
          tag_check.o multiorder.o idr-test.o iteration_check.o benchmark.o
@@ -25,6 +25,8 @@ main:	$(OFILES)
 idr-test.o: ../../../lib/test_ida.c
 idr-test: idr-test.o $(CORE_OFILES)
 
+xarray: $(CORE_OFILES)
+
 multiorder: multiorder.o $(CORE_OFILES)
 
 clean:
@@ -45,7 +47,7 @@ radix-tree.c: ../../../lib/radix-tree.c
 idr.c: ../../../lib/idr.c
         sed -e 's/^static //' -e 's/__always_inline //' -e 's/inline //' < $< > $@
 
-xarray.o: ../../../lib/xarray.c
+xarray.o: ../../../lib/xarray.c ../../../lib/test_xarray.c
 
 generated/map-shift.h:
         @if ! grep -qws $(SHIFT) generated/map-shift.h; then            \
diff --git a/tools/testing/radix-tree/linux/kernel.h b/tools/testing/radix-tree/linux/kernel.h
index 426f32f28547..5d06ac75a14d 100644
--- a/tools/testing/radix-tree/linux/kernel.h
+++ b/tools/testing/radix-tree/linux/kernel.h
@@ -14,6 +14,7 @@
 #include "../../../include/linux/kconfig.h"
 
 #define printk printf
+#define pr_info printk
 #define pr_debug printk
 #define pr_cont printk
 
diff --git a/tools/testing/radix-tree/linux/rcupdate.h b/tools/testing/radix-tree/linux/rcupdate.h
index 73ed33658203..fd280b070fdb 100644
--- a/tools/testing/radix-tree/linux/rcupdate.h
+++ b/tools/testing/radix-tree/linux/rcupdate.h
@@ -6,5 +6,7 @@
 
 #define rcu_dereference_raw(p) rcu_dereference(p)
 #define rcu_dereference_protected(p, cond) rcu_dereference(p)
+#define rcu_dereference_check(p, cond) rcu_dereference(p)
+#define RCU_INIT_POINTER(p, v)  (p) = (v)
 
 #endif
diff --git a/tools/testing/radix-tree/main.c b/tools/testing/radix-tree/main.c
index b741686e53d6..09deaf4f0959 100644
--- a/tools/testing/radix-tree/main.c
+++ b/tools/testing/radix-tree/main.c
@@ -365,6 +365,7 @@ int main(int argc, char **argv)
         rcu_register_thread();
         radix_tree_init();
 
+        xarray_tests();
         regression1_test();
         regression2_test();
         regression3_test();
diff --git a/tools/testing/radix-tree/test.h b/tools/testing/radix-tree/test.h
index 92d901eacf49..e3dc7a16f09b 100644
--- a/tools/testing/radix-tree/test.h
+++ b/tools/testing/radix-tree/test.h
@@ -34,6 +34,7 @@ int tag_tagged_items(struct radix_tree_root *, pthread_mutex_t *,
                         unsigned iftag, unsigned thentag);
 unsigned long find_item(struct radix_tree_root *, void *item);
 
+void xarray_tests(void);
 void tag_check(void);
 void multiorder_checks(void);
 void iteration_test(unsigned order, unsigned duration);
diff --git a/tools/testing/radix-tree/xarray.c b/tools/testing/radix-tree/xarray.c
index 9bbd667172a7..e61e43efe463 100644
--- a/tools/testing/radix-tree/xarray.c
+++ b/tools/testing/radix-tree/xarray.c
@@ -4,4 +4,32 @@
  * Copyright (c) 2018 Matthew Wilcox <willy@infradead.org>
  */
 
+#define XA_DEBUG
+#include "test.h"
+
+#define module_init(x)
+#define module_exit(x)
+#define MODULE_AUTHOR(x)
+#define MODULE_LICENSE(x)
+#define dump_stack()    assert(0)
+
 #include "../../../lib/xarray.c"
+#undef XA_DEBUG
+#include "../../../lib/test_xarray.c"
+
+void xarray_tests(void)
+{
+        xarray_checks();
+        xarray_exit();
+}
+
+int __weak main(void)
+{
+        radix_tree_init();
+        xarray_tests();
+        radix_tree_cpu_dead(1);
+        rcu_barrier();
+        if (nr_allocated)
+                printf("nr_allocated = %d\n", nr_allocated);
+        return 0;
+}