cairo_lang_syntax/node/
iter.rs

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

cairo_lang_syntax/node/
iter.rs
