tree_house/
tree_cursor.rs

1use crate::tree_sitter::Node;
2use crate::{Layer, Syntax};
3
4pub struct TreeCursor<'tree> {
5    syntax: &'tree Syntax,
6    current: Layer,
7    cursor: tree_sitter::TreeCursor<'tree>,
8}
9
10impl<'tree> TreeCursor<'tree> {
11    pub(crate) fn new(syntax: &'tree Syntax) -> Self {
12        let cursor = syntax.tree().walk();
13
14        Self {
15            syntax,
16            current: syntax.root,
17            cursor,
18        }
19    }
20
21    pub fn node(&self) -> Node<'tree> {
22        self.cursor.node()
23    }
24
25    pub fn goto_parent(&mut self) -> bool {
26        if self.cursor.goto_parent() {
27            return true;
28        };
29
30        loop {
31            // Ascend to the parent layer if one exists.
32            let Some(parent) = self.syntax.layer(self.current).parent else {
33                return false;
34            };
35
36            self.current = parent;
37            if let Some(tree) = self.syntax.layer(self.current).tree() {
38                self.cursor = tree.walk();
39                break;
40            }
41        }
42
43        true
44    }
45
46    pub fn goto_parent_with<P>(&mut self, predicate: P) -> bool
47    where
48        P: Fn(&Node) -> bool,
49    {
50        while self.goto_parent() {
51            if predicate(&self.node()) {
52                return true;
53            }
54        }
55
56        false
57    }
58
59    pub fn goto_first_child(&mut self) -> bool {
60        let range = self.cursor.node().byte_range();
61        let layer = self.syntax.layer(self.current);
62        if let Some((layer, tree)) = layer
63            .injection_at_byte_idx(range.start)
64            .filter(|injection| injection.range.end >= range.end)
65            .and_then(|injection| {
66                Some((injection.layer, self.syntax.layer(injection.layer).tree()?))
67            })
68        {
69            // Switch to the child layer.
70            self.current = layer;
71            self.cursor = tree.walk();
72            return true;
73        }
74
75        self.cursor.goto_first_child()
76    }
77
78    pub fn goto_next_sibling(&mut self) -> bool {
79        self.cursor.goto_next_sibling()
80    }
81
82    pub fn goto_previous_sibling(&mut self) -> bool {
83        self.cursor.goto_previous_sibling()
84    }
85
86    pub fn reset_to_byte_range(&mut self, start: u32, end: u32) {
87        let (layer, tree) = self.syntax.layer_and_tree_for_byte_range(start, end);
88        self.current = layer;
89        self.cursor = tree.walk();
90
91        loop {
92            let node = self.cursor.node();
93            if start < node.start_byte() || end > node.end_byte() {
94                self.cursor.goto_parent();
95                break;
96            }
97            if self.cursor.goto_first_child_for_byte(start).is_none() {
98                break;
99            }
100        }
101    }
102
103    /// Returns an iterator over the children of the node the TreeCursor is on
104    /// at the time this is called.
105    pub fn children<'a>(&'a mut self) -> ChildIter<'a, 'tree> {
106        let parent = self.node();
107
108        ChildIter {
109            cursor: self,
110            parent,
111        }
112    }
113}
114
115pub struct ChildIter<'a, 'tree> {
116    cursor: &'a mut TreeCursor<'tree>,
117    parent: Node<'tree>,
118}
119
120impl<'tree> Iterator for ChildIter<'_, 'tree> {
121    type Item = Node<'tree>;
122
123    fn next(&mut self) -> Option<Self::Item> {
124        // first iteration, just visit the first child
125        if self.cursor.node() == self.parent {
126            self.cursor.goto_first_child().then(|| self.cursor.node())
127        } else {
128            self.cursor.goto_next_sibling().then(|| self.cursor.node())
129        }
130    }
131}