Skip to main content

tuika_codeformatters/
lib.rs

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