quickmark_core/
tree_sitter_walker.rs

1use tree_sitter::{Node, Tree};
2
3#[derive(Copy, Clone, Debug)]
4pub enum TraversalOrder {
5    PreOrder,
6    PostOrder,
7}
8
9#[derive(Debug)]
10pub struct TreeSitterWalker<'a> {
11    pub order: TraversalOrder,
12    pub tree: &'a Tree,
13}
14
15impl<'a> TreeSitterWalker<'a> {
16    pub fn new(tree: &'a Tree) -> Self {
17        Self {
18            tree,
19            order: TraversalOrder::PreOrder,
20        }
21    }
22
23    pub fn with_order(tree: &'a Tree, order: TraversalOrder) -> Self {
24        Self { tree, order }
25    }
26
27    pub fn walk(&self, mut callback: impl FnMut(Node)) {
28        let root = self.tree.root_node();
29        match self.order {
30            TraversalOrder::PreOrder => self.walk_pre_order(root, &mut callback),
31            TraversalOrder::PostOrder => self.walk_post_order(root, &mut callback),
32        }
33    }
34
35    #[allow(clippy::only_used_in_recursion)]
36    fn walk_pre_order(&self, node: Node, callback: &mut impl FnMut(Node)) {
37        callback(node);
38        for child in node.children(&mut node.walk()) {
39            self.walk_pre_order(child, callback);
40        }
41    }
42    #[allow(clippy::only_used_in_recursion)]
43    fn walk_post_order(&self, node: Node, callback: &mut impl FnMut(Node)) {
44        for child in node.children(&mut node.walk()) {
45            self.walk_post_order(child, callback);
46        }
47        callback(node);
48    }
49}