Skip to main content

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