Skip to main content

zeph_tui/
highlight.rs

1// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::collections::HashMap;
5use std::sync::LazyLock;
6
7use ratatui::style::{Color, Modifier, Style};
8use ratatui::text::Span;
9use tree_sitter::Language;
10use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter};
11
12use crate::theme::SyntaxTheme;
13
14const CAPTURE_NAMES: &[&str] = &[
15    "attribute",
16    "comment",
17    "constant",
18    "constant.builtin",
19    "constructor",
20    "function",
21    "function.builtin",
22    "keyword",
23    "number",
24    "operator",
25    "property",
26    "punctuation",
27    "punctuation.bracket",
28    "punctuation.delimiter",
29    "string",
30    "string.escape",
31    "type",
32    "type.builtin",
33    "variable",
34    "variable.builtin",
35    "variable.parameter",
36];
37
38const BASH_HIGHLIGHTS_QUERY: &str = r#"
39[(string) (raw_string) (heredoc_body) (heredoc_start)] @string
40(command_name) @function
41(variable_name) @property
42["case" "do" "done" "elif" "else" "esac" "export" "fi" "for" "function" "if" "in" "select" "then" "unset" "until" "while"] @keyword
43(comment) @comment
44(function_definition name: (word) @function)
45(file_descriptor) @number
46["$" "&&" ">" ">>" "<" "|"] @operator
47((command (_) @constant) (#match? @constant "^-"))
48"#;
49
50static LANG_ALIASES: LazyLock<HashMap<&'static str, &'static str>> = LazyLock::new(|| {
51    HashMap::from([
52        ("rs", "rust"),
53        ("py", "python"),
54        ("js", "javascript"),
55        ("sh", "bash"),
56        ("shell", "bash"),
57    ])
58});
59
60/// Process-wide singleton [`SyntaxHighlighter`], initialised on first access.
61///
62/// Use this instead of constructing a new highlighter to avoid redundant
63/// tree-sitter grammar compilation on each frame.
64///
65/// # Examples
66///
67/// ```rust
68/// use zeph_tui::highlight::SYNTAX_HIGHLIGHTER;
69/// use zeph_tui::theme::SyntaxTheme;
70///
71/// let theme = SyntaxTheme::default();
72/// let spans = SYNTAX_HIGHLIGHTER.highlight("rust", "let x = 1;", &theme);
73/// assert!(spans.is_some());
74/// ```
75pub static SYNTAX_HIGHLIGHTER: LazyLock<SyntaxHighlighter> = LazyLock::new(SyntaxHighlighter::new);
76
77/// Tree-sitter-based syntax highlighter for TUI code blocks.
78///
79/// Supports Rust, Python, JavaScript, JSON, TOML, and Bash out of the box.
80/// Language aliases (`"rs"` → `"rust"`, `"sh"` → `"bash"`, etc.) are
81/// resolved transparently.
82///
83/// Construct via the [`SYNTAX_HIGHLIGHTER`] static for process-level sharing,
84/// or call the private `new` method directly in tests.
85///
86/// # Supported languages
87///
88/// | Identifier | Aliases |
89/// |------------|---------|
90/// | `rust`     | `rs`    |
91/// | `python`   | `py`    |
92/// | `javascript` | `js` |
93/// | `bash`     | `sh`, `shell` |
94/// | `json`     | —       |
95/// | `toml`     | —       |
96///
97/// # Examples
98///
99/// ```rust
100/// use zeph_tui::highlight::SYNTAX_HIGHLIGHTER;
101/// use zeph_tui::theme::SyntaxTheme;
102///
103/// let theme = SyntaxTheme::default();
104///
105/// // Known language → styled spans
106/// let spans = SYNTAX_HIGHLIGHTER.highlight("rust", "fn main() {}", &theme);
107/// assert!(spans.is_some());
108///
109/// // Alias works the same way
110/// let spans = SYNTAX_HIGHLIGHTER.highlight("rs", "let x = 1;", &theme);
111/// assert!(spans.is_some());
112///
113/// // Unknown language → None
114/// assert!(SYNTAX_HIGHLIGHTER.highlight("brainfuck", "+++", &theme).is_none());
115/// ```
116pub struct SyntaxHighlighter {
117    configs: HashMap<&'static str, HighlightConfiguration>,
118}
119
120impl SyntaxHighlighter {
121    fn new() -> Self {
122        let mut configs = HashMap::new();
123
124        let mut register = |name: &'static str,
125                            language: Language,
126                            lang_name: &str,
127                            highlights_query: &str,
128                            injections_query: &str| {
129            let Ok(mut config) = HighlightConfiguration::new(
130                language,
131                lang_name.to_string(),
132                highlights_query,
133                injections_query,
134                "",
135            ) else {
136                return;
137            };
138            config.configure(CAPTURE_NAMES);
139            configs.insert(name, config);
140        };
141
142        register(
143            "rust",
144            tree_sitter_rust::LANGUAGE.into(),
145            "rust",
146            tree_sitter_rust::HIGHLIGHTS_QUERY,
147            tree_sitter_rust::INJECTIONS_QUERY,
148        );
149
150        register(
151            "python",
152            tree_sitter_python::LANGUAGE.into(),
153            "python",
154            tree_sitter_python::HIGHLIGHTS_QUERY,
155            "",
156        );
157
158        register(
159            "javascript",
160            tree_sitter_javascript::LANGUAGE.into(),
161            "javascript",
162            tree_sitter_javascript::HIGHLIGHT_QUERY,
163            tree_sitter_javascript::INJECTIONS_QUERY,
164        );
165
166        register(
167            "json",
168            tree_sitter_json::LANGUAGE.into(),
169            "json",
170            tree_sitter_json::HIGHLIGHTS_QUERY,
171            "",
172        );
173
174        register(
175            "toml",
176            tree_sitter_toml_ng::LANGUAGE.into(),
177            "toml",
178            tree_sitter_toml_ng::HIGHLIGHTS_QUERY,
179            "",
180        );
181
182        register(
183            "bash",
184            tree_sitter_bash::LANGUAGE.into(),
185            "bash",
186            BASH_HIGHLIGHTS_QUERY,
187            "",
188        );
189
190        Self { configs }
191    }
192
193    /// Highlight `code` for the given `lang` using `theme`.
194    ///
195    /// Returns `None` if the language is unsupported or if tree-sitter fails
196    /// to parse the input. The returned spans concatenate to the original
197    /// source text unchanged.
198    ///
199    /// # Arguments
200    ///
201    /// * `lang` — language identifier or alias (case-insensitive).
202    /// * `code` — source code to highlight.
203    /// * `theme` — style mapping for each token class.
204    ///
205    /// # Examples
206    ///
207    /// ```rust
208    /// use zeph_tui::highlight::SYNTAX_HIGHLIGHTER;
209    /// use zeph_tui::theme::SyntaxTheme;
210    ///
211    /// let theme = SyntaxTheme::default();
212    /// let spans = SYNTAX_HIGHLIGHTER.highlight("rust", "let x = 42;", &theme).unwrap();
213    /// let text: String = spans.iter().map(|s| s.content.as_ref()).collect();
214    /// assert_eq!(text, "let x = 42;");
215    /// ```
216    pub fn highlight(
217        &self,
218        lang: &str,
219        code: &str,
220        theme: &SyntaxTheme,
221    ) -> Option<Vec<Span<'static>>> {
222        let lang_lower = lang.to_lowercase();
223        let canonical = LANG_ALIASES
224            .get(lang_lower.as_str())
225            .copied()
226            .unwrap_or(lang_lower.as_str());
227        let config = self.configs.get(canonical)?;
228
229        let mut highlighter = Highlighter::new();
230        let events = highlighter
231            .highlight(config, code.as_bytes(), None, |_| None)
232            .ok()?;
233
234        let mut spans = Vec::new();
235        let mut style_stack: Vec<Style> = Vec::new();
236
237        for event in events {
238            match event.ok()? {
239                HighlightEvent::Source { start, end } => {
240                    let text = code.get(start..end).unwrap_or_default();
241                    let style = style_stack.last().copied().unwrap_or(theme.default);
242                    spans.push(Span::styled(text.to_string(), style));
243                }
244                HighlightEvent::HighlightStart(highlight) => {
245                    let style = capture_to_style(highlight.0, theme);
246                    style_stack.push(style);
247                }
248                HighlightEvent::HighlightEnd => {
249                    style_stack.pop();
250                }
251            }
252        }
253
254        Some(spans)
255    }
256}
257
258fn capture_to_style(index: usize, theme: &SyntaxTheme) -> Style {
259    match CAPTURE_NAMES.get(index).copied().unwrap_or_default() {
260        "attribute" => theme.attribute,
261        "comment" => theme.comment,
262        "constant" | "constant.builtin" => theme.constant,
263        "constructor" | "type" | "type.builtin" => theme.r#type,
264        "function" | "function.builtin" => theme.function,
265        "keyword" => theme.keyword,
266        "number" => theme.number,
267        "operator" => theme.operator,
268        "property" | "variable" | "variable.builtin" | "variable.parameter" => theme.variable,
269        "punctuation" | "punctuation.bracket" | "punctuation.delimiter" => theme.punctuation,
270        "string" | "string.escape" => theme.string,
271        _ => theme.default,
272    }
273}
274
275impl Default for SyntaxTheme {
276    fn default() -> Self {
277        Self {
278            keyword: Style::default()
279                .fg(Color::Rgb(198, 120, 221))
280                .add_modifier(Modifier::BOLD),
281            string: Style::default().fg(Color::Rgb(152, 195, 121)),
282            comment: Style::default()
283                .fg(Color::Rgb(92, 99, 112))
284                .add_modifier(Modifier::ITALIC),
285            function: Style::default().fg(Color::Rgb(97, 175, 239)),
286            r#type: Style::default().fg(Color::Rgb(229, 192, 123)),
287            number: Style::default().fg(Color::Rgb(209, 154, 102)),
288            operator: Style::default().fg(Color::Rgb(171, 178, 191)),
289            variable: Style::default().fg(Color::Rgb(224, 108, 117)),
290            attribute: Style::default().fg(Color::Rgb(229, 192, 123)),
291            punctuation: Style::default().fg(Color::Rgb(171, 178, 191)),
292            constant: Style::default().fg(Color::Rgb(209, 154, 102)),
293            default: Style::default().fg(Color::Rgb(190, 175, 145)),
294        }
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    #[test]
303    fn highlight_rust_code() {
304        let hl = &*SYNTAX_HIGHLIGHTER;
305        let theme = SyntaxTheme::default();
306        let spans = hl.highlight("rust", "let x = 42;", &theme);
307        assert!(spans.is_some());
308        let spans = spans.unwrap();
309        assert!(!spans.is_empty());
310        let text: String = spans.iter().map(|s| s.content.as_ref()).collect();
311        assert_eq!(text, "let x = 42;");
312    }
313
314    #[test]
315    fn highlight_python_code() {
316        let hl = &*SYNTAX_HIGHLIGHTER;
317        let theme = SyntaxTheme::default();
318        let spans = hl.highlight("python", "def foo():\n    pass", &theme);
319        assert!(spans.is_some());
320    }
321
322    #[test]
323    fn highlight_unknown_lang_returns_none() {
324        let hl = &*SYNTAX_HIGHLIGHTER;
325        let theme = SyntaxTheme::default();
326        assert!(hl.highlight("brainfuck", "+++", &theme).is_none());
327    }
328
329    #[test]
330    fn highlight_json_code() {
331        let hl = &*SYNTAX_HIGHLIGHTER;
332        let theme = SyntaxTheme::default();
333        let spans = hl.highlight("json", r#"{"key": "value"}"#, &theme);
334        assert!(spans.is_some());
335    }
336
337    #[test]
338    fn highlight_js_code() {
339        let hl = &*SYNTAX_HIGHLIGHTER;
340        let theme = SyntaxTheme::default();
341        let spans = hl.highlight("js", "const x = 1;", &theme);
342        assert!(spans.is_some());
343    }
344
345    #[test]
346    fn highlight_alias_rs() {
347        let hl = &*SYNTAX_HIGHLIGHTER;
348        let theme = SyntaxTheme::default();
349        assert!(hl.highlight("rs", "fn main() {}", &theme).is_some());
350    }
351
352    #[test]
353    fn highlight_empty_string() {
354        let hl = &*SYNTAX_HIGHLIGHTER;
355        let theme = SyntaxTheme::default();
356        let spans = hl.highlight("rust", "", &theme);
357        assert!(spans.is_some());
358        assert!(spans.unwrap().is_empty());
359    }
360
361    #[test]
362    fn highlight_malformed_code_no_panic() {
363        let hl = &*SYNTAX_HIGHLIGHTER;
364        let theme = SyntaxTheme::default();
365        // Malformed Rust — should not panic, tree-sitter is error-tolerant
366        let spans = hl.highlight("rust", "fn {{{{ let !!!", &theme);
367        assert!(spans.is_some());
368    }
369
370    #[test]
371    fn highlight_toml_code() {
372        let hl = &*SYNTAX_HIGHLIGHTER;
373        let theme = SyntaxTheme::default();
374        let spans = hl.highlight("toml", "[package]\nname = \"foo\"", &theme);
375        assert!(spans.is_some());
376    }
377
378    #[test]
379    fn highlight_bash_code() {
380        let hl = &*SYNTAX_HIGHLIGHTER;
381        let theme = SyntaxTheme::default();
382        let spans = hl.highlight("bash", "echo \"hello\"", &theme);
383        assert!(spans.is_some());
384    }
385
386    #[test]
387    fn rust_keywords_get_keyword_style() {
388        let hl = &*SYNTAX_HIGHLIGHTER;
389        let theme = SyntaxTheme::default();
390        let spans = hl.highlight("rust", "let x = 1;", &theme).unwrap();
391        let let_span = spans.iter().find(|s| s.content.as_ref() == "let").unwrap();
392        assert_eq!(let_span.style, theme.keyword);
393    }
394}