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::{CodeTheme, Theme};
33
34const 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
67fn 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
91fn 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 "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 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#[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
224fn 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 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 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}