Skip to main content

xei_core/
syntax.rs

1//! Tree-sitter syntax highlighting via **highlight queries** (`highlights.scm`).
2//!
3//! Tree-sitter columns are **byte** offsets; the editor uses **char** indices.
4//! Query captures map to [`TokenKind`] through [`crate::highlight::from_capture`].
5
6use streaming_iterator::StreamingIterator;
7use tree_sitter::{Language, Parser, Query, QueryCursor, Tree};
8
9use crate::highlight::{self, TokenKind};
10
11/// One highlight span: (kind, start_col, end_col, row) — char columns, end exclusive.
12pub type HlToken = (TokenKind, usize, usize, usize);
13
14struct LangBundle {
15    parser: Parser,
16    query: Query,
17}
18
19pub struct SyntaxEngine {
20    rust: Option<LangBundle>,
21    python: Option<LangBundle>,
22    javascript: Option<LangBundle>,
23    typescript: Option<LangBundle>,
24    tsx: Option<LangBundle>,
25    c: Option<LangBundle>,
26    go: Option<LangBundle>,
27    bash: Option<LangBundle>,
28    json: Option<LangBundle>,
29    tree: Option<Tree>,
30    last_ext: String,
31    last_len: usize,
32    last_fingerprint: u64,
33    /// Query-based tokens (char columns)
34    pub tokens: Vec<HlToken>,
35    pub active: bool,
36}
37
38impl Default for SyntaxEngine {
39    fn default() -> Self {
40        Self {
41            rust: make_lang(
42                tree_sitter_rust::LANGUAGE.into(),
43                tree_sitter_rust::HIGHLIGHTS_QUERY,
44            ),
45            python: make_lang(
46                tree_sitter_python::LANGUAGE.into(),
47                tree_sitter_python::HIGHLIGHTS_QUERY,
48            ),
49            javascript: make_lang(
50                tree_sitter_javascript::LANGUAGE.into(),
51                tree_sitter_javascript::HIGHLIGHT_QUERY,
52            ),
53            typescript: make_lang(
54                tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
55                tree_sitter_typescript::HIGHLIGHTS_QUERY,
56            ),
57            tsx: make_lang(
58                tree_sitter_typescript::LANGUAGE_TSX.into(),
59                // TSX uses the same highlights as TS + JSX patterns when available
60                tree_sitter_typescript::HIGHLIGHTS_QUERY,
61            ),
62            c: make_lang(
63                tree_sitter_c::LANGUAGE.into(),
64                tree_sitter_c::HIGHLIGHT_QUERY,
65            ),
66            go: make_lang(tree_sitter_go::LANGUAGE.into(), tree_sitter_go::HIGHLIGHTS_QUERY),
67            bash: make_lang(
68                tree_sitter_bash::LANGUAGE.into(),
69                tree_sitter_bash::HIGHLIGHT_QUERY,
70            ),
71            json: make_lang(
72                tree_sitter_json::LANGUAGE.into(),
73                tree_sitter_json::HIGHLIGHTS_QUERY,
74            ),
75            tree: None,
76            last_ext: String::new(),
77            last_len: 0,
78            last_fingerprint: 0,
79            tokens: Vec::new(),
80            active: false,
81        }
82    }
83}
84
85fn make_lang(language: Language, source: &str) -> Option<LangBundle> {
86    let mut parser = Parser::new();
87    if parser.set_language(&language).is_err() {
88        return None;
89    }
90    let query = Query::new(&language, source).ok()?;
91    Some(LangBundle { parser, query })
92}
93
94impl SyntaxEngine {
95    pub fn new() -> Self {
96        Self::default()
97    }
98
99    pub fn parse(&mut self, text: &str, ext: Option<&str>) {
100        let ext_str = ext.unwrap_or("");
101        let fingerprint = fingerprint_text(text);
102
103        // Skip full work when content unchanged
104        if self.active
105            && self.last_ext == ext_str
106            && self.last_len == text.len()
107            && self.last_fingerprint == fingerprint
108            && !self.tokens.is_empty()
109        {
110            return;
111        }
112
113        let bundle = match ext {
114            Some("rs") => self.rust.as_mut(),
115            Some("py" | "pyi") => self.python.as_mut(),
116            Some("js" | "mjs" | "cjs") => self.javascript.as_mut(),
117            Some("jsx") => self.javascript.as_mut(),
118            Some("ts" | "mts" | "cts") => self.typescript.as_mut(),
119            Some("tsx") => self.tsx.as_mut().or(self.typescript.as_mut()),
120            Some("c" | "h") => self.c.as_mut(),
121            Some("cpp" | "hpp" | "cc" | "cxx" | "hh" | "hxx") => self.c.as_mut(),
122            Some("go") => self.go.as_mut(),
123            Some("sh" | "bash" | "zsh") => self.bash.as_mut(),
124            Some("json" | "jsonc") => self.json.as_mut(),
125            _ => {
126                self.tokens.clear();
127                self.tree = None;
128                self.last_ext.clear();
129                self.last_len = 0;
130                self.last_fingerprint = 0;
131                self.active = false;
132                return;
133            }
134        };
135
136        let Some(bundle) = bundle else {
137            self.tokens.clear();
138            self.active = false;
139            return;
140        };
141
142        self.active = true;
143        let len = text.len();
144
145        // IMPORTANT: never pass a stale Tree without `tree.edit(...)`.
146        // Incremental parse without edit panics inside tree-sitter.
147        // Drop any previous tree and full-reparse; wrap ALL ts calls in
148        // catch_unwind so a binding panic never kills the editor process.
149        self.tree = None;
150        self.tokens.clear();
151        self.last_ext = ext_str.to_string();
152        self.last_len = len;
153        self.last_fingerprint = fingerprint;
154
155        let source = text.as_bytes();
156        let lines: Vec<&str> = text.split('\n').collect();
157        let capture_names = bundle.query.capture_names().to_vec();
158
159        let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
160            let Some(tree) = bundle.parser.parse(text, None) else {
161                return None;
162            };
163            let root = tree.root_node();
164            let mut cursor = QueryCursor::new();
165            let mut tokens: Vec<HlToken> = Vec::new();
166            let mut matches = cursor.matches(&bundle.query, root, source);
167            while let Some(m) = matches.next() {
168                for cap in m.captures {
169                    let name = capture_names
170                        .get(cap.index as usize)
171                        .copied()
172                        .unwrap_or("");
173                    let Some(kind) = highlight::from_capture(name) else {
174                        continue;
175                    };
176                    let node = cap.node;
177                    let end_byte = node.end_byte();
178                    if end_byte > source.len() || node.start_byte() > end_byte {
179                        continue;
180                    }
181                    push_node_tokens(node, &lines, kind, &mut tokens);
182                }
183            }
184            tokens.sort_by_key(|(_, st, ed, row)| (*row, ed.saturating_sub(*st), *st));
185            Some((tree, tokens))
186        }));
187
188        match result {
189            Ok(Some((tree, tokens))) => {
190                self.tree = Some(tree);
191                self.tokens = tokens;
192            }
193            Ok(None) => {
194                self.tree = None;
195                self.tokens.clear();
196            }
197            Err(_) => {
198                // tree-sitter panicked — stay alive with no highlight
199                self.tree = None;
200                self.tokens.clear();
201                self.active = false;
202            }
203        }
204    }
205
206    /// Contiguous slice of highlight tokens on `row`. `tokens` is kept sorted by
207    /// row (see the `sort_by_key` in `parse`), so this is an O(log n) binary
208    /// search instead of the O(n) whole-file filter the renderer used to run for
209    /// every visible row on every frame.
210    pub fn tokens_for_row(&self, row: usize) -> &[HlToken] {
211        let lo = self.tokens.partition_point(|t| t.3 < row);
212        let hi = self.tokens.partition_point(|t| t.3 <= row);
213        &self.tokens[lo..hi]
214    }
215}
216
217fn push_node_tokens(
218    node: tree_sitter::Node,
219    lines: &[&str],
220    kind: TokenKind,
221    tokens: &mut Vec<HlToken>,
222) {
223    let start = node.start_position();
224    let end = node.end_position();
225
226    // Safety: skip huge multi-line non-comment/string spans
227    if start.row != end.row && !matches!(kind, TokenKind::Comment | TokenKind::String) {
228        return;
229    }
230
231    if start.row == end.row {
232        if let Some(line) = lines.get(start.row) {
233            let scol = byte_col_to_char_col(line, start.column);
234            let ecol = byte_col_to_char_col(line, end.column);
235            if scol < ecol {
236                tokens.push((kind, scol, ecol, start.row));
237            }
238        }
239        return;
240    }
241
242    // Multi-line comments / strings
243    if let Some(line) = lines.get(start.row) {
244        let scol = byte_col_to_char_col(line, start.column);
245        let ecol = line.chars().count();
246        if scol < ecol {
247            tokens.push((kind, scol, ecol, start.row));
248        }
249    }
250    for row in start.row + 1..end.row {
251        if let Some(line) = lines.get(row) {
252            let ecol = line.chars().count();
253            if ecol > 0 {
254                tokens.push((kind, 0, ecol, row));
255            }
256        }
257    }
258    if let Some(line) = lines.get(end.row) {
259        let ecol = byte_col_to_char_col(line, end.column);
260        if ecol > 0 {
261            tokens.push((kind, 0, ecol, end.row));
262        }
263    }
264}
265
266fn byte_col_to_char_col(line: &str, byte_col: usize) -> usize {
267    if byte_col == 0 {
268        return 0;
269    }
270    if byte_col >= line.len() {
271        return line.chars().count();
272    }
273    let mut idx = byte_col;
274    while idx > 0 && !line.is_char_boundary(idx) {
275        idx -= 1;
276    }
277    line.get(..idx).map(|s| s.chars().count()).unwrap_or(0)
278}
279
280/// Full-content FNV-1a fingerprint (skip re-query only when bytes are identical).
281fn fingerprint_text(text: &str) -> u64 {
282    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
283    for &b in text.as_bytes() {
284        h ^= b as u64;
285        h = h.wrapping_mul(0x0100_0000_01b3);
286    }
287    // Mix length so empty vs non-empty always differ
288    for b in text.len().to_le_bytes() {
289        h ^= b as u64;
290        h = h.wrapping_mul(0x0100_0000_01b3);
291    }
292    h
293}
294
295#[cfg(test)]
296mod tests {
297    use super::*;
298
299    #[test]
300    fn byte_to_char_ascii() {
301        assert_eq!(byte_col_to_char_col("hello", 0), 0);
302        assert_eq!(byte_col_to_char_col("hello", 5), 5);
303    }
304
305    #[test]
306    fn byte_to_char_cjk() {
307        let line = "a한b";
308        assert_eq!(byte_col_to_char_col(line, 0), 0);
309        assert_eq!(byte_col_to_char_col(line, 1), 1);
310        assert_eq!(byte_col_to_char_col(line, 4), 2);
311    }
312
313    #[test]
314    fn parse_rust_produces_tokens() {
315        let mut eng = SyntaxEngine::new();
316        eng.parse("fn main() { let x = 1; }", Some("rs"));
317        assert!(eng.active);
318        assert!(!eng.tokens.is_empty(), "expected query tokens");
319    }
320
321    #[test]
322    fn rust_does_not_paint_whole_function_as_one_token() {
323        let mut eng = SyntaxEngine::new();
324        let src = "fn main() {\n    let x = 42;\n    let s = \"hi\";\n}\n";
325        eng.parse(src, Some("rs"));
326        let first_line_len = src.lines().next().unwrap().chars().count();
327        let paints_whole_line = eng.tokens.iter().any(|(k, st, ed, row)| {
328            *row == 0 && *st == 0 && *ed >= first_line_len && matches!(k, TokenKind::Keyword)
329        });
330        assert!(
331            !paints_whole_line,
332            "keyword token painted entire first line: {:?}",
333            eng.tokens
334        );
335        let has_number = eng
336            .tokens
337            .iter()
338            .any(|(k, _, _, _)| matches!(k, TokenKind::Number));
339        let has_string = eng
340            .tokens
341            .iter()
342            .any(|(k, _, _, _)| matches!(k, TokenKind::String));
343        assert!(
344            has_number || has_string,
345            "expected number/string tokens, got {:?}",
346            eng.tokens
347        );
348    }
349
350    #[test]
351    fn rust_highlights_fn_and_function_name() {
352        let mut eng = SyntaxEngine::new();
353        eng.parse("fn main() {}", Some("rs"));
354        assert!(eng.active);
355        assert!(!eng.tokens.is_empty());
356    }
357
358    #[test]
359    fn python_query_active() {
360        let mut eng = SyntaxEngine::new();
361        eng.parse("def foo(x):\n    return x + 1\n", Some("py"));
362        assert!(eng.active);
363        assert!(!eng.tokens.is_empty());
364    }
365
366    #[test]
367    fn go_and_json_query_active() {
368        let mut eng = SyntaxEngine::new();
369        eng.parse("package main\nfunc Hello() {}\n", Some("go"));
370        assert!(eng.active);
371        assert!(!eng.tokens.is_empty());
372        eng.parse(r#"{"a": 1, "b": "x"}"#, Some("json"));
373        assert!(eng.active);
374        assert!(!eng.tokens.is_empty());
375    }
376
377    #[test]
378    fn skip_reparse_when_unchanged() {
379        let mut eng = SyntaxEngine::new();
380        eng.parse("fn a() {}", Some("rs"));
381        let n = eng.tokens.len();
382        eng.parse("fn a() {}", Some("rs"));
383        assert_eq!(eng.tokens.len(), n);
384    }
385
386    #[test]
387    fn rapid_edits_do_not_panic() {
388        // Regression: incremental parse without tree.edit() panicked with
389        // "range start index N out of range for slice of length M".
390        let mut eng = SyntaxEngine::new();
391        let mut src = String::from("fn main() {\n    let x = 1;\n}\n");
392        eng.parse(&src, Some("rs"));
393        for i in 0..80 {
394            src.insert(src.len().saturating_sub(2), char::from(b'a' + (i % 26) as u8));
395            eng.parse(&src, Some("rs"));
396            // delete a char near the middle
397            if src.len() > 10 {
398                let mid = src.len() / 2;
399                if src.is_char_boundary(mid) {
400                    src.remove(mid);
401                }
402                eng.parse(&src, Some("rs"));
403            }
404        }
405        assert!(eng.active || eng.tokens.is_empty());
406    }
407
408    #[test]
409    fn switch_language_and_edit() {
410        let mut eng = SyntaxEngine::new();
411        eng.parse("fn foo() {}", Some("rs"));
412        eng.parse("def foo():\n  pass\n", Some("py"));
413        eng.parse("fn bar() { let y = 2; }", Some("rs"));
414        assert!(eng.tokens.iter().any(|t| t.0 == TokenKind::Keyword || t.0 == TokenKind::Function));
415    }
416}