declutter/
highlight.rs

1use std::ops::Range;
2use std::sync::OnceLock;
3
4use tree_sitter_highlight::{HighlightConfiguration, HighlightEvent, Highlighter};
5
6use crate::lang::{Lang, parsable};
7use crate::project::Projection;
8
9/// The syntax categories the viewer colours.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum Class {
12    Keyword,
13    String,
14    Comment,
15    Number,
16    Constant,
17    Type,
18    Function,
19    Attribute,
20}
21
22/// Capture names handed to tree-sitter-highlight; a query capture such as
23/// `function.method` resolves to the longest matching prefix here.
24const NAMES: [(&str, Class); 14] = [
25    ("keyword", Class::Keyword),
26    ("include", Class::Keyword),
27    ("conditional", Class::Keyword),
28    ("repeat", Class::Keyword),
29    ("string", Class::String),
30    ("comment", Class::Comment),
31    ("number", Class::Number),
32    ("float", Class::Number),
33    ("boolean", Class::Constant),
34    ("constant", Class::Constant),
35    ("type", Class::Type),
36    ("constructor", Class::Type),
37    ("function", Class::Function),
38    ("attribute", Class::Attribute),
39];
40
41/// Syntax classes per line, as byte ranges within that line.
42pub type LineClasses = Vec<(Range<usize>, Class)>;
43
44/// Colours a projection by highlighting its own text, so what is coloured is exactly
45/// what is shown — comments removed, or only comments — not the original file.
46pub fn annotate(projection: &mut Projection, lang: Lang) {
47    let Some(config) = config(lang) else {
48        return;
49    };
50    let source = projection.lines.join("\n");
51    let mut starts = Vec::with_capacity(projection.lines.len());
52    let mut offset = 0;
53    for line in &projection.lines {
54        starts.push(offset);
55        offset += line.len() + 1;
56    }
57
58    let mut classes: Vec<LineClasses> = vec![Vec::new(); projection.lines.len()];
59    let mut highlighter = Highlighter::new();
60    let parse_text = parsable(lang, &source);
61    let Ok(events) = highlighter.highlight(config, parse_text.as_bytes(), None, None, |_| None)
62    else {
63        return;
64    };
65    let mut active: Vec<Class> = Vec::new();
66    for event in events {
67        match event {
68            Ok(HighlightEvent::HighlightStart(highlight)) => active.push(NAMES[highlight.0].1),
69            Ok(HighlightEvent::HighlightEnd) => {
70                active.pop();
71            }
72            Ok(HighlightEvent::Source { start, end }) => {
73                let Some(&class) = active.last() else {
74                    continue;
75                };
76                // Split the span over the lines it crosses.
77                let mut line = starts.partition_point(|&s| s <= start).saturating_sub(1);
78                let mut from = start;
79                while from < end && line < starts.len() {
80                    let line_end = starts[line] + projection.lines[line].len();
81                    let to = end.min(line_end);
82                    if to > from {
83                        let span = from - starts[line]..to - starts[line];
84                        match classes[line].last_mut() {
85                            // Grammars split some tokens (a string and its quotes); join them back.
86                            Some((last, last_class))
87                                if last.end == span.start && *last_class == class =>
88                            {
89                                last.end = span.end
90                            }
91                            _ => classes[line].push((span, class)),
92                        }
93                    }
94                    line += 1;
95                    from = starts.get(line).copied().unwrap_or(end);
96                }
97            }
98            Err(_) => return,
99        }
100    }
101    projection.syntax = classes;
102}
103
104pub fn config(lang: Lang) -> Option<&'static HighlightConfiguration> {
105    static SWIFT: OnceLock<Option<HighlightConfiguration>> = OnceLock::new();
106    static TYPESCRIPT: OnceLock<Option<HighlightConfiguration>> = OnceLock::new();
107    static TSX: OnceLock<Option<HighlightConfiguration>> = OnceLock::new();
108    static PYTHON: OnceLock<Option<HighlightConfiguration>> = OnceLock::new();
109    static GO: OnceLock<Option<HighlightConfiguration>> = OnceLock::new();
110    static RUST: OnceLock<Option<HighlightConfiguration>> = OnceLock::new();
111    static KOTLIN: OnceLock<Option<HighlightConfiguration>> = OnceLock::new();
112
113    let build = |query: String| {
114        let mut config =
115            HighlightConfiguration::new(lang.grammar(), lang.name(), &query, "", "").ok()?;
116        let names: Vec<&str> = NAMES.iter().map(|(name, _)| *name).collect();
117        config.configure(&names);
118        Some(config)
119    };
120    // The TypeScript queries only add to the JavaScript ones, so both are needed.
121    let ecma = || {
122        format!(
123            "{}\n{}",
124            tree_sitter_typescript::HIGHLIGHTS_QUERY,
125            tree_sitter_javascript::HIGHLIGHT_QUERY
126        )
127    };
128    let cell = match lang {
129        Lang::Swift => &SWIFT,
130        Lang::TypeScript => &TYPESCRIPT,
131        Lang::Tsx | Lang::JavaScript => &TSX,
132        Lang::Python => &PYTHON,
133        Lang::Go => &GO,
134        Lang::Rust => &RUST,
135        Lang::Kotlin => &KOTLIN,
136    };
137    cell.get_or_init(|| match lang {
138        Lang::Swift => build(tree_sitter_swift::HIGHLIGHTS_QUERY.to_string()),
139        Lang::TypeScript => build(ecma()),
140        Lang::Tsx | Lang::JavaScript => build(format!(
141            "{}\n{}",
142            ecma(),
143            tree_sitter_javascript::JSX_HIGHLIGHT_QUERY
144        )),
145        Lang::Python => build(tree_sitter_python::HIGHLIGHTS_QUERY.to_string()),
146        Lang::Go => build(tree_sitter_go::HIGHLIGHTS_QUERY.to_string()),
147        Lang::Rust => build(tree_sitter_rust::HIGHLIGHTS_QUERY.to_string()),
148        Lang::Kotlin => build(include_str!("queries/kotlin.scm").to_string()),
149    })
150    .as_ref()
151}
declutter/highlight.rs

declutter/
highlight.rs
