1use 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
60pub static SYNTAX_HIGHLIGHTER: LazyLock<SyntaxHighlighter> = LazyLock::new(SyntaxHighlighter::new);
76
77pub 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 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 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}