[PATCH 03/13] idr, radix-tree: Add get_tag_batch function

[Sandhya Bankar]

To implement select() on top of the IDR, we need to be able to get the
tags which represent the open files in bulk.  For this user, it makes
sense to get a batch of BITS_PER_LONG tags at a time, and until another
user shows up that wants something different, let's enforce that instead
of coping with arbitrary offsets.

Signed-off-by: Matthew Wilcox <mawilcox@xxxxxxxxxxxxx>
---
 include/linux/idr.h                 |  6 ++++
 include/linux/radix-tree.h          |  2 ++
 lib/radix-tree.c                    | 55 ++++++++++++++++++++++++++++++++++++-
 tools/testing/radix-tree/idr-test.c | 20 ++++++++++++++
 tools/testing/radix-tree/test.h     |  2 +-
 5 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/include/linux/idr.h b/include/linux/idr.h
index 9f71e63..d43cf01 100644
--- a/include/linux/idr.h
+++ b/include/linux/idr.h
@@ -147,6 +147,12 @@ static inline bool idr_tag_get(const struct idr *idr, int id, unsigned int tag)
 	return radix_tree_tag_get(&idr->idr_rt, id, tag);
 }
 
+static inline unsigned long idr_get_tag_batch(const struct idr *idr, int id,
+							unsigned int tag)
+{
+	return radix_tree_get_tag_batch(&idr->idr_rt, id, tag);
+}
+
 static inline void idr_init(struct idr *idr)
 {
 	INIT_RADIX_TREE(&idr->idr_rt, IDR_RT_MARKER);
diff --git a/include/linux/radix-tree.h b/include/linux/radix-tree.h
index 3e57350..f701e0b 100644
--- a/include/linux/radix-tree.h
+++ b/include/linux/radix-tree.h
@@ -339,6 +339,8 @@ void radix_tree_iter_tag_set(struct radix_tree_root *,
 		const struct radix_tree_iter *iter, unsigned int tag);
 void radix_tree_iter_tag_clear(struct radix_tree_root *,
 		const struct radix_tree_iter *iter, unsigned int tag);
+unsigned long radix_tree_get_tag_batch(const struct radix_tree_root *,
+				unsigned long index, unsigned int tag);
 unsigned int radix_tree_gang_lookup_tag(const struct radix_tree_root *,
 		void **results, unsigned long first_index,
 		unsigned int max_items, unsigned int tag);
diff --git a/lib/radix-tree.c b/lib/radix-tree.c
index 6723384..855ac8e 100644
--- a/lib/radix-tree.c
+++ b/lib/radix-tree.c
@@ -181,7 +181,8 @@ static inline void root_tag_clear_all(struct radix_tree_root *root)
 	root->gfp_mask &= (1 << ROOT_TAG_SHIFT) - 1;
 }
 
-static inline int root_tag_get(const struct radix_tree_root *root, unsigned tag)
+static inline bool root_tag_get(const struct radix_tree_root *root,
+								unsigned tag)
 {
 	return (__force int)root->gfp_mask & (1 << (tag + ROOT_TAG_SHIFT));
 }
@@ -1571,6 +1572,58 @@ int radix_tree_tag_get(const struct radix_tree_root *root,
 }
 EXPORT_SYMBOL(radix_tree_tag_get);
 
