Skip to main content

diffler_core/syntax/
registry.rs

1//! Language registry: maps a file path to its tree-sitter grammar, a configured
2//! highlight configuration, and (where the grammar ships one) a tags query used
3//! for scope/definition lookup. Built once and reused.
4
5use std::collections::HashMap;
6use std::ops::Range;
7use std::path::Path;
8
9use tree_sitter::{Language, Parser, Query, QueryCursor, StreamingIterator, Tree};
10use tree_sitter_highlight::HighlightConfiguration;
11
12use crate::syntax::MAX_PARSE_BYTES;
13
14/// Capture names recognized during highlighting. A grammar capture like
15/// `function.method` resolves to the longest matching prefix here (`function`),
16/// so listing the general categories is enough to color every grammar.
17pub const HIGHLIGHT_NAMES: &[&str] = &[
18    "attribute",
19    "boolean",
20    "comment",
21    "conditional",
22    "constant",
23    "constant.builtin",
24    "constructor",
25    "field",
26    "function",
27    "function.builtin",
28    "keyword",
29    "label",
30    "number",
31    "operator",
32    "parameter",
33    "property",
34    "punctuation",
35    "punctuation.bracket",
36    "punctuation.delimiter",
37    "spell",
38    "storageclass",
39    "string",
40    "string.escape",
41    "string.special",
42    "tag",
43    "text.emphasis",
44    "text.literal",
45    "text.reference",
46    "text.strong",
47    "text.title",
48    "text.uri",
49    "type",
50    "type.builtin",
51    "variable",
52    "variable.builtin",
53    "variable.parameter",
54];
55
56pub struct LangEntry {
57    pub name: &'static str,
58    pub language: Language,
59    /// `None` when the grammar's highlight query failed to compile; the file
60    /// then renders plain instead of erroring.
61    pub config: Option<HighlightConfiguration>,
62    /// Definition/tags query for scope lookup; `None` when the grammar ships no
63    /// tags query.
64    pub tags: Option<Query>,
65}
66
67pub struct LanguageRegistry {
68    entries: Vec<LangEntry>,
69    by_ext: HashMap<&'static str, usize>,
70    by_name: HashMap<&'static str, usize>,
71    /// The inline markdown highlight query, applied by hand over the block
72    /// grammar's `(inline)` nodes: tree-sitter's generic injection does not
73    /// drive the split markdown grammar's inline pass.
74    markdown_inline_query: Option<Query>,
75}
76
77impl LanguageRegistry {
78    /// Build the registry with every bundled grammar, reused for the session.
79    // flat per-language registration table
80    #[allow(clippy::too_many_lines)]
81    pub fn build() -> Self {
82        let mut r = Self {
83            entries: Vec::new(),
84            by_ext: HashMap::new(),
85            by_name: HashMap::new(),
86            markdown_inline_query: None,
87        };
88
89        r.add(
90            "rust",
91            &["rs"],
92            tree_sitter_rust::LANGUAGE.into(),
93            tree_sitter_rust::HIGHLIGHTS_QUERY,
94            Some(tree_sitter_rust::TAGS_QUERY),
95        );
96        r.add(
97            "python",
98            &["py", "pyi"],
99            tree_sitter_python::LANGUAGE.into(),
100            tree_sitter_python::HIGHLIGHTS_QUERY,
101            Some(tree_sitter_python::TAGS_QUERY),
102        );
103        r.add(
104            "javascript",
105            &["js", "jsx", "mjs", "cjs"],
106            tree_sitter_javascript::LANGUAGE.into(),
107            &format!(
108                "{}\n{}",
109                tree_sitter_javascript::HIGHLIGHT_QUERY,
110                tree_sitter_javascript::JSX_HIGHLIGHT_QUERY
111            ),
112            Some(tree_sitter_javascript::TAGS_QUERY),
113        );
114        r.add(
115            "typescript",
116            &["ts", "mts", "cts"],
117            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
118            &format!(
119                "{}\n{}",
120                tree_sitter_javascript::HIGHLIGHT_QUERY,
121                tree_sitter_typescript::HIGHLIGHTS_QUERY
122            ),
123            Some(tree_sitter_typescript::TAGS_QUERY),
124        );
125        r.add(
126            "tsx",
127            &["tsx"],
128            tree_sitter_typescript::LANGUAGE_TSX.into(),
129            &format!(
130                "{}\n{}\n{}",
131                tree_sitter_javascript::HIGHLIGHT_QUERY,
132                tree_sitter_javascript::JSX_HIGHLIGHT_QUERY,
133                tree_sitter_typescript::HIGHLIGHTS_QUERY
134            ),
135            Some(tree_sitter_typescript::TAGS_QUERY),
136        );
137        r.add(
138            "go",
139            &["go"],
140            tree_sitter_go::LANGUAGE.into(),
141            tree_sitter_go::HIGHLIGHTS_QUERY,
142            Some(tree_sitter_go::TAGS_QUERY),
143        );
144        r.add(
145            "c",
146            &["c", "h"],
147            tree_sitter_c::LANGUAGE.into(),
148            tree_sitter_c::HIGHLIGHT_QUERY,
149            Some(tree_sitter_c::TAGS_QUERY),
150        );
151        r.add(
152            "cpp",
153            &["cpp", "cc", "cxx", "hpp", "hh", "hxx"],
154            tree_sitter_cpp::LANGUAGE.into(),
155            tree_sitter_cpp::HIGHLIGHT_QUERY,
156            Some(tree_sitter_cpp::TAGS_QUERY),
157        );
158        r.add(
159            "java",
160            &["java"],
161            tree_sitter_java::LANGUAGE.into(),
162            tree_sitter_java::HIGHLIGHTS_QUERY,
163            Some(tree_sitter_java::TAGS_QUERY),
164        );
165        r.add(
166            "c-sharp",
167            &["cs"],
168            tree_sitter_c_sharp::LANGUAGE.into(),
169            tree_sitter_c_sharp::HIGHLIGHTS_QUERY,
170            Some(tree_sitter_c_sharp::TAGS_QUERY),
171        );
172        r.add(
173            "ruby",
174            &["rb"],
175            tree_sitter_ruby::LANGUAGE.into(),
176            tree_sitter_ruby::HIGHLIGHTS_QUERY,
177            Some(tree_sitter_ruby::TAGS_QUERY),
178        );
179        r.add(
180            "php",
181            &["php"],
182            tree_sitter_php::LANGUAGE_PHP.into(),
183            tree_sitter_php::HIGHLIGHTS_QUERY,
184            Some(tree_sitter_php::TAGS_QUERY),
185        );
186        r.add(
187            "bash",
188            &["sh", "bash", "zsh"],
189            tree_sitter_bash::LANGUAGE.into(),
190            tree_sitter_bash::HIGHLIGHT_QUERY,
191            None,
192        );
193        r.add(
194            "json",
195            &["json"],
196            tree_sitter_json::LANGUAGE.into(),
197            tree_sitter_json::HIGHLIGHTS_QUERY,
198            None,
199        );
200        r.add(
201            "html",
202            &["html", "htm"],
203            tree_sitter_html::LANGUAGE.into(),
204            tree_sitter_html::HIGHLIGHTS_QUERY,
205            None,
206        );
207        r.add(
208            "css",
209            &["css"],
210            tree_sitter_css::LANGUAGE.into(),
211            tree_sitter_css::HIGHLIGHTS_QUERY,
212            None,
213        );
214        r.add(
215            "yaml",
216            &["yml", "yaml"],
217            tree_sitter_yaml::LANGUAGE.into(),
218            tree_sitter_yaml::HIGHLIGHTS_QUERY,
219            None,
220        );
221        // the grammar's own numeric patterns are guarded by Lua-style `%d`
222        // predicates that tree-sitter's regex engine never matches, leaving
223        // every number styled as a string; a later pattern wins, so this one
224        // restores number coloring
225        r.add(
226            "sql",
227            &["sql"],
228            tree_sitter_sequel::LANGUAGE.into(),
229            &format!(
230                "{}\n((literal) @number (#match? @number \"^[-+]?[0-9][0-9.]*$\"))\n",
231                tree_sitter_sequel::HIGHLIGHTS_QUERY
232            ),
233            None,
234        );
235        // The block grammar highlights headings/markers and injects fenced code
236        // into its own language; inline emphasis, code spans, and links come from
237        // the by-hand inline pass below.
238        r.register(
239            "markdown",
240            &["md", "markdown"],
241            tree_sitter_md::LANGUAGE.into(),
242            tree_sitter_md::HIGHLIGHT_QUERY_BLOCK,
243            tree_sitter_md::INJECTION_QUERY_BLOCK,
244            None,
245        );
246        let md_inline: Language = tree_sitter_md::INLINE_LANGUAGE.into();
247        r.markdown_inline_query =
248            Query::new(&md_inline, tree_sitter_md::HIGHLIGHT_QUERY_INLINE).ok();
249        r.register(
250            "markdown_inline",
251            &[],
252            md_inline,
253            tree_sitter_md::HIGHLIGHT_QUERY_INLINE,
254            "",
255            None,
256        );
257
258        r
259    }
260
261    fn add(
262        &mut self,
263        name: &'static str,
264        extensions: &'static [&'static str],
265        language: Language,
266        highlights: &str,
267        tags: Option<&str>,
268    ) {
269        self.register(name, extensions, language, highlights, "", tags);
270    }
271
272    fn register(
273        &mut self,
274        name: &'static str,
275        extensions: &'static [&'static str],
276        language: Language,
277        highlights: &str,
278        injections: &str,
279        tags: Option<&str>,
280    ) {
281        let config =
282            HighlightConfiguration::new(language.clone(), name, highlights, injections, "")
283                .ok()
284                .map(|mut c| {
285                    c.configure(HIGHLIGHT_NAMES);
286                    c
287                });
288        let tags = tags.and_then(|q| Query::new(&language, q).ok());
289        let idx = self.entries.len();
290        self.entries.push(LangEntry {
291            name,
292            language,
293            config,
294            tags,
295        });
296        self.by_name.insert(name, idx);
297        for ext in extensions {
298            self.by_ext.insert(ext, idx);
299        }
300    }
301
302    /// The entry whose grammar handles `path`, keyed by file extension.
303    pub fn for_path(&self, path: &str) -> Option<&LangEntry> {
304        let ext = Path::new(path).extension()?.to_str()?;
305        let &idx = self.by_ext.get(ext)?;
306        self.entries.get(idx)
307    }
308
309    /// The entry for a markdown fence token (`rust`, `py`, `c++`, ...), matched
310    /// by grammar name then extension.
311    pub fn for_token(&self, token: &str) -> Option<&LangEntry> {
312        let token = token.trim().to_ascii_lowercase();
313        let token = match token.as_str() {
314            "c++" => "cpp",
315            "c#" | "csharp" => "cs",
316            "shell" => "bash",
317            "golang" => "go",
318            other => other,
319        };
320        let &idx = self.by_name.get(token).or_else(|| self.by_ext.get(token))?;
321        self.entries.get(idx)
322    }
323
324    /// Highlight config for a tree-sitter injection language name (the inline
325    /// markdown grammar, or a fenced code block's language). `None` leaves the
326    /// injected region plain.
327    pub fn config_for_injection(&self, lang: &str) -> Option<&HighlightConfiguration> {
328        self.for_token(lang)?.config.as_ref()
329    }
330
331    /// Inline markdown captures (emphasis, code spans, links) as byte range plus
332    /// the recognized highlight name, narrowest span first so a first-match
333    /// renderer picks the most specific. The inline grammar parses only the
334    /// block grammar's `(inline)` node ranges, so block markers stay untouched.
335    pub fn markdown_inline_spans(&self, content: &str) -> Vec<(Range<usize>, &'static str)> {
336        if content.len() > MAX_PARSE_BYTES {
337            return Vec::new();
338        }
339        let (Some(query), Some(block), Some(inline)) = (
340            self.markdown_inline_query.as_ref(),
341            self.by_name
342                .get("markdown")
343                .and_then(|&i| self.entries.get(i)),
344            self.by_name
345                .get("markdown_inline")
346                .and_then(|&i| self.entries.get(i)),
347        ) else {
348            return Vec::new();
349        };
350
351        let mut bp = Parser::new();
352        if bp.set_language(&block.language).is_err() {
353            return Vec::new();
354        }
355        let Some(block_tree) = bp.parse(content, None) else {
356            return Vec::new();
357        };
358        let ranges = inline_node_ranges(&block_tree);
359        if ranges.is_empty() {
360            return Vec::new();
361        }
362
363        let mut ip = Parser::new();
364        if ip.set_included_ranges(&ranges).is_err() || ip.set_language(&inline.language).is_err() {
365            return Vec::new();
366        }
367        let Some(inline_tree) = ip.parse(content, None) else {
368            return Vec::new();
369        };
370
371        let names = query.capture_names();
372        let mut cursor = QueryCursor::new();
373        let mut spans = Vec::new();
374        let mut matches = cursor.matches(query, inline_tree.root_node(), content.as_bytes());
375        while let Some(m) = matches.next() {
376            for cap in m.captures {
377                let cname = names.get(cap.index as usize).copied().unwrap_or("");
378                if let Some(name) = recognized_highlight(cname) {
379                    spans.push((cap.node.byte_range(), name));
380                }
381            }
382        }
383        spans.sort_by_key(|(r, _)| r.end - r.start);
384        spans
385    }
386}
387
388/// Byte ranges of every `(inline)` node in a markdown block tree, in document
389/// order. These are the regions the inline grammar reparses.
390fn inline_node_ranges(tree: &Tree) -> Vec<tree_sitter::Range> {
391    let mut ranges = Vec::new();
392    let mut cursor = tree.walk();
393    loop {
394        let node = cursor.node();
395        if node.kind() == "inline" && node.end_byte() > node.start_byte() {
396            ranges.push(node.range());
397        }
398        if cursor.goto_first_child() {
399            continue;
400        }
401        loop {
402            if cursor.goto_next_sibling() {
403                break;
404            }
405            if !cursor.goto_parent() {
406                return ranges;
407            }
408        }
409    }
410}
411
412/// The longest `HIGHLIGHT_NAMES` entry that is a dotted prefix of `capture`,
413/// matching how tree-sitter resolves capture names to recognized highlights.
414fn recognized_highlight(capture: &str) -> Option<&'static str> {
415    HIGHLIGHT_NAMES
416        .iter()
417        .copied()
418        .filter(|name| {
419            capture == *name
420                || capture
421                    .strip_prefix(name)
422                    .is_some_and(|r| r.starts_with('.'))
423        })
424        .max_by_key(|name| name.len())
425}
426
427impl Default for LanguageRegistry {
428    fn default() -> Self {
429        Self::build()
430    }
431}