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