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