cairo_lang_syntax/node/
iter.rs

1use crate::node::SyntaxNode;
2use crate::node::db::SyntaxGroup;
3
4/// `WalkEvent` describes tree walking process.
5#[derive(Debug, Copy, Clone)]
6pub enum WalkEvent<T> {
7    /// Fired before traversing the node.
8    Enter(T),
9    /// Fired after the node is traversed.
10    Leave(T),
11}
12
13impl<T> WalkEvent<T> {
14    pub fn map<U>(self, f: impl FnOnce(T) -> U) -> WalkEvent<U> {
15        match self {
16            WalkEvent::Enter(it) => WalkEvent::Enter(f(it)),
17            WalkEvent::Leave(it) => WalkEvent::Leave(f(it)),
18        }
19    }
20}
21
22/// Traverse the subtree rooted at the current node (including the current node) in preorder,
23/// excluding tokens.
24pub struct Preorder<'a> {
25    db: &'a dyn SyntaxGroup,
26    // FIXME(mkaput): Is it possible to avoid allocating iterators in layers here?
27    //   This code does it because without fast parent & prev/next sibling operations it has to
28    //   maintain DFS trace.
29    layers: Vec<PreorderLayer>,
30}
31
32struct PreorderLayer {
33    start: SyntaxNode,
34    children: Option<(Vec<SyntaxNode>, usize)>,
35}
36
37impl<'a> Preorder<'a> {
38    pub(super) fn new(start: SyntaxNode, db: &'a dyn SyntaxGroup) -> Self {
39        let initial_layer = PreorderLayer { start, children: None };
40
41        // NOTE(mkaput): Reserving some capacity should help amortization and thus make this
42        // iterator more performant. This wasn't benchmarked though and the capacity is just an
43        // educated guess, based on typical depth of syntax files in test suites.
44        let mut layers = Vec::with_capacity(32);
45        layers.push(initial_layer);
46
47        Self { db, layers }
48    }
49}
50
51impl Iterator for Preorder<'_> {
52    type Item = WalkEvent<SyntaxNode>;
53
54    fn next(&mut self) -> Option<Self::Item> {
55        // Lack of layers to traverse means end of iteration, so early return here.
56        //
57        // The layer is popped here to gain exclusive ownership of it without taking exclusive
58        // ownership of the layers stack.
59        let mut layer = self.layers.pop()?;
60        match layer.children.take() {
61            None => {
62                // #1: If children iterator is not initialized, this means entire iteration just
63                // started, and the enter event for start node has to be emitted.
64                let event = WalkEvent::Enter(layer.start);
65                layer.children = Some((layer.start.get_children(self.db), 0));
66                self.layers.push(layer);
67                Some(event)
68            }
69            Some((nodes, index)) => {
70                match nodes.get(index) {
71                    None => {
72                        // #2: If children iterator is exhausted, this means iteration of start node
73                        // just finished, and the layer needs to be popped (i.e. not pushed back)
74                        // and leave event for this node needs to be
75                        // emitted.
76                        Some(WalkEvent::Leave(layer.start))
77                    }
78                    Some(start) => {
79                        // #3: Otherwise the iterator is just in the middle of visiting a child, so
80                        // push a new layer to iterate it. To avoid
81                        // recursion, step #1 is duplicated and
82                        // inlined here.
83                        let event = WalkEvent::Enter(*start);
84                        let new_layer = PreorderLayer {
85                            children: Some((start.get_children(self.db), 0)),
86                            start: *start,
87                        };
88                        layer.children = Some((nodes, index + 1));
89                        self.layers.extend([layer, new_layer]);
90                        Some(event)
91                    }
92                }
93            }
94        }
95    }
96}
cairo_lang_syntax/node/iter.rs

cairo_lang_syntax/node/
iter.rs
