tree_sitter_iter/
lib.rs

1//! A very simple pre-order iterator for tree-sitter CSTs.
2
3#![deny(rustdoc::broken_intra_doc_links)]
4#![deny(missing_docs)]
5#![allow(clippy::redundant_field_names)]
6#![forbid(unsafe_code)]
7
8use tree_sitter::{Node, Tree, TreeCursor};
9
10/// A pre-order iterator over the nodes of a tree-sitter syntax tree.
11pub struct TreeIter<'tree> {
12    cursor: Option<TreeCursor<'tree>>,
13}
14
15impl<'tree> TreeIter<'tree> {
16    /// Creates a new `TreeSitterIter` for the given syntax tree.
17    pub fn new(tree: &'tree Tree) -> Self {
18        Self {
19            cursor: Some(tree.root_node().walk()),
20        }
21    }
22}
23
24impl<'tree> Iterator for TreeIter<'tree> {
25    type Item = Node<'tree>;
26
27    fn next(&mut self) -> Option<Self::Item> {
28        let cursor = match &mut self.cursor {
29            Some(cursor) => cursor,
30            None => return None,
31        };
32
33        let node = cursor.node();
34
35        if cursor.goto_first_child() || cursor.goto_next_sibling() {
36            return Some(node);
37        }
38
39        loop {
40            if !cursor.goto_parent() {
41                // If we can't go to the parent, the walk will be
42                // complete *after* the current node.
43                self.cursor = None;
44                break;
45            }
46
47            if cursor.goto_next_sibling() {
48                break;
49            }
50        }
51
52        Some(node)
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    #[test]
59    fn test_iter_is_total() {
60        let anchors = r#"
61jobs:
62  job1:
63    env: &env_vars # Define the anchor on first use
64      NODE_ENV: production
65      DATABASE_URL: ${{ secrets.DATABASE_URL }}
66    steps:
67      - run: echo "Using production settings"
68
69  job2:
70    env: *env_vars # Reuse the environment variables
71    steps:
72      - run: echo "Same environment variables here"
73        "#;
74
75        // NOTE(ww): These node counts will probably change if
76        // tree-sitter-yaml changes its node structure. Hopefully
77        // that doesn't happen often.
78        let testcases = &[
79            ("foo:", 9),
80            ("foo: # comment", 10),
81            ("foo: bar", 12),
82            ("foo: bar # comment", 13),
83            ("foo: []", 13),
84            ("foo: [] # comment", 14),
85            (anchors, 100),
86        ];
87
88        for (src, expected_count) in testcases {
89            let mut parser = tree_sitter::Parser::new();
90            parser
91                .set_language(&tree_sitter_yaml::LANGUAGE.into())
92                .expect("Error loading YAML grammar");
93            let tree = parser.parse(src, None).expect("Failed to parse source");
94
95            let node_count = tree.root_node().descendant_count();
96            let iter_count = super::TreeIter::new(&tree).count();
97
98            assert_eq!(node_count, *expected_count);
99            assert_eq!(node_count, iter_count);
100        }
101    }
102}
tree_sitter_iter/lib.rs

tree_sitter_iter/
lib.rs
