[PATCH 02/13] idr: Add idr_for_each_entry_tagged()

[Sandhya Bankar]

Add the ability to iterate over tagged entries in the IDR with
idr_get_next_tag() and idr_for_each_entry_tagged().

Signed-off-by: Matthew Wilcox <mawilcox@xxxxxxxxxxxxx>
---
 include/linux/idr.h                 | 15 ++++++++++++++-
 lib/idr.c                           | 30 +++++++++++++++++++++++++++++-
 tools/testing/radix-tree/idr-test.c | 18 ++++++++++--------
 tools/testing/radix-tree/test.c     |  9 +++++++--
 tools/testing/radix-tree/test.h     |  1 +
 5 files changed, 61 insertions(+), 12 deletions(-)

diff --git a/include/linux/idr.h b/include/linux/idr.h
index 7eb4432..9f71e63 100644
--- a/include/linux/idr.h
+++ b/include/linux/idr.h
@@ -84,7 +84,8 @@ static inline void idr_set_cursor(struct idr *idr, unsigned int val)
 int idr_alloc_cyclic(struct idr *, void *entry, int start, int end, gfp_t);
 int idr_for_each(const struct idr *,
 		 int (*fn)(int id, void *p, void *data), void *data);
-void *idr_get_next(struct idr *, int *nextid);
+void *idr_get_next(const struct idr *, int *nextid);
+void *idr_get_next_tag(const struct idr *, int *nextid, unsigned int tag);
 void *idr_replace(struct idr *, void *, int id);
 void idr_destroy(struct idr *);
 
@@ -213,6 +214,18 @@ static inline void *idr_find(const struct idr *idr, int id)
 	     entry;							\
 	     ++id, (entry) = idr_get_next((idr), &(id)))
 
+/**
+ * idr_for_each_entry_tagged - iterate over IDs with a set tag
+ * @idr: IDR handle
+ * @entry: The pointer stored in @idr
+ * @id: The index of @entry in @idr
+ * @tag: tag to search for
+ */
+#define idr_for_each_entry_tagged(idr, entry, id, tag)			\
+	for (id = 0;							\
+	     ((entry) = idr_get_next_tag(idr, &(id), (tag))) != NULL;	\
+	     ++id)
+
 /*
  * IDA - IDR based id allocator, use when translation from id to
  * pointer isn't necessary.
diff --git a/lib/idr.c b/lib/idr.c
index b13682b..68e39c3 100644
--- a/lib/idr.c
+++ b/lib/idr.c
@@ -120,7 +120,7 @@ int idr_for_each(const struct idr *idr,
  * to the ID of the found value.  To use in a loop, the value pointed to by
  * nextid must be incremented by the user.
  */
-void *idr_get_next(struct idr *idr, int *nextid)
+void *idr_get_next(const struct idr *idr, int *nextid)
 {
 	struct radix_tree_iter iter;
 	void __rcu **slot;
@@ -135,6 +135,34 @@ void *idr_get_next(struct idr *idr, int *nextid)
 EXPORT_SYMBOL(idr_get_next);
 
 /**
+ * idr_get_next_tag - Find next tagged entry
+ * @idr: idr handle
+ * @nextid: Pointer to lowest possible ID to return
+ * @tag: tag to search for
+ *
+ * Returns the next tagged entry in the tree with an ID greater than
+ * or equal to the value pointed to by @nextid.  On exit, @nextid is updated
+ * to the ID of the found value.  To use in a loop, the value pointed to by
+ * nextid must be incremented by the user.  If a NULL entry is tagged, it
+ * will be returned.
+ */
+void *idr_get_next_tag(const struct idr *idr, int *nextid, unsigned int tag)
+{
+	struct radix_tree_iter iter;
+	void __rcu **slot;
+
+	radix_tree_iter_init(&iter, *nextid);
+	slot = radix_tree_next_chunk(&idr->idr_rt, &iter,
+					RADIX_TREE_ITER_TAGGED | tag);
+	if (!slot)
+		return NULL;
+
+	*nextid = iter.index;
+	return rcu_dereference_raw(*slot);
+}
+EXPORT_UNUSED_SYMBOL(idr_get_next_tag);
+
+/**
  * idr_replace - replace pointer for given id
  * @idr: idr handle
  * @ptr: New pointer to associate with the ID
diff --git a/tools/testing/radix-tree/idr-test.c b/tools/testing/radix-tree/idr-test.c
index fd94bee..334ce1c 100644
--- a/tools/testing/radix-tree/idr-test.c
+++ b/tools/testing/radix-tree/idr-test.c
@@ -23,19 +23,15 @@
 
 int item_idr_free(int id, void *p, void *data)
 {
-	struct item *item = p;
-	assert(item->index == id);
-	free(p);
-
+	item_free(p, id);
 	return 0;
 }
 
 void item_idr_remove(struct idr *idr, int id)
 {
 	struct item *item = idr_find(idr, id);
-	assert(item->index == id);
 	idr_remove(idr, id);
-	free(item);
+	item_free(item, id);
 }
 
 void idr_alloc_test(void)
@@ -139,11 +135,13 @@ void idr_null_test(void)
 
 void idr_tag_test(void)
 {
-	unsigned int i;
+	int i;
 	DEFINE_IDR(idr);
+	struct item *item;
 
 	for (i = 0; i < 100; i++) {
-		assert(idr_alloc(&idr, NULL, 0, 0, GFP_KERNEL) == i);
+		item = item_create(i, 0);
+		assert(idr_alloc(&idr, item, 0, 0, GFP_KERNEL) == i);
 		if (i % 7 == 0)
 			idr_tag_set(&idr, i, IDR_TEST);
 	}
@@ -157,6 +155,10 @@ void idr_tag_test(void)
 		assert(idr_tag_get(&idr, i, IDR_TEST) == (i % 14 == 7));
 	}
 
+	idr_for_each_entry_tagged(&idr, item, i, IDR_TEST) {
+		assert(item->index % 14 == 7);
+	}
+
 	idr_for_each(&idr, item_idr_free, &idr);
 	idr_destroy(&idr);
 }
diff --git a/tools/testing/radix-tree/test.c b/tools/testing/radix-tree/test.c
index 1a257d7..74f8e5c 100644
--- a/tools/testing/radix-tree/test.c
+++ b/tools/testing/radix-tree/test.c
@@ -62,13 +62,18 @@ void item_sanity(struct item *item, unsigned long index)
 	assert((item->index | mask) == (index | mask));
 }
 
+void item_free(struct item *item, unsigned long index)
+{
+	item_sanity(item, index);
+	free(item);
+}
+
 int item_delete(struct radix_tree_root *root, unsigned long index)
 {
 	struct item *item = radix_tree_delete(root, index);
 
 	if (item) {
-		item_sanity(item, index);
-		free(item);
+		item_free(item, index);
 		return 1;
 	}
 	return 0;
diff --git a/tools/testing/radix-tree/test.h b/tools/testing/radix-tree/test.h
index 0f8220c..cbabea1 100644
--- a/tools/testing/radix-tree/test.h
+++ b/tools/testing/radix-tree/test.h
@@ -13,6 +13,7 @@ struct item {
 int item_insert(struct radix_tree_root *root, unsigned long index);
 int item_insert_order(struct radix_tree_root *root, unsigned long index,
 			unsigned order);
+void item_free(struct item *item, unsigned long index);
 int item_delete(struct radix_tree_root *root, unsigned long index);
 struct item *item_lookup(struct radix_tree_root *root, unsigned long index);
 
-- 
1.8.3.1