idr: support storing NULL in the IDR

The radix tree interprets storing NULL as a deleted entry.  Several users
of the IDR API use NULL as a temporary placeholder, or intentionally
convert entries between NULL and non-NULL pointers to keep IDs reserved
but not necessarily pointing to a valid entry at any given moment.  Adapt
the radix tree to cope with NULL pointers when being used as an IDR.

Link: http://lkml.kernel.org/r/1481931156-23469-1-git-send-email-mawilcox@linuxonhyperv.com
Signed-off-by: Matthew Wilcox <mawilcox@microsoft.com>
Cc: Stephen Rothwell <sfr@canb.auug.org.au>
Signed-off-by: Andrew Morton <akpm@linux-foundation.org>

3 files changed, 65 insertions, 14 deletions

diff --git a/lib/idr.c b/lib/idr.c
index ce229e2815f8..d13b11ad0e49 100644
--- a/lib/idr.c
+++ b/lib/idr.c
@@ -148,17 +148,16 @@ EXPORT_SYMBOL(idr_get_next);
 void *idr_replace(struct idr *idr, void *ptr, int id)
 {
 	struct radix_tree_node *node;
-	void **slot;
+	void **slot = NULL;
 	void *entry;
 
 	if (id < 0)
 		return ERR_PTR(-EINVAL);
-	if (!ptr || radix_tree_is_internal_node(ptr))
+	if (radix_tree_is_internal_node(ptr))
 		return ERR_PTR(-EINVAL);
 
 	entry = __radix_tree_lookup(&idr->idr_rt, id, &node, &slot);
-
-	if (!entry)
+	if (!slot || radix_tree_tag_get(&idr->idr_rt, id, IDR_FREE))
 		return ERR_PTR(-ENOENT);
 
 	__radix_tree_replace(&idr->idr_rt, node, slot, ptr, NULL, NULL);
diff --git a/lib/radix-tree.c b/lib/radix-tree.c
index b8b5fc1d8a42..4384f8045ad1 100644
--- a/lib/radix-tree.c
+++ b/lib/radix-tree.c
@@ -605,7 +605,7 @@ static int radix_tree_extend(struct radix_tree_root *root, gfp_t gfp,
 		maxshift += RADIX_TREE_MAP_SHIFT;
 
 	slot = root->rnode;
-	if (!slot)
+	if (!slot && (!is_idr(root) || root_tag_get(root, IDR_FREE)))
 		goto out;
 
 	do {
@@ -616,9 +616,10 @@ static int radix_tree_extend(struct radix_tree_root *root, gfp_t gfp,
 
 		if (is_idr(root)) {
 			all_tag_set(node, IDR_FREE);
-			if (!root_tag_get(root, IDR_FREE))
+			if (!root_tag_get(root, IDR_FREE)) {
 				tag_clear(node, IDR_FREE, 0);
-			root_tag_set(root, IDR_FREE);
+				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++) {
@@ -1082,7 +1083,11 @@ static void replace_slot(struct radix_tree_root *root,
 
 	WARN_ON_ONCE(radix_tree_is_internal_node(item));
 
-	count = !!item - !!old;
+	/* NULL is a valid entry in an IDR; do not adjust count */
+	if (is_idr(root))
+		count = 0;
+	else
+		count = !!item - !!old;
 	exceptional = !!radix_tree_exceptional_entry(item) -
 		      !!radix_tree_exceptional_entry(old);
 
@@ -1189,6 +1194,8 @@ void radix_tree_replace_slot(struct radix_tree_root *root,
 void radix_tree_iter_replace(struct radix_tree_root *root,
 		const struct radix_tree_iter *iter, void **slot, void *item)
 {
+	if (is_idr(root) && iter->node)
+		iter->node->count++;
 	__radix_tree_replace(root, iter->node, slot, item, NULL, NULL);
 }
 
@@ -1512,8 +1519,6 @@ int radix_tree_tag_get(struct radix_tree_root *root,
 		parent = entry_to_node(node);
 		offset = radix_tree_descend(parent, &node, index);
 
-		if (!node)
-			return 0;
 		if (!tag_get(parent, tag, offset))
 			return 0;
 		if (node == RADIX_TREE_RETRY)
@@ -1965,7 +1970,7 @@ void *radix_tree_delete_item(struct radix_tree_root *root,
 	int tag;
 
 	entry = __radix_tree_lookup(root, index, &node, &slot);
-	if (!entry)
+	if (!entry && !is_idr(root))
 		return NULL;
 
 	if (item && entry != item)
@@ -1982,9 +1987,12 @@ void *radix_tree_delete_item(struct radix_tree_root *root,
 
 	offset = get_slot_offset(node, slot);
 
-	if (is_idr(root))
+	if (is_idr(root)) {
+		if (tag_get(node, IDR_FREE, offset))
+			return entry;
 		node_tag_set(root, node, IDR_FREE, offset);
-	else
+		node->count--;
+	} else
 		for (tag = 0; tag < RADIX_TREE_MAX_TAGS; tag++)
 			node_tag_clear(root, node, tag, offset);
 
diff --git a/tools/testing/radix-tree/idr-test.c b/tools/testing/radix-tree/idr-test.c
index 64d125c93545..b3a3d9e76fb7 100644
--- a/tools/testing/radix-tree/idr-test.c
+++ b/tools/testing/radix-tree/idr-test.c
@@ -74,6 +74,48 @@ void idr_replace_test(void)
 	idr_destroy(&idr);
 }
 
+/*
+ * Unlike the radix tree, you can put a NULL pointer -- with care -- into
+ * the IDR.  Some interfaces, like idr_find() do not distinguish between
+ * "present, value is NULL" and "not present", but that's exactly what some
+ * users want.
+ */
+void idr_null_test(void)
+{
+	int i;
+	DEFINE_IDR(idr);
+
+	for (i = 0; i < 10; i++) {
+		assert(idr_alloc(&idr, NULL, 0, 0, GFP_KERNEL) == i);
+	}
+
+	assert(idr_replace(&idr, DUMMY_PTR, 3) == NULL);
+	assert(idr_replace(&idr, DUMMY_PTR, 4) == NULL);
+	assert(idr_replace(&idr, NULL, 4) == DUMMY_PTR);
+	assert(idr_replace(&idr, DUMMY_PTR, 11) == ERR_PTR(-ENOENT));
+	idr_remove(&idr, 5);
+	assert(idr_alloc(&idr, NULL, 0, 0, GFP_KERNEL) == 5);
+	idr_remove(&idr, 5);
+
+	for (i = 0; i < 10; i++)
+		idr_remove(&idr, i);
+	assert(idr_is_empty(&idr));
+
+	assert(idr_alloc(&idr, NULL, 0, 0, GFP_KERNEL) == 0);
+	assert(idr_replace(&idr, DUMMY_PTR, 3) == ERR_PTR(-ENOENT));
+	assert(idr_replace(&idr, DUMMY_PTR, 0) == NULL);
+	assert(idr_replace(&idr, NULL, 0) == DUMMY_PTR);
+
+	idr_destroy(&idr);
+	assert(idr_is_empty(&idr));
+
+	for (i = 1; i < 10; i++) {
+		assert(idr_alloc(&idr, NULL, 1, 0, GFP_KERNEL) == i);
+	}
+
+	idr_destroy(&idr);
+}
+
 void idr_checks(void)
 {
 	unsigned long i;
@@ -112,8 +154,8 @@ void idr_checks(void)
 	idr_destroy(&idr);
 
 	idr_replace_test();
-
 	idr_alloc_test();
+	idr_null_test();
 }
 
 /*
@@ -176,10 +218,12 @@ void ida_checks(void)
 	ida_remove(&ida, id);
 	assert(ida_is_empty(&ida));
 	ida_destroy(&ida);
+	assert(ida_is_empty(&ida));
 
 	ida_pre_get(&ida, GFP_KERNEL);
 	ida_get_new_above(&ida, 1, &id);
 	ida_destroy(&ida);
+	assert(ida_is_empty(&ida));
 
 	for (j = 1; j < 65537; j *= 2) {
 		for (i = 0; i < j; i++) {