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 6cca2e75d2527d04467d5878bb2e2435fa8a7b62..699caf0779d6b0958fce387eeedb20441e638733 100644 (file)
--- 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 9aa27f5dd7890a2fc3aa7f3bac60bf06e97e7d30..d0f25a0b4bdc7497b7b1591635e54c3afcac9341 100644 (file)
--- 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 fd8803f51ff769950057d6dffbe38bd9e0e7f11e..6b091fc7d2e071a9c763ce7b540f11541324c020 100644 (file)
--- 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 }
 };