radicle_cli/terminal/
highlight.rs

1use std::{collections::HashMap, path::Path};
2
3use radicle_term as term;
4use tree_sitter_highlight as ts;
5
6/// Highlight groups enabled.
7const HIGHLIGHTS: &[&str] = &[
8    "attribute",
9    "constant",
10    "constant.builtin",
11    "comment",
12    "constructor",
13    "function.builtin",
14    "function",
15    "integer_literal",
16    "float.literal",
17    "keyword",
18    "label",
19    "number",
20    "operator",
21    "property",
22    "punctuation",
23    "punctuation.bracket",
24    "punctuation.delimiter",
25    "punctuation.special",
26    "string",
27    "string.special",
28    "tag",
29    "type",
30    "type.builtin",
31    "variable",
32    "variable.builtin",
33    "variable.parameter",
34    "text.literal",
35    "text.title",
36];
37
38/// Syntax highlighter based on `tree-sitter`.
39#[derive(Default)]
40pub struct Highlighter {
41    configs: HashMap<&'static str, ts::HighlightConfiguration>,
42}
43
44/// Syntax theme.
45pub struct Theme {
46    color: fn(&'static str) -> Option<term::Color>,
47}
48
49impl Default for Theme {
50    fn default() -> Self {
51        let color = if term::Paint::truecolor() {
52            term::colors::rgb::theme
53        } else {
54            term::colors::fixed::theme
55        };
56        Self { color }
57    }
58}
59
60impl Theme {
61    /// Get the named color.
62    pub fn color(&self, color: &'static str) -> term::Color {
63        if let Some(c) = (self.color)(color) {
64            c
65        } else {
66            term::Color::Unset
67        }
68    }
69
70    /// Return the color of a syntax group.
71    pub fn highlight(&self, group: &'static str) -> Option<term::Color> {
72        let color = match group {
73            "keyword" => self.color("red"),
74            "comment" => self.color("grey"),
75            "constant" => self.color("orange"),
76            "number" => self.color("blue"),
77            "string" => self.color("teal"),
78            "string.special" => self.color("green"),
79            "function" => self.color("purple"),
80            "operator" => self.color("blue"),
81            // Eg. `true` and `false` in rust.
82            "constant.builtin" => self.color("blue"),
83            "type.builtin" => self.color("teal"),
84            "punctuation.bracket" | "punctuation.delimiter" => term::Color::default(),
85            // Eg. the '#' in Markdown titles.
86            "punctuation.special" => self.color("dim"),
87            // Eg. Markdown code blocks.
88            "text.literal" => self.color("blue"),
89            "text.title" => self.color("orange"),
90            "variable.builtin" => term::Color::default(),
91            "property" => self.color("blue"),
92            // Eg. `#[derive(Debug)]` in rust
93            "attribute" => self.color("blue"),
94            "label" => self.color("green"),
95            // `Option`
96            "type" => self.color("grey.light"),
97            "variable.parameter" => term::Color::default(),
98            "constructor" => self.color("orange"),
99
100            _ => return None,
101        };
102        Some(color)
103    }
104}
105
106/// Syntax highlighted file builder.
107#[derive(Default)]
108struct Builder {
109    /// Output lines.
110    lines: Vec<term::Line>,
111    /// Current output line.
112    line: Vec<term::Label>,
113    /// Current label.
114    label: Vec<u8>,
115    /// Current stack of styles.
116    styles: Vec<term::Style>,
117}
118
119impl Builder {
120    /// Run the builder to completion.
121    fn run(
122        mut self,
123        highlights: impl Iterator<Item = Result<ts::HighlightEvent, ts::Error>>,
124        code: &[u8],
125        theme: &Theme,
126    ) -> Result<Vec<term::Line>, ts::Error> {
127        for event in highlights {
128            match event? {
129                ts::HighlightEvent::Source { start, end } => {
130                    for (i, byte) in code.iter().enumerate().skip(start).take(end - start) {
131                        if *byte == b'\n' {
132                            self.advance();
133                            // Start on new line.
134                            self.lines.push(term::Line::from(self.line.clone()));
135                            self.line.clear();
136                        } else if i == code.len() - 1 {
137                            // File has no `\n` at the end.
138                            self.label.push(*byte);
139                            self.advance();
140                            self.lines.push(term::Line::from(self.line.clone()));
141                        } else {
142                            // Add to existing label.
143                            self.label.push(*byte);
144                        }
145                    }
146                }
147                ts::HighlightEvent::HighlightStart(h) => {
148                    let name = HIGHLIGHTS[h.0];
149                    let style =
150                        term::Style::default().fg(theme.highlight(name).unwrap_or_default());
151
152                    self.advance();
153                    self.styles.push(style);
154                }
155                ts::HighlightEvent::HighlightEnd => {
156                    self.advance();
157                    self.styles.pop();
158                }
159            }
160        }
161        Ok(self.lines)
162    }
163
164    /// Advance the state by pushing the current label onto the current line,
165    /// using the current styling.
166    fn advance(&mut self) {
167        if !self.label.is_empty() {
168            // Take the top-level style when there are more than one.
169            let style = self.styles.first().cloned().unwrap_or_default();
170            self.line
171                .push(term::Label::new(String::from_utf8_lossy(&self.label).as_ref()).style(style));
172            self.label.clear();
173        }
174    }
175}
176
177impl Highlighter {
178    /// Highlight a source code file.
179    pub fn highlight(&mut self, path: &Path, code: &[u8]) -> Result<Vec<term::Line>, ts::Error> {
180        let theme = Theme::default();
181        let mut highlighter = ts::Highlighter::new();
182        let Some(config) = self.detect(path, code) else {
183            let Ok(code) = std::str::from_utf8(code) else {
184                return Err(ts::Error::Unknown);
185            };
186            return Ok(code.lines().map(term::Line::new).collect());
187        };
188        config.configure(HIGHLIGHTS);
189
190        let highlights = highlighter.highlight(config, code, None, |_| {
191            // Language injection callback.
192            None
193        })?;
194
195        Builder::default().run(highlights, code, &theme)
196    }
197
198    /// Detect language.
199    fn detect(&mut self, path: &Path, _code: &[u8]) -> Option<&mut ts::HighlightConfiguration> {
200        match path.extension().and_then(|e| e.to_str()) {
201            Some("rs") => self.config("rust"),
202            Some("ts" | "js") => self.config("typescript"),
203            Some("json") => self.config("json"),
204            Some("sh" | "bash") => self.config("shell"),
205            Some("md" | "markdown") => self.config("markdown"),
206            Some("go") => self.config("go"),
207            Some("c") => self.config("c"),
208            Some("py") => self.config("python"),
209            Some("rb") => self.config("ruby"),
210            Some("tsx") => self.config("tsx"),
211            Some("html") | Some("htm") | Some("xml") => self.config("html"),
212            Some("css") => self.config("css"),
213            Some("toml") => self.config("toml"),
214            _ => None,
215        }
216    }
217
218    /// Get a language configuration.
219    fn config(&mut self, language: &'static str) -> Option<&mut ts::HighlightConfiguration> {
220        match language {
221            "rust" => Some(self.configs.entry(language).or_insert_with(|| {
222                ts::HighlightConfiguration::new(
223                    tree_sitter_rust::LANGUAGE.into(),
224                    language,
225                    tree_sitter_rust::HIGHLIGHTS_QUERY,
226                    tree_sitter_rust::INJECTIONS_QUERY,
227                    "",
228                )
229                .expect("Highlighter::config: highlight configuration must be valid")
230            })),
231            "json" => Some(self.configs.entry(language).or_insert_with(|| {
232                ts::HighlightConfiguration::new(
233                    tree_sitter_json::LANGUAGE.into(),
234                    language,
235                    tree_sitter_json::HIGHLIGHTS_QUERY,
236                    "",
237                    "",
238                )
239                .expect("Highlighter::config: highlight configuration must be valid")
240            })),
241            "typescript" => Some(self.configs.entry(language).or_insert_with(|| {
242                ts::HighlightConfiguration::new(
243                    tree_sitter_typescript::LANGUAGE_TYPESCRIPT.into(),
244                    language,
245                    tree_sitter_typescript::HIGHLIGHTS_QUERY,
246                    "",
247                    tree_sitter_typescript::LOCALS_QUERY,
248                )
249                .expect("Highlighter::config: highlight configuration must be valid")
250            })),
251            "markdown" => Some(self.configs.entry(language).or_insert_with(|| {
252                ts::HighlightConfiguration::new(
253                    tree_sitter_md::LANGUAGE.into(),
254                    language,
255                    tree_sitter_md::HIGHLIGHT_QUERY_BLOCK,
256                    tree_sitter_md::INJECTION_QUERY_BLOCK,
257                    "",
258                )
259                .expect("Highlighter::config: highlight configuration must be valid")
260            })),
261            "css" => Some(self.configs.entry(language).or_insert_with(|| {
262                ts::HighlightConfiguration::new(
263                    tree_sitter_css::LANGUAGE.into(),
264                    language,
265                    tree_sitter_css::HIGHLIGHTS_QUERY,
266                    "",
267                    "",
268                )
269                .expect("Highlighter::config: highlight configuration must be valid")
270            })),
271            "go" => Some(self.configs.entry(language).or_insert_with(|| {
272                ts::HighlightConfiguration::new(
273                    tree_sitter_go::LANGUAGE.into(),
274                    language,
275                    tree_sitter_go::HIGHLIGHTS_QUERY,
276                    "",
277                    "",
278                )
279                .expect("Highlighter::config: highlight configuration must be valid")
280            })),
281            "shell" => Some(self.configs.entry(language).or_insert_with(|| {
282                ts::HighlightConfiguration::new(
283                    tree_sitter_bash::LANGUAGE.into(),
284                    language,
285                    tree_sitter_bash::HIGHLIGHT_QUERY,
286                    "",
287                    "",
288                )
289                .expect("Highlighter::config: highlight configuration must be valid")
290            })),
291            "c" => Some(self.configs.entry(language).or_insert_with(|| {
292                ts::HighlightConfiguration::new(
293                    tree_sitter_c::LANGUAGE.into(),
294                    language,
295                    tree_sitter_c::HIGHLIGHT_QUERY,
296                    "",
297                    "",
298                )
299                .expect("Highlighter::config: highlight configuration must be valid")
300            })),
301            "python" => Some(self.configs.entry(language).or_insert_with(|| {
302                ts::HighlightConfiguration::new(
303                    tree_sitter_python::LANGUAGE.into(),
304                    language,
305                    tree_sitter_python::HIGHLIGHTS_QUERY,
306                    "",
307                    "",
308                )
309                .expect("Highlighter::config: highlight configuration must be valid")
310            })),
311            "ruby" => Some(self.configs.entry(language).or_insert_with(|| {
312                ts::HighlightConfiguration::new(
313                    tree_sitter_ruby::LANGUAGE.into(),
314                    language,
315                    tree_sitter_ruby::HIGHLIGHTS_QUERY,
316                    "",
317                    tree_sitter_ruby::LOCALS_QUERY,
318                )
319                .expect("Highlighter::config: highlight configuration must be valid")
320            })),
321            "tsx" => Some(self.configs.entry(language).or_insert_with(|| {
322                ts::HighlightConfiguration::new(
323                    tree_sitter_typescript::LANGUAGE_TSX.into(),
324                    language,
325                    tree_sitter_typescript::HIGHLIGHTS_QUERY,
326                    "",
327                    tree_sitter_typescript::LOCALS_QUERY,
328                )
329                .expect("Highlighter::config: highlight configuration must be valid")
330            })),
331            "html" => Some(self.configs.entry(language).or_insert_with(|| {
332                ts::HighlightConfiguration::new(
333                    tree_sitter_html::LANGUAGE.into(),
334                    language,
335                    tree_sitter_html::HIGHLIGHTS_QUERY,
336                    tree_sitter_html::INJECTIONS_QUERY,
337                    "",
338                )
339                .expect("Highlighter::config: highlight configuration must be valid")
340            })),
341            "toml" => Some(self.configs.entry(language).or_insert_with(|| {
342                ts::HighlightConfiguration::new(
343                    tree_sitter_toml_ng::language(),
344                    language,
345                    tree_sitter_toml_ng::HIGHLIGHTS_QUERY,
346                    "",
347                    "",
348                )
349                .expect("Highlighter::config: highlight configuration must be valid")
350            })),
351            _ => None,
352        }
353    }
354}