+static unsigned long
+__radix_tree_get_tag_batch(const struct radix_tree_root *root,
+				unsigned long index, unsigned int tag)
+{
+	struct radix_tree_node *node;
+	void __rcu **slot = NULL;
+	bool idr_free = is_idr(root) && (tag == IDR_FREE);
+
+	__radix_tree_lookup(root, index, &node, &slot);
+	if (!slot)
+		return idr_free ? ~0UL : 0;
+	if (!node)
+		return root_tag_get(root, tag) | (idr_free ? ~1UL : 0);
+	if (node->shift)
+		return idr_free ? ~0UL : 0;
+	return node->tags[tag][(index / BITS_PER_LONG) &
+						(RADIX_TREE_TAG_LONGS - 1)];
+}
+
+/**
+ * radix_tree_get_tag_batch() - get a batch of tags
+ * @root: radix tree root
+ * @index: start index of batch
+ * @tag: tag to get
+ *
+ * Get a batch of BITS_PER_LONG tags.  The only values of @index
+ * permitted are multiples of BITS_PER_LONG.
+ *
+ * Return: The tags for the next BITS_PER_LONG indices.
+ */
+unsigned long radix_tree_get_tag_batch(const struct radix_tree_root *root,
+				unsigned long index, unsigned int tag)
+{
+	unsigned long bits = 0;
+	unsigned shift = BITS_PER_LONG > RADIX_TREE_MAP_SIZE ? \
+						RADIX_TREE_MAP_SIZE : 0;
+
+	if (WARN_ON_ONCE(index & (BITS_PER_LONG - 1)))
+		return bits;
+
+	index += BITS_PER_LONG;
+	for (;;) {
+		index -= RADIX_TREE_MAP_SIZE;
+		bits |= __radix_tree_get_tag_batch(root, index, tag);
+		if (!(index & (BITS_PER_LONG - 1)))
+			break;
+		bits <<= shift;
+	}
+
+	return bits;
+}
+
 static inline void __set_iter_shift(struct radix_tree_iter *iter,
 					unsigned int shift)
 {
diff --git a/tools/testing/radix-tree/idr-test.c b/tools/testing/radix-tree/idr-test.c
index 334ce1c..3f9f429 100644
--- a/tools/testing/radix-tree/idr-test.c
+++ b/tools/testing/radix-tree/idr-test.c
@@ -139,6 +139,8 @@ void idr_tag_test(void)
 	DEFINE_IDR(idr);
 	struct item *item;
 
+	assert(idr_get_tag_batch(&idr, 0, IDR_FREE) == ~0UL);
+
 	for (i = 0; i < 100; i++) {
 		item = item_create(i, 0);
 		assert(idr_alloc(&idr, item, 0, 0, GFP_KERNEL) == i);
@@ -146,6 +148,20 @@ void idr_tag_test(void)
 			idr_tag_set(&idr, i, IDR_TEST);
 	}
 
+	assert(idr_get_tag_batch(&idr, 0, IDR_FREE) == 0);
+#if BITS_PER_LONG == 64
+	assert(idr_get_tag_batch(&idr, 0, IDR_TEST) == 0x8102040810204081UL);
+	assert(idr_get_tag_batch(&idr, 64, IDR_TEST) == 0x408102040UL);
+#else
+	assert(idr_get_tag_batch(&idr, 0, IDR_TEST) == 0x10204081UL);
+	assert(idr_get_tag_batch(&idr, 32, IDR_TEST) == 0x81020408UL);
+	assert(idr_get_tag_batch(&idr, 64, IDR_TEST) == 0x08102040UL);
+	assert(idr_get_tag_batch(&idr, 96, IDR_TEST) == 0x4UL);
+#endif
+	assert((int)idr_get_tag_batch(&idr, 64, IDR_FREE) == 0);
+	assert(idr_get_tag_batch(&idr, 128, IDR_FREE) == ~0UL);
+	assert(idr_get_tag_batch(&idr, 128, IDR_TEST) == 0);
+
 	for (i = 0; i < 100; i += 14) {
 		assert(idr_tag_get(&idr, i, IDR_TEST));
 		idr_tag_clear(&idr, i, IDR_TEST);
@@ -159,6 +175,10 @@ void idr_tag_test(void)
 		assert(item->index % 14 == 7);
 	}
 
+	item_free(idr_remove(&idr, 7), 7);
+	assert((int)idr_get_tag_batch(&idr, 0, IDR_TEST) == 0x00200000UL);
+	assert((int)idr_get_tag_batch(&idr, 0, IDR_FREE) == 0x00000080);
+
 	idr_for_each(&idr, item_idr_free, &idr);
 	idr_destroy(&idr);
 }
diff --git a/tools/testing/radix-tree/test.h b/tools/testing/radix-tree/test.h
index cbabea1..09616ff 100644
--- a/tools/testing/radix-tree/test.h
+++ b/tools/testing/radix-tree/test.h
@@ -52,7 +52,7 @@ struct item *
 /* Normally private parts of lib/radix-tree.c */
 struct radix_tree_node *entry_to_node(void *ptr);
 void radix_tree_dump(struct radix_tree_root *root);
-int root_tag_get(struct radix_tree_root *root, unsigned int tag);
+bool root_tag_get(struct radix_tree_root *root, unsigned int tag);
 unsigned long node_maxindex(struct radix_tree_node *);
 unsigned long shift_maxindex(unsigned int shift);
 int radix_tree_cpu_dead(unsigned int cpu);
-- 
1.8.3.1