From: James Jones <jejones3141@gmail.com>
Date: Fri, 19 Feb 2021 19:28:25 +0000 (-0600)
Subject: Add rbtree_iter_delete() and corresponding test. (#3949)
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=77b181f0da7649fb01788141758b7d2dc0cae7ab;p=thirdparty%2Ffreeradius-server.git

Add rbtree_iter_delete() and corresponding test. (#3949)

rbtree_iter_delete() deletes the current node in an iterator
and remembers the next one; rbtree_iter_next_inorder() notices
when that's happened and uses the saved node, otherwise it
behaves as usual.
---

diff --git a/src/lib/util/rbtree.c b/src/lib/util/rbtree.c
index 6cca2e75d25..699caf0779d 100644
--- a/src/lib/util/rbtree.c
+++ b/src/lib/util/rbtree.c
@@ -916,6 +916,17 @@ void *rbtree_iter_next_inorder(fr_rb_tree_iter_t *iter)
 	 */
 	if (unlikely(iter->node == NIL)) return NULL;
 
+	/*
+	 *	rbtree_iter_delete() has already deleted this node,
+	 *	and saved the next one for us. (We check for NULL;
+	 *	NIL just means we're at the end.)
+	 */
+	if (!iter->node) {
+		iter->node = iter->next;
+		iter->next = NULL;
+		return iter->node->data;
+	}
+
 	if (x->right != NIL) {
 		x = x->right;
 
@@ -1122,6 +1133,23 @@ void *rbtree_iter_next_postorder(fr_rb_tree_iter_t *iter)
 	return x->data;
 }
 
+/** Remove the current node from the tree
+ *
+ * @note Only makes sense for in-order traversals.
+ *
+ * @param[in] iter	previously initialised with #rbtree_iter_init
+ */
+void rbtree_iter_delete(fr_rb_tree_iter_t *iter)
+{
+	fr_rb_node_t *x = iter->node;
+
+	if (unlikely(x == NIL)) return;
+	(void) rbtree_iter_next_inorder(iter);
+	iter->next = iter->node;
+	iter->node = NULL;
+	rbtree_delete_internal(iter->tree, x, true);
+}
+
 /** Explicitly unlock the tree
  *
  * @note Must be called if iterating over the tree ends early.
diff --git a/src/lib/util/rbtree.h b/src/lib/util/rbtree.h
index 9aa27f5dd78..d0f25a0b4bd 100644
--- a/src/lib/util/rbtree.h
+++ b/src/lib/util/rbtree.h
@@ -170,7 +170,8 @@ int		rbtree_walk(rbtree_t *tree, fr_rb_order_t order, fr_rb_walker_t compare, vo
  */
 typedef struct {
 	rbtree_t	*tree;			//!< Tree being iterated over.
-	fr_rb_node_t	*node;			///< current node.
+	fr_rb_node_t	*node;			///< current node--set to NULL (not NIL) by rbtree_iter_delete()
+	fr_rb_node_t	*next;			///< if non-NULL, next node cached by rbtree_iter_delete()
 } fr_rb_tree_iter_t;
 
 void		*rbtree_iter_init_inorder(fr_rb_tree_iter_t *iter, rbtree_t *tree) CC_HINT(nonnull);
@@ -187,6 +188,8 @@ void		*rbtree_iter_next_postorder(fr_rb_tree_iter_t *iter) CC_HINT(nonnull);
 
 void		rbtree_iter_done(fr_rb_tree_iter_t *iter) CC_HINT(nonnull);
 
+void		rbtree_iter_delete(fr_rb_tree_iter_t *iter) CC_HINT(nonnull);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/lib/util/rbtree_tests.c b/src/lib/util/rbtree_tests.c
index fd8803f51ff..6b091fc7d2e 100644
--- a/src/lib/util/rbtree_tests.c
+++ b/src/lib/util/rbtree_tests.c
@@ -158,10 +158,75 @@ static void test_rbtree_iter_postorder(void)
 	talloc_free(t);
 }
 
+/*
+ *	primality test used in rbtree_delete_iter() test.
+ */
+static bool is_prime(uint32_t n)
+{
+	uint32_t	i, q;
+
+	if (n < 2) return false;
+
+	for (i = 2; (q = n / i) >= i; i++) {
+		if (i * q == n) return false;
+	}
+	return true;
+}
+
+uint32_t	non_primes[] = { 1,  4,  6,  8,  9, 10, 12, 14, 15, 16, 18, 20, 21, 22, 24, 25, 26, 27, 28,
+				30, 32, 33, 34, 35, 36, 38, 39, 40, 42, 44, 45, 46, 48, 49, 50};
+
+static void test_rbtree_iter_delete(void)
+{
+	rbtree_t 		*t;
+	size_t			i;
+	fr_rb_test_node_t	*p;
+	fr_rb_tree_iter_t	iter;
+
+	t = rbtree_alloc(NULL, fr_rb_test_node_t, node, fr_rb_test_cmp, NULL, RBTREE_FLAG_LOCK);
+	TEST_CHECK(t != NULL);
+
+ 	/*
+ 	 *	Initialise the test nodes
+ 	 *	with integers from 1 to 50.
+ 	 */
+	for (i = 1; i <= 50; i++) {
+		p = talloc(t, fr_rb_test_node_t);
+		p->num = i;
+		rbtree_insert(t, p);
+	}
+
+	/*
+	 *	Remove the primes.
+	 */
+	for (p = rbtree_iter_init_inorder(&iter, t);
+	     p;
+	     p = rbtree_iter_next_inorder(&iter)) {
+		if (is_prime(p->num)) rbtree_iter_delete(&iter);
+	}
+
+	/*
+	 *	Check that all the non-primes are still there.
+	 */
+	for (p = rbtree_iter_init_inorder(&iter, t), i = 0;
+	     p;
+	     p = rbtree_iter_next_inorder(&iter), i++) {
+		TEST_MSG("Checking non_primes[%zu] = %u vs p->num = %u", i, non_primes[i], p->num);
+		TEST_CHECK(non_primes[i] == p->num);
+		TEST_CHECK(pthread_mutex_trylock(&t->mutex) != 0);	/* Lock still held */
+	}
+	TEST_CHECK(pthread_mutex_trylock(&t->mutex) == 0);		/* Lock released */
+	pthread_mutex_unlock(&t->mutex);
+
+	talloc_free(t);
+}
+
+
 TEST_LIST = {
 	{ "rbtree_iter_inorder",            test_rbtree_iter_inorder },
 	{ "rbtree_iter_preorder",           test_rbtree_iter_preorder },
 	{ "rbtree_iter_postorder",          test_rbtree_iter_postorder },
+	{ "rbtree_iter_delete",             test_rbtree_iter_delete },
 
 	{ NULL }
 };