[PATCH 03/10] lib: idr: implement Linux idr_alloc/_u32 API

[Ahmad Fatoum <a.fatoum@barebox.org>]

idr_alloc_one is a barebox invention that was adequate for the use case
in the ARM SCMI driver, but is not enough for incoming 9PFS support.

Reimplement it int terms of idr_alloc like in Linux.

Signed-off-by: Ahmad Fatoum <a.fatoum@barebox.org>
---
 drivers/firmware/arm_scmi/bus.c    |   5 +-
 drivers/firmware/arm_scmi/driver.c |  10 ++-
 include/linux/idr.h                |   5 +-
 lib/idr.c                          | 116 ++++++++++++++++++++++++-----
 test/self/idr.c                    |   5 ++
 5 files changed, 115 insertions(+), 26 deletions(-)

diff --git a/drivers/firmware/arm_scmi/bus.c b/drivers/firmware/arm_scmi/bus.c
index 0d7c404ae5ed..cf58b7088281 100644
--- a/drivers/firmware/arm_scmi/bus.c
+++ b/drivers/firmware/arm_scmi/bus.c
@@ -114,8 +114,9 @@ static int scmi_protocol_device_request(const struct scmi_device_id *id_table)
 		}
 		INIT_LIST_HEAD(phead);
 
-		ret = idr_alloc_one(&scmi_requested_devices, (void *)phead,
-				id_table->protocol_id);
+		ret = idr_alloc(&scmi_requested_devices, (void *)phead,
+				id_table->protocol_id,
+				id_table->protocol_id + 1, GFP_KERNEL);
 		if (ret != id_table->protocol_id) {
 			pr_err("Failed to save SCMI device - ret:%d\n", ret);
 			kfree(rdev);
diff --git a/drivers/firmware/arm_scmi/driver.c b/drivers/firmware/arm_scmi/driver.c
index e602f7a4404a..0048cc012223 100644
--- a/drivers/firmware/arm_scmi/driver.c
+++ b/drivers/firmware/arm_scmi/driver.c
@@ -173,7 +173,8 @@ int scmi_protocol_register(const struct scmi_protocol *proto)
 	}
 
 	spin_lock(&protocol_lock);
