Mercurial > hg > ucx / file diff

--- a/src/tree.c	Wed Mar 20 23:35:32 2024 +0100
+++ b/src/tree.c	Wed Apr 03 21:22:23 2024 +0200
@@ -186,8 +186,17 @@
     // check if we are currently exiting or entering nodes
     if (iter->exiting) {
         children = NULL;
+        // skipping on exit is pointless, just clear the flag
+        iter->skip = false;
     } else {
-        children = tree_children(iter->node);
+        if (iter->skip) {
+            // skip flag is set, pretend that there are no children
+            iter->skip = false;
+            children = NULL;
+        } else {
+            // try to enter the children (if any)
+            children = tree_children(iter->node);
+        }
     }
 
     if (children == NULL) {
@@ -263,10 +272,12 @@
 
     // visit the root node
     iter.node = root;
+    iter.next = NULL;
     iter.counter = 1;
     iter.depth = 1;
     iter.stack[0] = root;
     iter.exiting = false;
+    iter.skip = false;
 
     // assign base iterator functions
     iter.base.mutating = false;
@@ -310,6 +321,30 @@
     ptrdiff_t const loc_next = iter->loc_next;
     ptrdiff_t const loc_children = iter->loc_children;
 
+    // add the children of the current node to the queue
+    // unless the skip flag is set
+    void *children;
+    if (iter->skip) {
+        iter->skip = false;
+        children = NULL;
+    } else {
+        children = tree_children(iter->node);
+    }
+    if (children != NULL) {
+        struct cx_tree_visitor_queue_s *q;
+        q = malloc(sizeof(struct cx_tree_visitor_queue_s));
+        q->depth = iter->depth + 1;
+        q->node = children;
+        if (iter->queue_last == NULL) {
+            assert(iter->queue_next == NULL);
+            iter->queue_next = q;
+        } else {
+            iter->queue_last->next = q;
+        }
+        iter->queue_last = q;
+        cx_tree_visitor_enqueue_siblings(iter, children, loc_next);
+    }
+
     // check if there is a next node
     if (iter->queue_next == NULL) {
         iter->node = NULL;
@@ -331,23 +366,6 @@
 
     // increment the node counter
     iter->counter++;
-
-    // add the children of the new node to the queue
-    void *children = tree_children(iter->node);
-    if (children != NULL) {
-        struct cx_tree_visitor_queue_s *q;
-        q = malloc(sizeof(struct cx_tree_visitor_queue_s));
-        q->depth = iter->depth + 1;
-        q->node = children;
-        if (iter->queue_last == NULL) {
-            assert(iter->queue_next == NULL);
-            iter->queue_next = q;
-        } else {
-            iter->queue_last->next = q;
-        }
-        iter->queue_last = q;
-        cx_tree_visitor_enqueue_siblings(iter, children, loc_next);
-    }
 }
 
 CxTreeVisitor cx_tree_visitor(
@@ -366,19 +384,9 @@
     iter.node = root;
     iter.counter = 1;
     iter.depth = 1;
-
-    // put all children of root into the queue
-    void *children = tree_children(root);
-    if (children == NULL) {
-        iter.queue_next = NULL;
-        iter.queue_last = NULL;
-    } else {
-        iter.queue_next = malloc(sizeof(struct cx_tree_visitor_queue_s));
-        iter.queue_next->depth = 2;
-        iter.queue_next->node = children;
-        iter.queue_last = iter.queue_next;
-        cx_tree_visitor_enqueue_siblings(&iter, children, loc_next);
-    }
+    iter.skip = false;
+    iter.queue_next = NULL;
+    iter.queue_last = NULL;
 
     // assign base iterator functions
     iter.base.mutating = false;