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    stack: Vec<RedTree<'a, L>>,
32}
33
34impl<'a, L: Language> PreOrder<'a, L> {
35    /// Creates a new pre-order iterator starting from the given root.
36    pub fn new(root: RedTree<'a, L>) -> Self {
37        Self { stack: vec![root] }
38    }
39}
40
41impl<'a, L: Language> Iterator for PreOrder<'a, L> {
42    type Item = RedTree<'a, L>;
43
44    fn next(&mut self) -> Option<Self::Item> {
45        let next = self.stack.pop()?;
46
47        if let RedTree::Node(node) = next {
48            // Push children in reverse order so they are popped in correct order
49            let children = node.green.children();
50            let mut offset = node.offset + node.green.text_len() as usize;
51
52            for child in children.iter().rev() {
53                offset -= child.len() as usize;
54                match child {
55                    crate::GreenTree::Node(n) => {
56                        self.stack.push(RedTree::Node(RedNode::new(n, offset)));
57                    }
58                    crate::GreenTree::Leaf(t) => {
59                        self.stack.push(RedTree::Leaf(RedLeaf { kind: t.kind, span: core::range::Range { start: offset, end: offset + t.length as usize } }));
60                    }
61                }
62            }
63        }
64
65        Some(next)
66    }
67}
oak_core/visitor/mod.rs

oak_core/visitor/
mod.rs
