1use streaming_iterator::StreamingIterator;
7use tree_sitter::{Language, Parser, Query, QueryCursor, Tree};
8
9use crate::highlight::{self, TokenKind};
10
11pub 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 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 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 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 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 self.tree = None;
200 self.tokens.clear();
201 self.active = false;
202 }
203 }
204 }
205
206 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 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 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
280fn 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 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 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 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}