Mercurial > hg > ucx / file diff

--- a/src/tree.c	Thu Mar 14 22:07:19 2024 +0100
+++ b/src/tree.c	Wed Mar 20 23:31:41 2024 +0100
@@ -178,10 +178,8 @@
 
 static void cx_tree_iter_next(void *it) {
     struct cx_tree_iterator_s *iter = it;
-    off_t const loc_next = iter->loc_next;
-    off_t const loc_children = iter->loc_children;
-
-    // TODO: support mutating iterator
+    ptrdiff_t const loc_next = iter->loc_next;
+    ptrdiff_t const loc_children = iter->loc_children;
 
     void *children;
 
@@ -280,3 +278,116 @@
 
     return iter;
 }
+
+static bool cx_tree_visitor_valid(void const *it) {
+    struct cx_tree_visitor_s const *iter = it;
+    return iter->node != NULL;
+}
+
+static void *cx_tree_visitor_current(void const *it) {
+    struct cx_tree_visitor_s const *iter = it;
+    return iter->node;
+}
+
+__attribute__((__nonnull__))
+static void cx_tree_visitor_enqueue_siblings(
+        struct cx_tree_visitor_s *iter, void *node, ptrdiff_t loc_next) {
+    node = tree_next(node);
+    while (node != NULL) {
+        struct cx_tree_visitor_queue_s *q;
+        q = malloc(sizeof(struct cx_tree_visitor_queue_s));
+        q->depth = iter->queue_last->depth;
+        q->node = node;
+        iter->queue_last->next = q;
+        iter->queue_last = q;
+        node = tree_next(node);
+    }
+    iter->queue_last->next = NULL;
+}
+
+static void cx_tree_visitor_next(void *it) {
+    struct cx_tree_visitor_s *iter = it;
+    ptrdiff_t const loc_next = iter->loc_next;
+    ptrdiff_t const loc_children = iter->loc_children;
+
+    // check if there is a next node
+    if (iter->queue_next == NULL) {
+        iter->node = NULL;
+        return;
+    }
+
+    // dequeue the next node
+    iter->node = iter->queue_next->node;
+    iter->depth = iter->queue_next->depth;
+    {
+        struct cx_tree_visitor_queue_s *q = iter->queue_next;
+        iter->queue_next = q->next;
+        if (iter->queue_next == NULL) {
+            assert(iter->queue_last == q);
+            iter->queue_last = NULL;
+        }
+        free(q);
+    }
+
+    // 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(
+        void *root,
+        ptrdiff_t loc_children,
+        ptrdiff_t loc_next
+) {
+    CxTreeVisitor iter;
+    iter.loc_children = loc_children;
+    iter.loc_next = loc_next;
+
+    // allocate stack
+    iter.depth = 0;
+
+    // visit the root node
+    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);
+    }
+
+    // assign base iterator functions
+    iter.base.mutating = false;
+    iter.base.remove = false;
+    iter.base.current_impl = NULL;
+    iter.base.valid = cx_tree_visitor_valid;
+    iter.base.next = cx_tree_visitor_next;
+    iter.base.current = cx_tree_visitor_current;
+
+    return iter;
+}
+