oak_core/visitor/
mod.rs

1//! Tree traversal and transformation utilities.
2//!
3//! This module provides traits and utilities for visiting and transforming red-green trees.
4
5use crate::{
6    Language,
7    tree::red_tree::{RedLeaf, RedNode, RedTree},
8};
9
10/// A visitor for traversing a red-green tree.
11pub trait Visitor<'a, L: Language> {
12    /// Visits a red node.
13    fn visit_node(&mut self, node: RedNode<'a, L>);
14
15    /// Visits a red leaf kind.
16    fn visit_token(&mut self, token: RedLeaf<L>);
17
18    /// Helper to walk children of a node.
19    fn walk_node(&mut self, node: RedNode<'a, L>) {
20        for child in node.children() {
21            match child {
22                RedTree::Node(n) => self.visit_node(n),
23                RedTree::Leaf(t) => self.visit_token(t),
24            }
25        }
26    }
27}
28
29/// A pre-order traversal iterator for red trees.
30pub struct PreOrder<'a, L: Language> {
31    /// The stack of nodes to visit.
32    stack: Vec<RedTree<'a, L>>,
33}
34
35impl<'a, L: Language> PreOrder<'a, L> {
36    /// Creates a new pre-order iterator starting from the given root.
37    pub fn new(root: RedTree<'a, L>) -> Self {
38        Self { stack: vec![root] }
39    }
40}
41
42impl<'a, L: Language> Iterator for PreOrder<'a, L> {
43    type Item = RedTree<'a, L>;
44
45    fn next(&mut self) -> Option<Self::Item> {
46        let next = self.stack.pop()?;
47
48        if let RedTree::Node(node) = next {
49            // Push children in reverse order so they are popped in correct order
50            let children = node.green.children();
51            let mut offset = node.offset + node.green.text_len() as usize;
52
53            for child in children.iter().rev() {
54                offset -= child.len() as usize;
55                match child {
56                    crate::GreenTree::Node(n) => self.stack.push(RedTree::Node(RedNode::new(n, offset))),
57                    crate::GreenTree::Leaf(t) => self.stack.push(RedTree::Leaf(RedLeaf { kind: t.kind, span: core::range::Range { start: offset, end: offset + t.length as usize } })),
58                }
59            }
60        }
61
62        Some(next)
63    }
64}
oak_core/visitor/mod.rs

oak_core/visitor/
mod.rs