-	ret = idr_alloc_one(&scmi_protocols, (void *)proto, proto->id);
+	ret = idr_alloc(&scmi_protocols, (void *)proto, proto->id, proto->id + 1,
+			GFP_KERNEL);
 	spin_unlock(&protocol_lock);
 	if (ret != proto->id) {
 		pr_err("unable to allocate SCMI idr slot for 0x%x - err %d\n",
@@ -1090,7 +1091,7 @@ scmi_alloc_init_protocol_instance(struct scmi_info *info,
 	if (ret)
 		goto clean;
 
-	ret = idr_alloc_one(&info->protocols, pi, proto->id);
+	ret = idr_alloc(&info->protocols, pi, proto->id, proto->id + 1, GFP_KERNEL);
 	if (ret != proto->id)
 		goto clean;
 
@@ -1444,7 +1445,7 @@ static int scmi_chan_setup(struct scmi_info *info, struct device_node *of_node,
 	}
 
 idr_alloc:
-	ret = idr_alloc_one(idr, cinfo, prot_id);
+	ret = idr_alloc(idr, cinfo, prot_id, prot_id + 1, GFP_KERNEL);
 	if (ret != prot_id) {
 		dev_err(info->dev,
 			"unable to allocate SCMI idr slot err %d\n", ret);
@@ -1693,7 +1694,8 @@ static int scmi_probe(struct device *dev)
 		 * Save this valid DT protocol descriptor amongst
 		 * @active_protocols for this SCMI instance/
 		 */
-		ret = idr_alloc_one(&info->active_protocols, child, prot_id);
+		ret = idr_alloc(&info->active_protocols, child,
+				prot_id, prot_id + 1, GFP_KERNEL);
 		if (ret != prot_id) {
 			dev_err(dev, "SCMI protocol %d already activated. Skip\n",
 				prot_id);
diff --git a/include/linux/idr.h b/include/linux/idr.h
index 9939085d0e8d..8573a2b05ee6 100644
--- a/include/linux/idr.h
+++ b/include/linux/idr.h
@@ -14,6 +14,7 @@
 
 #include <errno.h>
 #include <linux/list.h>
+#include <linux/gfp.h>
 
 struct idr {
 	int			id;
@@ -58,7 +59,9 @@ static inline void *idr_find(struct idr *head, int id)
 	return idr ? idr->ptr : NULL;
 }
 
-int idr_alloc_one(struct idr *head, void *ptr, int start);
+int idr_alloc(struct idr *, void *ptr, int start, int end, gfp_t);
+int __must_check idr_alloc_u32(struct idr *, void *ptr, u32 *id,
+				unsigned long max, gfp_t);
 
 static inline void idr_init(struct idr *idr)
 {
diff --git a/lib/idr.c b/lib/idr.c
index a25e46b17b95..03bb6eadd533 100644
--- a/lib/idr.c
+++ b/lib/idr.c
@@ -6,8 +6,10 @@
 
 #include <errno.h>
 #include <linux/idr.h>
-#include <malloc.h>
+#include <linux/bug.h>
+#include <linux/limits.h>
 #include <linux/minmax.h>
+#include <xfuncs.h>
 
 struct idr *__idr_find(struct idr *head, int lookup_id)
 {
@@ -48,31 +50,107 @@ int idr_for_each(const struct idr *idr,
 	return 0;
 }
 
-static int idr_compare(struct list_head *a, struct list_head *b)
-{
-	int id_a = list_entry(a, struct idr, list)->id;
-	int id_b = list_entry(b, struct idr, list)->id;
-
-	return __compare3(id_a, id_b);
-}
-
-int idr_alloc_one(struct idr *head, void *ptr, int start)
+static struct idr *__idr_alloc(void *ptr, u32 start)
 {
 	struct idr *idr;
 
-	if (__idr_find(head, start))
-		return -EBUSY;
-
-	idr = malloc(sizeof(*idr));
-	if (!idr)
-		return -ENOMEM;
-
+	idr = xmalloc(sizeof(*idr));
 	idr->id = start;
 	idr->ptr = ptr;
 
-	list_add_sort(&idr->list, &head->list, idr_compare);
+	return idr;
+}
 
-	return start;
+/**
+ * idr_alloc_u32() - Allocate an ID.
+ * @head: IDR handle.
+ * @ptr: Pointer to be associated with the new ID.
+ * @nextid: Pointer to an ID.
+ * @max: The maximum ID to allocate (inclusive).
+ * @gfp: Memory allocation flags.
+ *
+ * Allocates an unused ID in the range specified by @nextid and @max.
+ * Note that @max is inclusive whereas the @end parameter to idr_alloc()
+ * is exclusive. The new ID is assigned to @nextid before the pointer
+ * is inserted into the IDR.
+ *
+ * Return: 0 if an ID was allocated, -ENOMEM if memory allocation failed,
+ * or -ENOSPC if no free IDs could be found.  If an error occurred,
+ * @nextid is unchanged.
+ */
+int idr_alloc_u32(struct idr *head, void *ptr, u32 *nextid,
+		  unsigned long max, gfp_t gfp)
+{
+	struct idr *idr, *prev, *next;
+
+	if (list_empty(&head->list)) {
+		idr = __idr_alloc(ptr, *nextid);
+		list_add_tail(&idr->list, &head->list);
+		return 0;
+	}
+
+	next = list_first_entry(&head->list, struct idr, list);
+	if (*nextid < next->id) {
+		idr = __idr_alloc(ptr, *nextid);
+		list_add(&idr->list, &head->list);
+		return 0;
+	}
+
+	prev = list_last_entry(&head->list, struct idr, list);
+	if (prev->id < max) {
+		*nextid = max_t(u32, prev->id + 1, *nextid);
+		idr = __idr_alloc(ptr, *nextid);
+		list_add_tail(&idr->list, &head->list);
+		return 0;
+	}
+
+	list_for_each_entry_safe(prev, next, &head->list, list) {
+		/* at the end of the list */
+		if (&next->list == &head->list)
+			break;
+
+		if (prev->id > max)
+			break;
+		if (next->id < *nextid)
+			continue;
+
+		*nextid = max_t(u32, prev->id + 1, *nextid);
+		idr = __idr_alloc(ptr, *nextid);
+		list_add(&idr->list, &prev->list);
+		return 0;
+	}
+
+	return -ENOSPC;
+}
+
+/**
+ * idr_alloc() - Allocate an ID.
+ * @head: IDR handle.
+ * @ptr: Pointer to be associated with the new ID.
+ * @start: The minimum ID (inclusive).
+ * @end: The maximum ID (exclusive).
+ * @gfp: Memory allocation flags.
+ *
+ * Allocates an unused ID in the range specified by @start and @end.  If
+ * @end is <= 0, it is treated as one larger than %INT_MAX.  This allows
+ * callers to use @start + N as @end as long as N is within integer range.
+ *
+ * Return: The newly allocated ID, -ENOMEM if memory allocation failed,
+ * or -ENOSPC if no free IDs could be found.
+ */
+int idr_alloc(struct idr *head, void *ptr, int start, int end, gfp_t gfp)
+{
+	u32 id = start;
+	int ret;
+
+	if (WARN_ON_ONCE(start < 0))
+		return -EINVAL;
+
+	ret = idr_alloc_u32(head, ptr, &id, end > 0 ? end - 1 : INT_MAX, gfp);
+	if (ret)
+		return ret;
+
+	return id;
 }
 
 static void __idr_remove(struct idr *idr)
diff --git a/test/self/idr.c b/test/self/idr.c
index bb902b041bb4..cc3217bb06e4 100644
--- a/test/self/idr.c
+++ b/test/self/idr.c
@@ -42,6 +42,11 @@ static int count_idr(int id, void *p, void *data)
 	return 0;
 }
 
+static inline int idr_alloc_one(struct idr *idr, void *ptr, int id)
+{
+	return idr_alloc_u32(idr, ptr, &id, id + 1, GFP_KERNEL) ?: id;
+}
+
 static void test_idr(void)
 {
 	void *ptr;
-- 
2.39.5