[PATCH 01/13] idr: Add ability to set/clear tags

[Matthew Wilcox <bankarsandhya512@gmail.com>]

Now that the IDR uses the radix tree, we can expose the radix tree tags
to users of the IDR.  A few spots in the radix tree needed to be changed
to cope with the fact that the IDR can have NULL pointers with tags set.
One of the more notable changes is that IDR_FREE really is special -- an
index which is out of range of the current tree height is free.

Signed-off-by: Matthew Wilcox <mawilcox@microsoft.com>
---
 include/linux/idr.h                 | 57 +++++++++++++++++++++++++++++++++++--
 lib/radix-tree.c                    | 31 +++++++++-----------
 tools/testing/radix-tree/idr-test.c | 27 ++++++++++++++++++
 3 files changed, 95 insertions(+), 20 deletions(-)

diff --git a/include/linux/idr.h b/include/linux/idr.h
index bf70b3e..7eb4432 100644
--- a/include/linux/idr.h
+++ b/include/linux/idr.h
@@ -22,8 +22,8 @@ struct idr {
 };
 
 /*
- * The IDR API does not expose the tagging functionality of the radix tree
- * to users.  Use tag 0 to track whether a node has free space below it.
+ * Reserve one of the radix tree tags to track whether a node has free
+ * space below it.
  */
 #define IDR_FREE	0
 
@@ -93,6 +93,59 @@ static inline void *idr_remove(struct idr *idr, int id)
 	return radix_tree_delete_item(&idr->idr_rt, id, NULL);
 }
 
+/**
+ * idr_tag_set - Set a tag on an entry
+ * @idr: IDR pointer
+ * @id: ID of entry to tag
+ * @tag: Tag index to set
+ *
+ * If there is an entry at @id in this IDR, set a tag on it and return
+ * the address of the entry.  If @id is outside the range of the IDR,
+ * return NULL.  This API does not allow you to set IDR_FREE on an entry;
+ * use idr_remove() for that.  The implementation does not currently check
+ * that IDR_FREE is clear, so it is possible to set a tag on a free entry.
+ * This is not recommended and may change in the future.
+ */
+static inline void *idr_tag_set(struct idr *idr, int id, unsigned int tag)
+{
+	BUG_ON(tag == IDR_FREE);
+	return radix_tree_tag_set(&idr->idr_rt, id, tag);
+}
+
+/**
+ * idr_tag_clear - Clear a tag on an entry
+ * @idr: IDR pointer
+ * @id: ID of entry to tag
+ * @tag: Tag index to clear
+ *
+ * If there is an entry at @id in this IDR, clear its tag and return
+ * the address of the entry.  If @id is outside the range of the IDR,
+ * return NULL.  This API does not allow you to clear IDR_FREE on an entry;
+ * use idr_alloc() for that.  The implementation does not currently check
+ * that IDR_FREE is clear, so it is possible to clear a tag on a free entry.
+ * This is not recommended and may change in the future.
+ */
+static inline void *idr_tag_clear(struct idr *idr, int id, unsigned int tag)
+{
+	BUG_ON(tag == IDR_FREE);
+	return radix_tree_tag_clear(&idr->idr_rt, id, tag);
+}
+
+/**
+ * idr_tag_get - Return whether a particular entry has a tag set
+ * @idr: IDR pointer
+ * @id: ID of entry to check
+ * @tag: Tag index to check
+ *
+ * Returns true/false depending whether @tag is set on this ID.  Unlike
+ * idr_tag_set() or idr_tag_clear(), you can use the IDR_FREE tag value,
+ * as it can be useful to know whether a particular ID has been allocated.
+ */
+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 void idr_init(struct idr *idr)
 {
 	INIT_RADIX_TREE(&idr->idr_rt, IDR_RT_MARKER);
diff --git a/lib/radix-tree.c b/lib/radix-tree.c
index 691a9ad..6723384 100644
--- a/lib/radix-tree.c
+++ b/lib/radix-tree.c
@@ -638,18 +638,17 @@ static int radix_tree_extend(struct radix_tree_root *root, gfp_t gfp,
 		if (!node)
 			return -ENOMEM;
 
+		/* Propagate the aggregated tag info to the new child */
+		for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++) {
+			if (root_tag_get(root, tag))
+				tag_set(node, tag, 0);
+		}
 		if (is_idr(root)) {
 			all_tag_set(node, IDR_FREE);
 			if (!root_tag_get(root, IDR_FREE)) {
 				tag_clear(node, IDR_FREE, 0);
 				root_tag_set(root, IDR_FREE);
 			}
-		} else {
-			/* Propagate the aggregated tag info to the new child */
-			for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++) {
-				if (root_tag_get(root, tag))
-					tag_set(node, tag, 0);
-			}
 		}
 
 		BUG_ON(shift > BITS_PER_LONG);
