Skip to main content

tuika_codeformatters/
lib.rs

1//! Tree-sitter syntax highlighting for [`tuika`]'s
2//! [`CodeBlock`](tuika::components::CodeBlock) and [`Markdown`](tuika::components::Markdown) components.
3//!
4//! tuika owns the *presentation* of code (framing, background, language label,
5//! wrapping) but deliberately depends on no grammar. This crate fills the gap: a
6//! ready-made [`Highlighter`](tuika::highlight::Highlighter) backed by the same tree-sitter
7//! grammars a coding tool already carries, mapping token classes onto the host's
8//! [`Theme`]'s [`code`](tuika::style::CodeTheme) palette so highlighted code follows the
9//! theme.
10//!
11//! ```
12//! use tuika::prelude::*;
13//! use tuika_codeformatters::TreeSitterHighlighter;
14//!
15//! let hl = TreeSitterHighlighter::new();
16//! let _block = CodeBlock::new("rust", "fn main() {}").highlighter(&hl);
17//! let _ = Theme::default();
18//! ```
19//!
20//! Supported languages (with common aliases): Rust, Python, TypeScript/
21//! JavaScript, TSX/JSX, Go, Java, Ruby, CSS, HTML, C#, PHP, Zig, Scala, and SQL.
22//! Anything else — or source that fails to parse — returns [`None`], and the
23//! caller renders it as plain code.
24
25use std::cell::RefCell;
26use std::collections::HashMap;
27use std::rc::Rc;
28
29use ratatui::style::{Modifier, Style};
30use ratatui::text::Span;
31use tree_sitter_highlight::{Highlight, HighlightConfiguration, HighlightEvent, Highlighter};
32use tuika::Theme;
33use tuika::style::CodeTheme;
34
35/// Highlight capture names we recognize, longest-specific first so
36/// tree-sitter-highlight resolves the most precise style. Kept in sync with
37/// [`style_for_name`].
38const HIGHLIGHT_NAMES: &[&str] = &[
39    "keyword",
40    "function.builtin",
41    "function.method",
42    "function",
43    "constructor",
44    "type.builtin",
45    "type",
46    "constant.builtin",
47    "constant.numeric",
48    "constant",
49    "number",
50    "string.special",
51    "string",
52    "escape",
53    "comment",
54    "operator",
55    "punctuation.bracket",
56    "punctuation.delimiter",
57    "punctuation.special",
58    "punctuation",
59    "property",
60    "attribute",
61    "tag",
62    "label",
63    "variable.builtin",
64    "variable.parameter",
65    "variable",
66];
67
68/// Map a tree-sitter capture name to a themed [`Style`] using the host palette.
69fn style_for_name(name: &str, code: &CodeTheme) -> Style {
70    let base = Style::default();
71    match name {
72        "keyword" => base.fg(code.keyword).add_modifier(Modifier::BOLD),
73        "function" | "function.builtin" | "function.method" | "constructor" | "label" => {
74            base.fg(code.function)
75        }
76        "type" | "type.builtin" => base.fg(code.type_name),
77        "constant" | "constant.builtin" | "constant.numeric" | "number" | "variable.builtin" => {
78            base.fg(code.constant)
79        }
80        "string" | "string.special" | "escape" => base.fg(code.string),
81        "comment" => base.fg(code.comment).add_modifier(Modifier::ITALIC),
82        "operator"
83        | "punctuation"
84        | "punctuation.bracket"
85        | "punctuation.delimiter"
86        | "punctuation.special" => base.fg(code.punctuation),
87        "attribute" | "tag" => base.fg(code.keyword),
88        _ => base.fg(code.text),
89    }
90}
91
92/// Map a fence info string (e.g. `rust`, `py`, `ts`, `c#`) to the canonical
93/// grammar key, or `None` when we don't highlight that language.
94fn canonical_language(lang: &str) -> Option<&'static str> {
95    let lang = lang.trim().to_ascii_lowercase();
96    let key = match lang.as_str() {
97        "rust" | "rs" => "rust",
98        "python" | "py" => "python",
99        // The TypeScript grammar is a superset that parses plain JS cleanly.
100        "typescript" | "ts" | "javascript" | "js" | "mjs" | "cjs" => "typescript",
101        "tsx" | "jsx" => "tsx",
102        "go" | "golang" => "go",
103        "java" => "java",
104        "ruby" | "rb" => "ruby",
105        "css" => "css",
106        "html" | "htm" => "html",
107        "c#" | "cs" | "csharp" | "c_sharp" => "csharp",
108        "php" => "php",
109        "zig" => "zig",
110        "scala" => "scala",
111        "sql" => "sql",
112        _ => return None,
113    };
114    Some(key)
115}
116
117fn build_configuration(key: &str) -> Option<HighlightConfiguration> {
118    let (language, query): (tree_sitter::Language, &str) = match key {
119        "rust" => (
120            tree_sitter_rust::LANGUAGE.into(),
121            tree_sitter_rust::HIGHLIGHTS_QUERY,
122        ),
123        "python" => (
124            tree_sitter_python::LANGUAGE.into(),
125            tree_sitter_python::HIGHLIGHTS_QUERY,
126        ),
127        "typescript" => (
128            tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
129            tree_sitter_typescript::HIGHLIGHTS_QUERY,
130        ),
131        "tsx" => (
132            tree_sitter_typescript::LANGUAGE_TSX.into(),
133            tree_sitter_typescript::HIGHLIGHTS_QUERY,
134        ),
135        "go" => (
136            tree_sitter_go::LANGUAGE.into(),
137            tree_sitter_go::HIGHLIGHTS_QUERY,
138        ),
139        "java" => (
140            tree_sitter_java::LANGUAGE.into(),
141            tree_sitter_java::HIGHLIGHTS_QUERY,
142        ),
143        "ruby" => (
144            tree_sitter_ruby::LANGUAGE.into(),
145            tree_sitter_ruby::HIGHLIGHTS_QUERY,
146        ),
147        "css" => (
148            tree_sitter_css::LANGUAGE.into(),
149            tree_sitter_css::HIGHLIGHTS_QUERY,
150        ),
151        "html" => (
152            tree_sitter_html::LANGUAGE.into(),
153            tree_sitter_html::HIGHLIGHTS_QUERY,
154        ),
155        "csharp" => (
156            tree_sitter_c_sharp::LANGUAGE.into(),
157            tree_sitter_c_sharp::HIGHLIGHTS_QUERY,
158        ),
159        "php" => (
160            tree_sitter_php::LANGUAGE_PHP.into(),
161            tree_sitter_php::HIGHLIGHTS_QUERY,
162        ),
163        "zig" => (
164            tree_sitter_zig::LANGUAGE.into(),
165            tree_sitter_zig::HIGHLIGHTS_QUERY,
166        ),
167        "scala" => (
168            tree_sitter_scala::LANGUAGE.into(),
169            tree_sitter_scala::HIGHLIGHTS_QUERY,
170        ),
171        "sql" => (
172            tree_sitter_sequel::LANGUAGE.into(),
173            tree_sitter_sequel::HIGHLIGHTS_QUERY,
174        ),
175        _ => return None,
176    };
177    let mut config = HighlightConfiguration::new(language, key, query, "", "").ok()?;
178    let names: Vec<String> = HIGHLIGHT_NAMES.iter().map(|n| n.to_string()).collect();
179    config.configure(&names);
180    Some(config)
181}
182
183thread_local! {
184    /// Per-language highlight configs, built on first use. `None` marks a
185    /// language whose config failed to build so we don't retry it every render.
186    static CONFIGS: RefCell<HashMap<&'static str, Option<Rc<HighlightConfiguration>>>> =
187        RefCell::new(HashMap::new());
188}
189
190fn config_for(key: &'static str) -> Option<Rc<HighlightConfiguration>> {
191    CONFIGS.with(|configs| {
192        configs
193            .borrow_mut()
194            .entry(key)
195            .or_insert_with(|| build_configuration(key).map(Rc::new))
196            .clone()
197    })
198}
199
200/// A tree-sitter-backed [`Highlighter`](tuika::highlight::Highlighter).
201///
202/// Zero-sized and cheap to construct; the per-language parser configurations are
203/// built lazily and cached thread-locally on first use, so keeping one around is
204/// no better than making one per frame.
205#[derive(Clone, Copy, Debug, Default)]
206pub struct TreeSitterHighlighter;
207
208impl TreeSitterHighlighter {
209    pub fn new() -> Self {
210        Self
211    }
212}
213
214impl tuika::highlight::Highlighter for TreeSitterHighlighter {
215    fn highlight(
216        &self,
217        lang: &str,
218        lines: &[&str],
219        theme: &Theme,
220    ) -> Option<Vec<Vec<Span<'static>>>> {
221        highlight_lines(lang, lines, &theme.code)
222    }
223}
224
225/// Highlight a fenced code block, returning one span vector per input line, or
226/// `None` for unsupported languages / parse failures (see the crate docs).
227fn highlight_lines(
228    lang: &str,
229    lines: &[&str],
230    code: &CodeTheme,
231) -> Option<Vec<Vec<Span<'static>>>> {
232    if lines.is_empty() {
233        return None;
234    }
235    let key = canonical_language(lang)?;
236    let config = config_for(key)?;
237    let source = lines.join("\n");
238
239    let mut highlighter = Highlighter::new();
240    let events = highlighter
241        .highlight(&config, source.as_bytes(), None, |_| None)
242        .ok()?;
243
244    let default_style = Style::default().fg(code.text);
245    let mut style_stack: Vec<Style> = Vec::new();
246    let mut out: Vec<Vec<Span<'static>>> = vec![Vec::new()];
247    for event in events {
248        match event.ok()? {
249            HighlightEvent::HighlightStart(Highlight(index)) => {
250                let style = HIGHLIGHT_NAMES
251                    .get(index)
252                    .map(|name| style_for_name(name, code))
253                    .unwrap_or(default_style);
254                style_stack.push(style);
255            }
256            HighlightEvent::HighlightEnd => {
257                style_stack.pop();
258            }
259            HighlightEvent::Source { start, end } => {
260                let style = style_stack.last().copied().unwrap_or(default_style);
261                let text = source.get(start..end)?;
262                let mut segments = text.split('\n');
263                if let Some(first) = segments.next()
264                    && !first.is_empty()
265                {
266                    out.last_mut()?.push(Span::styled(first.to_string(), style));
267                }
268                for segment in segments {
269                    out.push(Vec::new());
270                    if !segment.is_empty() {
271                        out.last_mut()?
272                            .push(Span::styled(segment.to_string(), style));
273                    }
274                }
275            }
276        }
277    }
278
279    // The event stream must reproduce exactly one output line per source line;
280    // if it doesn't (unexpected), bail so the caller falls back deterministically.
281    if out.len() == lines.len() {
282        Some(out)
283    } else {
284        None
285    }
286}
287
288#[cfg(test)]
289mod tests {
290    use super::*;
291    use tuika::highlight::Highlighter as _;
292
293    fn plain(spans: &[Span<'static>]) -> String {
294        spans.iter().map(|s| s.content.as_ref()).collect()
295    }
296
297    #[test]
298    fn highlights_rust_keywords_distinctly() {
299        let theme = Theme::default();
300        let hl = TreeSitterHighlighter::new();
301        let lines = vec!["fn main() {", "    let x = 1;", "}"];
302        let out = hl
303            .highlight("rust", &lines, &theme)
304            .expect("rust highlights");
305        assert_eq!(out.len(), 3);
306        for (rendered, source) in out.iter().zip(lines.iter()) {
307            assert_eq!(&plain(rendered), source);
308        }
309        let fn_span = out[0]
310            .iter()
311            .find(|s| s.content.as_ref() == "fn")
312            .expect("fn span present");
313        assert_eq!(fn_span.style.fg, Some(theme.code.keyword));
314    }
315
316    #[test]
317    fn follows_the_host_theme() {
318        // A different palette restyles the same token.
319        let mut theme = Theme::default();
320        theme.code.keyword = ratatui::style::Color::Indexed(200);
321        let hl = TreeSitterHighlighter::new();
322        let out = hl
323            .highlight("rust", &["fn f() {}"], &theme)
324            .expect("rust highlights");
325        let fn_span = out[0]
326            .iter()
327            .find(|s| s.content.as_ref() == "fn")
328            .expect("fn span");
329        assert_eq!(fn_span.style.fg, Some(ratatui::style::Color::Indexed(200)));
330    }
331
332    #[test]
333    fn aliases_resolve_and_js_uses_typescript_grammar() {
334        let theme = Theme::default();
335        let hl = TreeSitterHighlighter::new();
336        assert_eq!(canonical_language("py"), Some("python"));
337        assert_eq!(canonical_language("js"), Some("typescript"));
338        assert_eq!(canonical_language("c#"), Some("csharp"));
339        assert!(hl.highlight("js", &["const x = 1;"], &theme).is_some());
340    }
341
342    #[test]
343    fn unsupported_language_returns_none() {
344        let theme = Theme::default();
345        let hl = TreeSitterHighlighter::new();
346        assert!(hl.highlight("brainfuck", &["+++"], &theme).is_none());
347        assert_eq!(canonical_language("whatever"), None);
348    }
349
350    #[test]
351    fn blank_lines_inside_a_block_are_preserved() {
352        let theme = Theme::default();
353        let hl = TreeSitterHighlighter::new();
354        let lines = vec!["x = 1", "", "y = 2"];
355        let out = hl
356            .highlight("python", &lines, &theme)
357            .expect("python highlights");
358        assert_eq!(out.len(), 3);
359        assert_eq!(plain(&out[1]), "");
360    }
361}