Skip to main content

md_tui/highlight/
mod.rs

1use std::cell::RefCell;
2use std::collections::HashMap;
3
4use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter};
5
6use ratatui::style::Color;
7
8#[allow(dead_code)]
9static HIGHLIGHT_NAMES: [&str; 18] = [
10    "attribute",
11    "constant",
12    "function.builtin",
13    "function",
14    "keyword",
15    "operator",
16    "property",
17    "punctuation",
18    "punctuation.bracket",
19    "punctuation.delimiter",
20    "string",
21    "string.special",
22    "tag",
23    "type",
24    "type.builtin",
25    "variable",
26    "variable.builtin",
27    "variable.parameter",
28];
29
30pub static COLOR_MAP: [Color; 18] = [
31    Color::Yellow,
32    Color::Yellow,
33    Color::Green,
34    Color::Green,
35    Color::Red,
36    Color::Red,
37    Color::Blue,
38    Color::Blue,
39    Color::Blue,
40    Color::Blue,
41    Color::Magenta,
42    Color::Magenta,
43    Color::Cyan,
44    Color::Cyan,
45    Color::Cyan,
46    Color::Reset,
47    Color::Reset,
48    Color::Reset,
49];
50
51#[derive(Debug)]
52pub enum HighlightInfo {
53    Highlighted(Vec<HighlightEvent>),
54    Mermaid,
55    Unhighlighted,
56}
57
58// With every `tree-sitter-*` feature disabled all numbered match arms gate
59// out, leaving only the early-return arms — that makes the trailing match
60// reachable only with at least one language feature, which is exactly the
61// supported configuration.
62#[allow(unused_variables, unreachable_code)]
63#[must_use]
64pub fn highlight_code(language: &str, lines: &[u8]) -> HighlightInfo {
65    // The mapping from alias to (tree-sitter language, query) lives inline
66    // here because each entry needs its own `#[cfg]` and the upstream crates
67    // disagree on the query constant name (HIGHLIGHTS_QUERY vs HIGHLIGHT_QUERY)
68    // and on the language constant name. The cache keyed by lang_name avoids
69    // rebuilding the `HighlightConfiguration` on every call.
70    let result: Result<Vec<HighlightEvent>, String> = match language {
71        #[cfg(feature = "tree-sitter-bash")]
72        "bash" | "sh" => highlight_with_language(
73            lines,
74            tree_sitter_bash::LANGUAGE.into(),
75            "bash",
76            tree_sitter_bash::HIGHLIGHT_QUERY,
77        ),
78
79        #[cfg(feature = "tree-sitter-c")]
80        "c" => highlight_with_language(
81            lines,
82            tree_sitter_c::LANGUAGE.into(),
83            "c",
84            tree_sitter_c::HIGHLIGHT_QUERY,
85        ),
86
87        #[cfg(feature = "tree-sitter-cpp")]
88        "cpp" => highlight_with_language(
89            lines,
90            tree_sitter_cpp::LANGUAGE.into(),
91            "cpp",
92            tree_sitter_cpp::HIGHLIGHT_QUERY,
93        ),
94
95        #[cfg(feature = "tree-sitter-css")]
96        "css" => highlight_with_language(
97            lines,
98            tree_sitter_css::LANGUAGE.into(),
99            "css",
100            tree_sitter_css::HIGHLIGHTS_QUERY,
101        ),
102
103        #[cfg(feature = "tree-sitter-diff")]
104        "diff" | "patch" => highlight_with_language(
105            lines,
106            tree_sitter_diff::LANGUAGE.into(),
107            "diff",
108            tree_sitter_diff::HIGHLIGHTS_QUERY,
109        ),
110
111        #[cfg(feature = "tree-sitter-elixir")]
112        "elixir" => highlight_with_language(
113            lines,
114            tree_sitter_elixir::LANGUAGE.into(),
115            "elixir",
116            tree_sitter_elixir::HIGHLIGHTS_QUERY,
117        ),
118
119        #[cfg(feature = "tree-sitter-go")]
120        "go" => highlight_with_language(
121            lines,
122            tree_sitter_go::LANGUAGE.into(),
123            "go",
124            tree_sitter_go::HIGHLIGHTS_QUERY,
125        ),
126
127        #[cfg(feature = "tree-sitter-html")]
128        "html" => highlight_with_language(
129            lines,
130            tree_sitter_html::LANGUAGE.into(),
131            "html",
132            tree_sitter_html::HIGHLIGHTS_QUERY,
133        ),
134
135        #[cfg(feature = "tree-sitter-java")]
136        "java" => highlight_with_language(
137            lines,
138            tree_sitter_java::LANGUAGE.into(),
139            "java",
140            tree_sitter_java::HIGHLIGHTS_QUERY,
141        ),
142
143        #[cfg(feature = "tree-sitter-javascript")]
144        "javascript" | "js" => highlight_with_language(
145            lines,
146            tree_sitter_javascript::LANGUAGE.into(),
147            "javascript",
148            tree_sitter_javascript::HIGHLIGHT_QUERY,
149        ),
150
151        #[cfg(feature = "tree-sitter-json")]
152        "json" => highlight_with_language(
153            lines,
154            tree_sitter_json::LANGUAGE.into(),
155            "json",
156            tree_sitter_json::HIGHLIGHTS_QUERY,
157        ),
158
159        #[cfg(feature = "tree-sitter-lua")]
160        "lua" => highlight_with_language(
161            lines,
162            tree_sitter_lua::LANGUAGE.into(),
163            "lua",
164            tree_sitter_lua::HIGHLIGHTS_QUERY,
165        ),
166
167        #[cfg(feature = "tree-sitter-ocaml")]
168        "ocaml" => highlight_with_language(
169            lines,
170            tree_sitter_ocaml::LANGUAGE_OCAML_TYPE.into(),
171            "ocaml",
172            tree_sitter_ocaml::HIGHLIGHTS_QUERY,
173        ),
174
175        #[cfg(feature = "tree-sitter-php")]
176        "php" => highlight_with_language(
177            lines,
178            tree_sitter_php::LANGUAGE_PHP.into(),
179            "php",
180            tree_sitter_php::HIGHLIGHTS_QUERY,
181        ),
182
183        #[cfg(feature = "tree-sitter-python")]
184        "python" => highlight_with_language(
185            lines,
186            tree_sitter_python::LANGUAGE.into(),
187            "python",
188            tree_sitter_python::HIGHLIGHTS_QUERY,
189        ),
190
191        #[cfg(feature = "tree-sitter-rust")]
192        "rust" => highlight_with_language(
193            lines,
194            tree_sitter_rust::LANGUAGE.into(),
195            "rust",
196            tree_sitter_rust::HIGHLIGHTS_QUERY,
197        ),
198
199        #[cfg(feature = "tree-sitter-scala")]
200        "scala" => highlight_with_language(
201            lines,
202            tree_sitter_scala::LANGUAGE.into(),
203            "scala",
204            tree_sitter_scala::HIGHLIGHTS_QUERY,
205        ),
206
207        #[cfg(feature = "tree-sitter-typescript")]
208        "tsx" => highlight_with_language(
209            lines,
210            tree_sitter_typescript::LANGUAGE_TSX.into(),
211            "tsx",
212            tree_sitter_typescript::HIGHLIGHTS_QUERY,
213        ),
214
215        #[cfg(feature = "tree-sitter-typescript")]
216        "typescript" | "ts" => highlight_with_language(
217            lines,
218            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
219            "typescript",
220            tree_sitter_typescript::HIGHLIGHTS_QUERY,
221        ),
222
223        #[cfg(feature = "tree-sitter-yaml")]
224        "yaml" | "yml" => highlight_with_language(
225            lines,
226            tree_sitter_yaml::LANGUAGE.into(),
227            "yaml",
228            tree_sitter_yaml::HIGHLIGHTS_QUERY,
229        ),
230
231        "mermaid" => return HighlightInfo::Mermaid,
232
233        _ => return HighlightInfo::Unhighlighted,
234    };
235
236    match result {
237        Ok(events) => HighlightInfo::Highlighted(events),
238        Err(_) => HighlightInfo::Unhighlighted,
239    }
240}
241
242thread_local! {
243    /// Per-language `HighlightConfiguration` cache: lock-free, build-on-demand
244    /// memoization that needs no synchronization at any thread count and matches
245    /// tree-sitter's per-thread `Highlighter` model. (The type is `Sync`, so a
246    /// shared `OnceLock`-per-language cache is also possible; not worth it since
247    /// highlighting is main-thread-only.)
248    static HIGHLIGHT_CONFIGS: RefCell<HashMap<&'static str, HighlightConfiguration>> =
249        RefCell::new(HashMap::new());
250}
251
252pub fn highlight_with_language(
253    lines: &[u8],
254    language: tree_sitter::Language,
255    lang_name: &'static str,
256    query: &str,
257) -> Result<Vec<HighlightEvent>, String> {
258    HIGHLIGHT_CONFIGS.with(|cell| {
259        let mut configs = cell.borrow_mut();
260        if !configs.contains_key(lang_name) {
261            let mut config = HighlightConfiguration::new(language, lang_name, query, "", "")
262                .map_err(|e| e.to_string())?;
263            config.configure(&HIGHLIGHT_NAMES);
264            configs.insert(lang_name, config);
265        }
266        let config = configs.get(lang_name).expect("inserted above if missing");
267
268        let mut highlighter = Highlighter::new();
269        let events = highlighter
270            .highlight(config, lines, None, |_| None)
271            .map_err(|e| e.to_string())?;
272        events
273            .collect::<Result<Vec<_>, _>>()
274            .map_err(|e| e.to_string())
275    })
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    #[test]
283    fn test_equal_length() {
284        assert_eq!(HIGHLIGHT_NAMES.len(), COLOR_MAP.len());
285    }
286
287    #[test]
288    #[cfg(feature = "tree-sitter-typescript")]
289    fn test_highlight_typescript() {
290        let code = b"const x: number = 1;";
291        let result = highlight_code("typescript", code);
292        if let HighlightInfo::Highlighted(events) = result {
293            assert!(!events.is_empty());
294        } else {
295            panic!("Expected Highlighted, got {:?}", result);
296        }
297    }
298
299    #[test]
300    #[cfg(feature = "tree-sitter-typescript")]
301    fn test_highlight_tsx() {
302        let code = b"const x = <div>hello</div>;";
303        let result = highlight_code("tsx", code);
304        if let HighlightInfo::Highlighted(events) = result {
305            assert!(!events.is_empty());
306        } else {
307            panic!("Expected Highlighted, got {:?}", result);
308        }
309    }
310}