Mercurial > hg > ucx / changeset

--- a/src/tree.c	Mon Feb 19 22:12:13 2024 +0100
+++ b/src/tree.c	Wed Feb 21 18:32:38 2024 +0100
@@ -182,40 +182,58 @@
     off_t const loc_children = iter->loc_children;
 
     // TODO: support mutating iterator
-    // TODO: support visit_on_exit
+
+    void *children;
 
-    void *children = tree_children(iter->node);
+    // check if we are currently exiting or entering nodes
+    if (iter->exiting) {
+        children = NULL;
+    } else {
+        children = tree_children(iter->node);
+    }
+
     if (children == NULL) {
         // search for the next node
-        void *current = iter->node;
-        do {
-            // check if there is a sibling
-            void *next = tree_next(current);
-            if (next == NULL) {
-                // no sibling, check again for parent node
-                --iter->depth;
-                if (iter->depth == 0) {
+        void *next;
+        cx_tree_iter_search_next:
+        // check if there is a sibling
+        next = tree_next(iter->node);
+        if (next == NULL) {
+            // no sibling, we are done with this node and exit
+            if (iter->visit_on_exit && !iter->exiting) {
+                // iter is supposed to visit the node again
+                iter->exiting = true;
+            } else {
+                iter->exiting = false;
+                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;
-                    current = NULL;
-                    iter->stack_capacity = 0;
+                    iter->stack_capacity = iter->depth = 0;
                     free(iter->stack);
                     iter->stack = NULL;
                 } else {
                     // the parent node can be obtained from the top of stack
                     // this way we can avoid the loc_parent in the iterator
-                    current = iter->stack[iter->depth - 1];
+                    iter->depth--;
+                    iter->node = iter->stack[iter->depth - 1];
+                    // retry with the parent node to find a sibling
+                    goto cx_tree_iter_search_next;
                 }
+            }
+        } else {
+            if (iter->visit_on_exit && !iter->exiting) {
+                // iter is supposed to visit the node again
+                iter->exiting = true;
             } else {
-                // move to the sibling and break the loop
-                current = NULL;
+                iter->exiting = false;
+                // move to the sibling
                 iter->counter++;
                 iter->node = next;
                 // new top of stack is the sibling
                 iter->stack[iter->depth - 1] = next;
             }
-        } while (current != NULL);
+        }
     } else {
         // node has children, push the first child onto the stack and enter it
         cx_array_simple_add(iter->stack, children);

--- a/tests/test_tree.c	Mon Feb 19 22:12:13 2024 +0100
+++ b/tests/test_tree.c	Wed Feb 21 18:32:38 2024 +0100
@@ -271,7 +271,7 @@
     }
 }
 
-CX_TEST(test_tree_iterator_basic_test_only_enter) {
+CX_TEST(test_tree_iterator_basic_only_enter) {
     tree_node root = {0};
     tree_node a = {0};
     tree_node b = {0};
@@ -330,6 +330,64 @@
     }
 }
 
+CX_TEST(test_tree_iterator_basic_enter_and_exit) {
+    tree_node root = {0};
+    tree_node a = {0};
+    tree_node b = {0};
+    tree_node c = {0};
+    tree_node aa = {0};
+    tree_node ab = {0};
+    tree_node ba = {0};
+    tree_node ca = {0};
+    tree_node cb = {0};
+    tree_node cc = {0};
+    tree_node cba = {0};
+
+    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);
+    CX_TEST_DO {
+        CxTreeIterator iter = cx_tree_iterator(&root, true, tree_child_list);
+        unsigned chk = 0;
+        cx_foreach(tree_node*, node, iter) {
+            CX_TEST_ASSERT(iter.exiting || node->data == 0);
+            node->data++;
+            chk++;
+            CX_TEST_ASSERT(node == iter.node);
+            if (node == &root) {
+                CX_TEST_ASSERT(iter.depth == 1);
+            } else if (node == &a || node == &b || node == &c) {
+                CX_TEST_ASSERT(iter.depth == 2);
+            } else if (node == &cba) {
+                CX_TEST_ASSERT(iter.depth == 4);
+            } else {
+                CX_TEST_ASSERT(iter.depth == 3);
+            }
+        }
+        CX_TEST_ASSERT(iter.counter == 11);
+        CX_TEST_ASSERT(chk == 22);
+        CX_TEST_ASSERT(iter.stack == NULL);
+        CX_TEST_ASSERT(root.data == 2);
+        CX_TEST_ASSERT(a.data == 2);
+        CX_TEST_ASSERT(b.data == 2);
+        CX_TEST_ASSERT(c.data == 2);
+        CX_TEST_ASSERT(aa.data == 2);
+        CX_TEST_ASSERT(ab.data == 2);
+        CX_TEST_ASSERT(ba.data == 2);
+        CX_TEST_ASSERT(ca.data == 2);
+        CX_TEST_ASSERT(cb.data == 2);
+        CX_TEST_ASSERT(cc.data == 2);
+        CX_TEST_ASSERT(cba.data == 2);
+    }
+}
+
 CxTestSuite *cx_test_suite_tree_low_level(void) {
     CxTestSuite *suite = cx_test_suite_new("tree (low level)");
 
@@ -339,7 +397,8 @@
     cx_test_register(suite, test_tree_unlink);
     cx_test_register(suite, test_tree_search);
     cx_test_register(suite, test_tree_iterator_create_and_dispose);
-    cx_test_register(suite, test_tree_iterator_basic_test_only_enter);
+    cx_test_register(suite, test_tree_iterator_basic_only_enter);
+    cx_test_register(suite, test_tree_iterator_basic_enter_and_exit);
 
     return suite;
 }