tree_house/
tree_cursor.rs

1use std::collections::VecDeque;
2
3use crate::tree_sitter::Node;
4use crate::{Layer, Syntax};
5
6pub struct TreeCursor<'tree> {
7    syntax: &'tree Syntax,
8    current: Layer,
9    cursor: tree_sitter::TreeCursor<'tree>,
10}
11
12impl<'tree> TreeCursor<'tree> {
13    pub(crate) fn new(syntax: &'tree Syntax) -> Self {
14        let cursor = syntax.tree().walk();
15
16        Self {
17            syntax,
18            current: syntax.root,
19            cursor,
20        }
21    }
22
23    pub fn node(&self) -> Node<'tree> {
24        self.cursor.node()
25    }
26
27    pub fn goto_parent(&mut self) -> bool {
28        if self.cursor.goto_parent() {
29            return true;
30        };
31
32        loop {
33            // Ascend to the parent layer if one exists.
34            let Some(parent) = self.syntax.layer(self.current).parent else {
35                return false;
36            };
37
38            self.current = parent;
39            if let Some(tree) = self.syntax.layer(self.current).tree() {
40                self.cursor = tree.walk();
41                break;
42            }
43        }
44
45        true
46    }
47
48    pub fn goto_parent_with<P>(&mut self, predicate: P) -> bool
49    where
50        P: Fn(&Node) -> bool,
51    {
52        while self.goto_parent() {
53            if predicate(&self.node()) {
54                return true;
55            }
56        }
57
58        false
59    }
60
61    pub fn goto_first_child(&mut self) -> bool {
62        let range = self.cursor.node().byte_range();
63        let layer = self.syntax.layer(self.current);
64        if let Some((layer, tree)) = layer
65            .injection_at_byte_idx(range.start)
66            .filter(|injection| injection.range.end >= range.end)
67            .and_then(|injection| {
68                Some((injection.layer, self.syntax.layer(injection.layer).tree()?))
69            })
70        {
71            // Switch to the child layer.
72            self.current = layer;
73            self.cursor = tree.walk();
74            return true;
75        }
76
77        self.cursor.goto_first_child()
78    }
79
80    pub fn goto_next_sibling(&mut self) -> bool {
81        self.cursor.goto_next_sibling()
82    }
83
84    pub fn goto_previous_sibling(&mut self) -> bool {
85        self.cursor.goto_previous_sibling()
86    }
87
88    pub fn reset_to_byte_range(&mut self, start: u32, end: u32) {
89        let (layer, tree) = self.syntax.layer_and_tree_for_byte_range(start, end);
90        self.current = layer;
91        self.cursor = tree.walk();
92
93        loop {
94            let node = self.cursor.node();
95            if start < node.start_byte() || end > node.end_byte() {
96                self.cursor.goto_parent();
97                break;
98            }
99            if self.cursor.goto_first_child_for_byte(start).is_none() {
100                break;
101            }
102        }
103    }
104
105    /// Returns an iterator over the children of the node the TreeCursor is on
106    /// at the time this is called.
107    pub fn children<'a>(&'a mut self) -> ChildIter<'a, 'tree> {
108        let parent = self.node();
109
110        ChildIter {
111            cursor: self,
112            parent,
113        }
114    }
115}
116
117pub struct ChildIter<'a, 'tree> {
118    cursor: &'a mut TreeCursor<'tree>,
119    parent: Node<'tree>,
120}
121
122impl<'tree> Iterator for ChildIter<'_, 'tree> {
123    type Item = Node<'tree>;
124
125    fn next(&mut self) -> Option<Self::Item> {
126        // first iteration, just visit the first child
127        if self.cursor.node() == self.parent {
128            self.cursor.goto_first_child().then(|| self.cursor.node())
129        } else {
130            self.cursor.goto_next_sibling().then(|| self.cursor.node())
131        }
132    }
133}
134
135impl<'cursor, 'tree> IntoIterator for &'cursor mut TreeCursor<'tree> {
136    type Item = Node<'tree>;
137    type IntoIter = TreeRecursiveWalker<'cursor, 'tree>;
138
139    fn into_iter(self) -> Self::IntoIter {
140        let mut queue = VecDeque::new();
141        let root = self.node();
142        queue.push_back(root.clone());
143
144        TreeRecursiveWalker {
145            cursor: self,
146            queue,
147            root,
148        }
149    }
150}
151
152pub struct TreeRecursiveWalker<'cursor, 'tree> {
153    cursor: &'cursor mut TreeCursor<'tree>,
154    queue: VecDeque<Node<'tree>>,
155    root: Node<'tree>,
156}
157
158impl<'tree> Iterator for TreeRecursiveWalker<'_, 'tree> {
159    type Item = Node<'tree>;
160
161    fn next(&mut self) -> Option<Self::Item> {
162        let current = self.cursor.node();
163
164        if current != self.root && self.cursor.goto_next_sibling() {
165            self.queue.push_back(current);
166            return Some(self.cursor.node());
167        }
168
169        while let Some(queued) = self.queue.pop_front() {
170            self.cursor.cursor.reset(&queued);
171
172            if !self.cursor.goto_first_child() {
173                continue;
174            }
175
176            return Some(self.cursor.node());
177        }
178
179        None
180    }
181}
tree_house/tree_cursor.rs

tree_house/
tree_cursor.rs
