1use 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
35const 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
68fn 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
92fn 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 "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 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#[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
225fn 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 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 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}