[PATCH 03/18] ioasid: Convert ioasid_idr to XArray

[Jacob Pan]

IDR is to be replaced by XArray, keep up with the changes.
XArray has internal locking for normal APIs used here, also removed
radix tree related preload.

Suggested-by: Ira Weiny <ira.weiny@xxxxxxxxx>
Signed-off-by: Jacob Pan <jacob.jun.pan@xxxxxxxxxxxxxxx>
---
 drivers/base/ioasid.c | 35 ++++++++++++++++-------------------
 1 file changed, 16 insertions(+), 19 deletions(-)

diff --git a/drivers/base/ioasid.c b/drivers/base/ioasid.c
index 294e856..35bdc680 100644
--- a/drivers/base/ioasid.c
+++ b/drivers/base/ioasid.c
@@ -4,7 +4,7 @@
  * subsets. Users create a subset with DECLARE_IOASID_SET, then allocate and
  * free IOASIDs with ioasid_alloc and ioasid_free.
  */
-#include <linux/idr.h>
+#include <linux/xarray.h>
 #include <linux/ioasid.h>
 #include <linux/slab.h>
 #include <linux/spinlock.h>
@@ -16,7 +16,7 @@ struct ioasid_data {
 	struct rcu_head rcu;
 };
 
-static DEFINE_IDR(ioasid_idr);
+static DEFINE_XARRAY_ALLOC(ioasid_xa);
 static DEFINE_MUTEX(ioasid_allocator_lock);
 static const struct ioasid_allocator *ioasid_allocator;
 
@@ -24,9 +24,9 @@ static const struct ioasid_allocator *ioasid_allocator;
 /**
  * ioasid_set_allocator - register a custom allocator
  *
- * Custom allocator take precedence over the default IDR based allocator.
+ * Custom allocator take precedence over the default xarray based allocator.
  * Private data associated with the ASID are managed by ASID common code
- * similar to IDR data.
+ * similar to data stored in xa.
  */
 int ioasid_set_allocator(struct ioasid_allocator *allocator)
 {
@@ -72,13 +72,12 @@ int ioasid_set_data(ioasid_t ioasid, void *data)
 	struct ioasid_data *ioasid_data;
 	int ret = 0;
 
-	idr_lock(&ioasid_idr);
-	ioasid_data = idr_find(&ioasid_idr, ioasid);
+	ioasid_data = xa_load(&ioasid_xa, ioasid);
 	if (ioasid_data)
 		ioasid_data->private = data;
 	else
 		ret = -ENOENT;
-	idr_unlock(&ioasid_idr);
+
 	/* getter may use the private data */
 	synchronize_rcu();
 
@@ -110,7 +109,7 @@ ioasid_t ioasid_alloc(struct ioasid_set *set, ioasid_t min, ioasid_t max,
 	data->set = set;
 	data->private = private;
 
-	/* Use custom allocator if available, otherwise default to IDR */
+	/* Use custom allocator if available, otherwise default to xarray */
 	if (ioasid_allocator) {
 		mutex_lock(&ioasid_allocator_lock);
 		id = ioasid_allocator->alloc(min, max, ioasid_allocator->pdata);
@@ -120,17 +119,18 @@ ioasid_t ioasid_alloc(struct ioasid_set *set, ioasid_t min, ioasid_t max,
 			goto exit_free;
 		}
 		/*
-		 * Use IDR to manage private data also sanitiy check custom
+		 * Use XA to manage private data also sanitiy check custom
 		 * allocator for duplicates.
 		 */
 		min = id;
 		max = id + 1;
 	}
-	idr_preload(GFP_KERNEL);
-	idr_lock(&ioasid_idr);
-	data->id = id = idr_alloc(&ioasid_idr, data, min, max, GFP_ATOMIC);
-	idr_unlock(&ioasid_idr);
-	idr_preload_end();
+
+	if (xa_alloc(&ioasid_xa, &id, data, XA_LIMIT(min, max), GFP_KERNEL)) {
+		pr_err("Failed to alloc ioasid from %d to %d\n", min, max);
+		goto exit_free;
+	}
+	data->id = id;
 
 exit_free:
 	if (id < 0 || id == INVALID_IOASID) {
@@ -160,10 +160,7 @@ int ioasid_free(ioasid_t ioasid)
 		return ret;
 	}
 
-	idr_lock(&ioasid_idr);
-	ioasid_data = idr_remove(&ioasid_idr, ioasid);
-	idr_unlock(&ioasid_idr);
-
+	ioasid_data = xa_erase(&ioasid_xa, ioasid);
 	if (ioasid_data)
 		kfree_rcu(ioasid_data, rcu);
 
@@ -192,7 +189,7 @@ void *ioasid_find(struct ioasid_set *set, ioasid_t ioasid,
 	struct ioasid_data *ioasid_data;
 
 	rcu_read_lock();
-	ioasid_data = idr_find(&ioasid_idr, ioasid);
+	ioasid_data = xa_load(&ioasid_xa, ioasid);
 	if (!ioasid_data) {
 		priv = ERR_PTR(-ENOENT);
 		goto unlock;
-- 
2.7.4