Mercurial > hg > ucx / changeset

--- a/src/cx/tree.h	Wed Feb 21 18:53:55 2024 +0100
+++ b/src/cx/tree.h	Mon Feb 26 21:07:23 2024 +0100
@@ -90,6 +90,11 @@
      */
     void *node;
     /**
+     * Stores a copy of the next pointer of the visited node.
+     * Allows freeing a node on exit without corrupting the iteration.
+     */
+    void *next;
+    /**
      * Internal stack.
      * Will be automatically freed once the iterator becomes invalid.
      *

--- a/src/tree.c	Wed Feb 21 18:53:55 2024 +0100
+++ b/src/tree.c	Mon Feb 26 21:07:23 2024 +0100
@@ -197,7 +197,12 @@
         void *next;
         cx_tree_iter_search_next:
         // check if there is a sibling
-        next = tree_next(iter->node);
+        if (iter->exiting) {
+            next = iter->next;
+        } else {
+            next = tree_next(iter->node);
+            iter->next = next;
+        }
         if (next == NULL) {
             // no sibling, we are done with this node and exit
             if (iter->visit_on_exit && !iter->exiting) {
@@ -208,7 +213,7 @@
                 if (iter->depth == 1) {
                     // there is no parent - we have iterated the entire tree
                     // invalidate the iterator and free the node stack
-                    iter->node = NULL;
+                    iter->node = iter->next = NULL;
                     iter->stack_capacity = iter->depth = 0;
                     free(iter->stack);
                     iter->stack = NULL;

--- a/tests/test_tree.c	Wed Feb 21 18:53:55 2024 +0100
+++ b/tests/test_tree.c	Mon Feb 26 21:07:23 2024 +0100
@@ -30,6 +30,8 @@
 
 #include "cx/test.h"
 
+#include "util_allocator.h"
+
 typedef struct tree_node {
     struct tree_node *parent;
     struct tree_node *next;
@@ -465,6 +467,46 @@
     free(actual);
 }
 
+CX_TEST(test_tree_iterator_free_nodes) {
+    CxTestingAllocator talloc;
+    cx_testing_allocator_init(&talloc);
+    CxAllocator *alloc = &talloc.base;
+    CX_TEST_DO {
+        tree_node *root = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *a = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *b = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *c = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *aa = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *ab = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *ba = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *ca = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *cb = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *cc = cxCalloc(alloc, 1, sizeof(tree_node));
+        tree_node *cba = cxCalloc(alloc, 1, sizeof(tree_node));
+
+        cx_tree_link(root, a, tree_node_layout);
+        cx_tree_link(root, b, tree_node_layout);
+        cx_tree_link(root, c, tree_node_layout);
+        cx_tree_link(a, aa, tree_node_layout);
+        cx_tree_link(a, ab, tree_node_layout);
+        cx_tree_link(b, ba, tree_node_layout);
+        cx_tree_link(c, ca, tree_node_layout);
+        cx_tree_link(c, cb, tree_node_layout);
+        cx_tree_link(c, cc, tree_node_layout);
+        cx_tree_link(cb, cba, tree_node_layout);
+
+        CxTreeIterator iter = cx_tree_iterator(root, true, tree_child_list);
+        cx_foreach(tree_node *, node, iter) {
+            if (iter.exiting) {
+                cxFree(alloc,node);
+            }
+        }
+
+        CX_TEST_ASSERT(cx_testing_allocator_verify(&talloc));
+    }
+    cx_testing_allocator_destroy(&talloc);
+}
+
 CxTestSuite *cx_test_suite_tree_low_level(void) {
     CxTestSuite *suite = cx_test_suite_new("tree (low level)");
 
@@ -477,6 +519,7 @@
     cx_test_register(suite, test_tree_iterator_basic_only_enter);
     cx_test_register(suite, test_tree_iterator_basic_enter_and_exit);
     cx_test_register(suite, test_tree_iterator_xml);
+    cx_test_register(suite, test_tree_iterator_free_nodes);
 
     return suite;
 }