cairo_lang_syntax/node/
iter.rs

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

cairo_lang_syntax/node/
iter.rs
