1use std::cell::RefCell;
2use std::collections::HashMap;
3
4use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter};
5
6use ratatui::style::Color;
7
8#[allow(dead_code)]
9static HIGHLIGHT_NAMES: [&str; 18] = [
10 "attribute",
11 "constant",
12 "function.builtin",
13 "function",
14 "keyword",
15 "operator",
16 "property",
17 "punctuation",
18 "punctuation.bracket",
19 "punctuation.delimiter",
20 "string",
21 "string.special",
22 "tag",
23 "type",
24 "type.builtin",
25 "variable",
26 "variable.builtin",
27 "variable.parameter",
28];
29
30pub static COLOR_MAP: [Color; 18] = [
31 Color::Yellow,
32 Color::Yellow,
33 Color::Green,
34 Color::Green,
35 Color::Red,
36 Color::Red,
37 Color::Blue,
38 Color::Blue,
39 Color::Blue,
40 Color::Blue,
41 Color::Magenta,
42 Color::Magenta,
43 Color::Cyan,
44 Color::Cyan,
45 Color::Cyan,
46 Color::Reset,
47 Color::Reset,
48 Color::Reset,
49];
50
51#[derive(Debug)]
52pub enum HighlightInfo {
53 Highlighted(Vec<HighlightEvent>),
54 Mermaid,
55 Unhighlighted,
56}
57
58#[allow(unused_variables, unreachable_code)]
63#[must_use]
64pub fn highlight_code(language: &str, lines: &[u8]) -> HighlightInfo {
65 let result: Result<Vec<HighlightEvent>, String> = match language {
71 #[cfg(feature = "tree-sitter-bash")]
72 "bash" | "sh" => highlight_with_language(
73 lines,
74 tree_sitter_bash::LANGUAGE.into(),
75 "bash",
76 tree_sitter_bash::HIGHLIGHT_QUERY,
77 ),
78
79 #[cfg(feature = "tree-sitter-c")]
80 "c" => highlight_with_language(
81 lines,
82 tree_sitter_c::LANGUAGE.into(),
83 "c",
84 tree_sitter_c::HIGHLIGHT_QUERY,
85 ),
86
87 #[cfg(feature = "tree-sitter-cpp")]
88 "cpp" => highlight_with_language(
89 lines,
90 tree_sitter_cpp::LANGUAGE.into(),
91 "cpp",
92 tree_sitter_cpp::HIGHLIGHT_QUERY,
93 ),
94
95 #[cfg(feature = "tree-sitter-css")]
96 "css" => highlight_with_language(
97 lines,
98 tree_sitter_css::LANGUAGE.into(),
99 "css",
100 tree_sitter_css::HIGHLIGHTS_QUERY,
101 ),
102
103 #[cfg(feature = "tree-sitter-diff")]
104 "diff" | "patch" => highlight_with_language(
105 lines,
106 tree_sitter_diff::LANGUAGE.into(),
107 "diff",
108 tree_sitter_diff::HIGHLIGHTS_QUERY,
109 ),
110
111 #[cfg(feature = "tree-sitter-elixir")]
112 "elixir" => highlight_with_language(
113 lines,
114 tree_sitter_elixir::LANGUAGE.into(),
115 "elixir",
116 tree_sitter_elixir::HIGHLIGHTS_QUERY,
117 ),
118
119 #[cfg(feature = "tree-sitter-go")]
120 "go" => highlight_with_language(
121 lines,
122 tree_sitter_go::LANGUAGE.into(),
123 "go",
124 tree_sitter_go::HIGHLIGHTS_QUERY,
125 ),
126
127 #[cfg(feature = "tree-sitter-html")]
128 "html" => highlight_with_language(
129 lines,
130 tree_sitter_html::LANGUAGE.into(),
131 "html",
132 tree_sitter_html::HIGHLIGHTS_QUERY,
133 ),
134
135 #[cfg(feature = "tree-sitter-java")]
136 "java" => highlight_with_language(
137 lines,
138 tree_sitter_java::LANGUAGE.into(),
139 "java",
140 tree_sitter_java::HIGHLIGHTS_QUERY,
141 ),
142
143 #[cfg(feature = "tree-sitter-javascript")]
144 "javascript" | "js" => highlight_with_language(
145 lines,
146 tree_sitter_javascript::LANGUAGE.into(),
147 "javascript",
148 tree_sitter_javascript::HIGHLIGHT_QUERY,
149 ),
150
151 #[cfg(feature = "tree-sitter-json")]
152 "json" => highlight_with_language(
153 lines,
154 tree_sitter_json::LANGUAGE.into(),
155 "json",
156 tree_sitter_json::HIGHLIGHTS_QUERY,
157 ),
158
159 #[cfg(feature = "tree-sitter-lua")]
160 "lua" => highlight_with_language(
161 lines,
162 tree_sitter_lua::LANGUAGE.into(),
163 "lua",
164 tree_sitter_lua::HIGHLIGHTS_QUERY,
165 ),
166
167 #[cfg(feature = "tree-sitter-ocaml")]
168 "ocaml" => highlight_with_language(
169 lines,
170 tree_sitter_ocaml::LANGUAGE_OCAML_TYPE.into(),
171 "ocaml",
172 tree_sitter_ocaml::HIGHLIGHTS_QUERY,
173 ),
174
175 #[cfg(feature = "tree-sitter-php")]
176 "php" => highlight_with_language(
177 lines,
178 tree_sitter_php::LANGUAGE_PHP.into(),
179 "php",
180 tree_sitter_php::HIGHLIGHTS_QUERY,
181 ),
182
183 #[cfg(feature = "tree-sitter-python")]
184 "python" => highlight_with_language(
185 lines,
186 tree_sitter_python::LANGUAGE.into(),
187 "python",
188 tree_sitter_python::HIGHLIGHTS_QUERY,
189 ),
190
191 #[cfg(feature = "tree-sitter-rust")]
192 "rust" => highlight_with_language(
193 lines,
194 tree_sitter_rust::LANGUAGE.into(),
195 "rust",
196 tree_sitter_rust::HIGHLIGHTS_QUERY,
197 ),
198
199 #[cfg(feature = "tree-sitter-scala")]
200 "scala" => highlight_with_language(
201 lines,
202 tree_sitter_scala::LANGUAGE.into(),
203 "scala",
204 tree_sitter_scala::HIGHLIGHTS_QUERY,
205 ),
206
207 #[cfg(feature = "tree-sitter-typescript")]
208 "tsx" => highlight_with_language(
209 lines,
210 tree_sitter_typescript::LANGUAGE_TSX.into(),
211 "tsx",
212 tree_sitter_typescript::HIGHLIGHTS_QUERY,
213 ),
214
215 #[cfg(feature = "tree-sitter-typescript")]
216 "typescript" | "ts" => highlight_with_language(
217 lines,
218 tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
219 "typescript",
220 tree_sitter_typescript::HIGHLIGHTS_QUERY,
221 ),
222
223 #[cfg(feature = "tree-sitter-yaml")]
224 "yaml" | "yml" => highlight_with_language(
225 lines,
226 tree_sitter_yaml::LANGUAGE.into(),
227 "yaml",
228 tree_sitter_yaml::HIGHLIGHTS_QUERY,
229 ),
230
231 "mermaid" => return HighlightInfo::Mermaid,
232
233 _ => return HighlightInfo::Unhighlighted,
234 };
235
236 match result {
237 Ok(events) => HighlightInfo::Highlighted(events),
238 Err(_) => HighlightInfo::Unhighlighted,
239 }
240}
241
242thread_local! {
243 static HIGHLIGHT_CONFIGS: RefCell<HashMap<&'static str, HighlightConfiguration>> =
249 RefCell::new(HashMap::new());
250}
251
252pub fn highlight_with_language(
253 lines: &[u8],
254 language: tree_sitter::Language,
255 lang_name: &'static str,
256 query: &str,
257) -> Result<Vec<HighlightEvent>, String> {
258 HIGHLIGHT_CONFIGS.with(|cell| {
259 let mut configs = cell.borrow_mut();
260 if !configs.contains_key(lang_name) {
261 let mut config = HighlightConfiguration::new(language, lang_name, query, "", "")
262 .map_err(|e| e.to_string())?;
263 config.configure(&HIGHLIGHT_NAMES);
264 configs.insert(lang_name, config);
265 }
266 let config = configs.get(lang_name).expect("inserted above if missing");
267
268 let mut highlighter = Highlighter::new();
269 let events = highlighter
270 .highlight(config, lines, None, |_| None)
271 .map_err(|e| e.to_string())?;
272 events
273 .collect::<Result<Vec<_>, _>>()
274 .map_err(|e| e.to_string())
275 })
276}
277
278#[cfg(test)]
279mod tests {
280 use super::*;
281
282 #[test]
283 fn test_equal_length() {
284 assert_eq!(HIGHLIGHT_NAMES.len(), COLOR_MAP.len());
285 }
286
287 #[test]
288 #[cfg(feature = "tree-sitter-typescript")]
289 fn test_highlight_typescript() {
290 let code = b"const x: number = 1;";
291 let result = highlight_code("typescript", code);
292 if let HighlightInfo::Highlighted(events) = result {
293 assert!(!events.is_empty());
294 } else {
295 panic!("Expected Highlighted, got {:?}", result);
296 }
297 }
298
299 #[test]
300 #[cfg(feature = "tree-sitter-typescript")]
301 fn test_highlight_tsx() {
302 let code = b"const x = <div>hello</div>;";
303 let result = highlight_code("tsx", code);
304 if let HighlightInfo::Highlighted(events) = result {
305 assert!(!events.is_empty());
306 } else {
307 panic!("Expected Highlighted, got {:?}", result);
308 }
309 }
310}