tree_house/
locals.rs

1use std::{
2    borrow::Cow,
3    ops::{Index, IndexMut},
4};
5
6use hashbrown::HashMap;
7use kstring::KString;
8use ropey::RopeSlice;
9use tree_sitter::{Capture, InactiveQueryCursor};
10
11use crate::{LanguageConfig, LanguageLoader, Layer, Range, Syntax, TREE_SITTER_MATCH_LIMIT};
12
13#[derive(Debug, PartialEq, Eq, Hash, Clone, Copy)]
14pub struct Scope(u32);
15
16impl Scope {
17    const ROOT: Scope = Scope(0);
18    fn idx(self) -> usize {
19        self.0 as usize
20    }
21}
22
23#[derive(Debug)]
24pub struct Locals {
25    scopes: Vec<ScopeData>,
26}
27
28impl Default for Locals {
29    fn default() -> Self {
30        let mut scopes = Vec::with_capacity(4);
31        scopes.push(ScopeData {
32            definitions: HashMap::new(),
33            range: 0..u32::MAX,
34            inherit: false,
35            children: Vec::new(),
36            parent: None,
37        });
38
39        Self { scopes }
40    }
41}
42
43impl Locals {
44    fn push(&mut self, scope: ScopeData) -> Scope {
45        let new_scope_id = Scope(self.scopes.len() as u32);
46        let parent = scope
47            .parent
48            .expect("push cannot be used for the root layer");
49        self[parent].children.push(new_scope_id);
50        self.scopes.push(scope);
51        new_scope_id
52    }
53
54    pub fn lookup_reference(&self, mut scope: Scope, name: &str) -> Option<&Definition> {
55        loop {
56            let scope_data = &self[scope];
57            if let Some(def) = scope_data.definitions.get(name) {
58                return Some(def);
59            }
60            if !scope_data.inherit {
61                break;
62            }
63            scope = scope_data.parent?;
64        }
65
66        None
67    }
68
69    pub fn scope_cursor(&self, pos: u32) -> ScopeCursor<'_> {
70        let mut scope = Scope::ROOT;
71        let mut scope_stack = Vec::with_capacity(8);
72        loop {
73            let scope_data = &self[scope];
74            let child_idx = scope_data
75                .children
76                .partition_point(|&child| self[child].range.end < pos);
77            scope_stack.push((scope, child_idx as u32));
78            let Some(&child) = scope_data.children.get(child_idx) else {
79                break;
80            };
81            if pos < self[child].range.start {
82                break;
83            }
84            scope = child;
85        }
86        ScopeCursor {
87            locals: self,
88            scope_stack,
89        }
90    }
91}
92
93impl Index<Scope> for Locals {
94    type Output = ScopeData;
95
96    fn index(&self, scope: Scope) -> &Self::Output {
97        &self.scopes[scope.idx()]
98    }
99}
100
101impl IndexMut<Scope> for Locals {
102    fn index_mut(&mut self, scope: Scope) -> &mut Self::Output {
103        &mut self.scopes[scope.idx()]
104    }
105}
106
107#[derive(Debug)]
108pub struct ScopeCursor<'a> {
109    pub locals: &'a Locals,
110    scope_stack: Vec<(Scope, u32)>,
111}
112
113impl ScopeCursor<'_> {
114    pub fn advance(&mut self, to: u32) -> Scope {
115        let (mut active_scope, mut child_idx) = self.scope_stack.pop().unwrap();
116        loop {
117            let scope_data = &self.locals[active_scope];
118            if to < scope_data.range.end {
119                break;
120            }
121            (active_scope, child_idx) = self.scope_stack.pop().unwrap();
122            child_idx += 1;
123        }
124        'outer: loop {
125            let scope_data = &self.locals[active_scope];
126            loop {
127                let Some(&child) = scope_data.children.get(child_idx as usize) else {
128                    break 'outer;
129                };
130                if self.locals[child].range.start > to {
131                    break 'outer;
132                }
133                if to < self.locals[child].range.end {
134                    self.scope_stack.push((active_scope, child_idx));
135                    active_scope = child;
136                    child_idx = 0;
137                    break;
138                }
139                child_idx += 1;
140            }
141        }
142        self.scope_stack.push((active_scope, child_idx));
143        active_scope
144    }
145
146    pub fn current_scope(&self) -> Scope {
147        // The root scope is always active so `scope_stack` is never empty.
148        self.scope_stack.last().unwrap().0
149    }
150}
151
152#[derive(Debug)]
153pub struct Definition {
154    pub capture: Capture,
155    pub range: Range,
156}
157
158#[derive(Debug)]
159pub struct ScopeData {
160    definitions: HashMap<KString, Definition>,
161    range: Range,
162    inherit: bool,
163    /// A list of sorted, non-overlapping child scopes.
164    ///
165    /// See the docs of the `Locals` type: locals information is laid out like a tree - similar
166    /// to injections - per injection layer.
167    children: Vec<Scope>,
168    parent: Option<Scope>,
169}
170
171impl Syntax {
172    pub(crate) fn run_local_query(
173        &mut self,
174        layer: Layer,
175        source: RopeSlice<'_>,
176        loader: &impl LanguageLoader,
177    ) {
178        let layer_data = &mut self.layer_mut(layer);
179        let Some(LanguageConfig {
180            ref injection_query,
181            ..
182        }) = loader.get_config(layer_data.language)
183        else {
184            return;
185        };
186        let definition_captures = injection_query.local_definition_captures.load();
187        if definition_captures.is_empty() {
188            return;
189        }
190
191        let root = layer_data.parse_tree.as_ref().unwrap().root_node();
192        let mut cursor = InactiveQueryCursor::new(0..u32::MAX, TREE_SITTER_MATCH_LIMIT)
193            .execute_query(&injection_query.local_query, &root, source);
194        let mut locals = Locals::default();
195        let mut scope = Scope::ROOT;
196
197        while let Some((query_match, node_idx)) = cursor.next_matched_node() {
198            let matched_node = query_match.matched_node(node_idx);
199            let range = matched_node.node.byte_range();
200            let capture = matched_node.capture;
201
202            while range.start >= locals[scope].range.end {
203                scope = locals[scope].parent.expect("root node covers entire range");
204            }
205
206            if Some(capture) == injection_query.local_scope_capture {
207                scope = locals.push(ScopeData {
208                    definitions: HashMap::new(),
209                    range: matched_node.node.byte_range(),
210                    inherit: !injection_query
211                        .not_scope_inherits
212                        .contains(&query_match.pattern()),
213                    children: Vec::new(),
214                    parent: Some(scope),
215                });
216            } else if definition_captures.contains_key(&capture) {
217                let text = match source
218                    .byte_slice(range.start as usize..range.end as usize)
219                    .into()
220                {
221                    Cow::Borrowed(inner) => KString::from_ref(inner),
222                    Cow::Owned(inner) => KString::from_string(inner),
223                };
224                locals[scope]
225                    .definitions
226                    .insert(text, Definition { capture, range });
227            }
228            // NOTE: `local.reference` captures are handled by the highlighter and are not
229            // considered during parsing.
230        }
231
232        layer_data.locals = locals;
233    }
234}
235
236#[cfg(test)]
237mod test {
238    use super::*;
239
240    #[test]
241    fn cursor() {
242        let mut locals = Locals::default();
243        let scope1 = locals.push(ScopeData {
244            definitions: Default::default(),
245            range: 5..105,
246            inherit: true,
247            // NOTE: the subsequent call to `push` below will add scope2 to scope1's children.
248            children: Default::default(),
249            parent: Some(Scope::ROOT),
250        });
251        let scope2 = locals.push(ScopeData {
252            definitions: Default::default(),
253            range: 10..100,
254            inherit: true,
255            children: Default::default(),
256            parent: Some(scope1),
257        });
258
259        let mut cursor = locals.scope_cursor(0);
260        assert_eq!(cursor.current_scope(), Scope::ROOT);
261        assert_eq!(cursor.advance(3), Scope::ROOT);
262        assert_eq!(cursor.advance(5), scope1);
263        assert_eq!(cursor.advance(8), scope1);
264        assert_eq!(cursor.advance(10), scope2);
265        assert_eq!(cursor.advance(50), scope2);
266        assert_eq!(cursor.advance(100), scope1);
267        assert_eq!(cursor.advance(105), Scope::ROOT);
268        assert_eq!(cursor.advance(110), Scope::ROOT);
269
270        let mut cursor = locals.scope_cursor(8);
271        assert_eq!(cursor.current_scope(), scope1);
272        assert_eq!(cursor.advance(10), scope2);
273        assert_eq!(cursor.advance(100), scope1);
274        assert_eq!(cursor.advance(110), Scope::ROOT);
275
276        let mut cursor = locals.scope_cursor(10);
277        assert_eq!(cursor.current_scope(), scope2);
278        assert_eq!(cursor.advance(100), scope1);
279        assert_eq!(cursor.advance(110), Scope::ROOT);
280    }
281}
tree_house/locals.rs

tree_house/
locals.rs