@@ -1434,8 +1433,6 @@ void *radix_tree_tag_set(struct radix_tree_root *root,
 
 		parent = entry_to_node(node);
 		offset = radix_tree_descend(parent, &node, index);
-		BUG_ON(!node);
-
 		if (!tag_get(parent, tag, offset))
 			tag_set(parent, tag, offset);
 	}
@@ -1489,7 +1486,7 @@ static void node_tag_clear(struct radix_tree_root *root,
  *	Clear the search tag (which must be < RADIX_TREE_MAX_TAGS)
  *	corresponding to @index in the radix tree.  If this causes
  *	the leaf node to have no tags set then clear the tag in the
- *	next-to-leaf node, etc.
+ *	parent node, etc.
  *
  *	Returns the address of the tagged item on success, else NULL.  ie:
  *	has the same return value and semantics as radix_tree_lookup().
@@ -1512,8 +1509,7 @@ void *radix_tree_tag_clear(struct radix_tree_root *root,
 		offset = radix_tree_descend(parent, &node, index);
 	}
 
-	if (node)
-		node_tag_clear(root, parent, tag, offset);
+	node_tag_clear(root, parent, tag, offset);
 
 	return node;
 }
@@ -1552,11 +1548,11 @@ int radix_tree_tag_get(const struct radix_tree_root *root,
 	struct radix_tree_node *node, *parent;
 	unsigned long maxindex;
 
-	if (!root_tag_get(root, tag))
-		return 0;
-
 	radix_tree_load_root(root, &node, &maxindex);
 	if (index > maxindex)
+		return is_idr(root) && (tag == IDR_FREE);
+
+	if (!root_tag_get(root, tag))
 		return 0;
 
 	while (radix_tree_is_internal_node(node)) {
@@ -1783,7 +1779,7 @@ void __rcu **radix_tree_next_chunk(const struct radix_tree_root *root,
 			child = rcu_dereference_raw(node->slots[offset]);
 		}
 
-		if (!child)
+		if (!child && !is_idr(root))
 			goto restart;
 		if (child == RADIX_TREE_RETRY)
 			break;
@@ -1994,11 +1990,10 @@ static bool __radix_tree_delete(struct radix_tree_root *root,
 	unsigned offset = get_slot_offset(node, slot);
 	int tag;
 
+	for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++)
+		node_tag_clear(root, node, tag, offset);
 	if (is_idr(root))
 		node_tag_set(root, node, IDR_FREE, offset);
-	else
-		for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++)
-			node_tag_clear(root, node, tag, offset);
 
 	replace_slot(slot, NULL, node, -1, exceptional);
 	return node && delete_node(root, node, NULL, NULL);
diff --git a/tools/testing/radix-tree/idr-test.c b/tools/testing/radix-tree/idr-test.c
index 30cd0b2..fd94bee 100644
--- a/tools/testing/radix-tree/idr-test.c
+++ b/tools/testing/radix-tree/idr-test.c
@@ -135,6 +135,32 @@ void idr_null_test(void)
 	assert(idr_is_empty(&idr));
 }
 
+#define IDR_TEST	1
+
+void idr_tag_test(void)
+{
+	unsigned int i;
+	DEFINE_IDR(idr);
+
+	for (i = 0; i < 100; i++) {
+		assert(idr_alloc(&idr, NULL, 0, 0, GFP_KERNEL) == i);
+		if (i % 7 == 0)
+			idr_tag_set(&idr, i, IDR_TEST);
+	}
+
+	for (i = 0; i < 100; i += 14) {
+		assert(idr_tag_get(&idr, i, IDR_TEST));
+		idr_tag_clear(&idr, i, IDR_TEST);
+	}
+
+	for (i = 0; i < 100; i++) {
+		assert(idr_tag_get(&idr, i, IDR_TEST) == (i % 14 == 7));
+	}
+
+	idr_for_each(&idr, item_idr_free, &idr);
+	idr_destroy(&idr);
+}
+
 void idr_nowait_test(void)
 {
 	unsigned int i;
@@ -225,6 +251,7 @@ void idr_checks(void)
 	idr_replace_test();
 	idr_alloc_test();
 	idr_null_test();
+	idr_tag_test();
 	idr_nowait_test();
 	idr_get_next_test();
 }
-- 
1.8.3.1