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 <mawil...@microsoft.com>
---
 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

Reply via email to