clankerdiff_syntax/
highlight.rs

1//! Tree-sitter syntax highlighting with UTF-8 byte spans and a bounded cache.
2
3use crate::{
4    SyntaxError,
5    incremental::{AppendContext, Grammars, IncrementalDocument, SyntaxWorkStats},
6    language::{LanguageHint, resolve_language},
7    spans::{Span, spans_to_flat_tokens},
8};
9use arborium_theme::tag_to_name;
10use clankerdiff_theme::{Fingerprint, HighlightSpan, SyntaxTheme};
11use imbl::Vector;
12use lru::LruCache;
13use std::{
14    fmt, mem,
15    num::NonZeroUsize,
16    ops::Range,
17    sync::{Arc, OnceLock},
18};
19
20const DEFAULT_MAX_DOCUMENTS: usize = 512;
21const DEFAULT_SOURCE_BYTES: usize = 8 * 1024 * 1024;
22const MAX_INJECTION_DEPTH: usize = 3;
23const DOCUMENT_KEY_DOMAIN: &[u8] = b"syntax-document-v2";
24
25#[must_use]
26pub fn empty_spans() -> Arc<[HighlightSpan]> {
27    static EMPTY: OnceLock<Arc<[HighlightSpan]>> = OnceLock::new();
28    Arc::clone(EMPTY.get_or_init(|| Arc::from(Vec::<HighlightSpan>::new())))
29}
30
31#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
32pub struct HighlightStats {
33    pub calls: u64,
34    pub hits: u64,
35    pub misses: u64,
36    pub evictions: u64,
37    pub bytes: usize,
38}
39
40#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
41pub struct CacheKey {
42    fingerprint: Fingerprint,
43}
44
45impl CacheKey {
46    #[must_use]
47    pub const fn fingerprint(self) -> Fingerprint {
48        self.fingerprint
49    }
50
51    fn document(language: &str, source_id: Fingerprint) -> Self {
52        Self {
53            fingerprint: Fingerprint::of([
54                DOCUMENT_KEY_DOMAIN,
55                language.as_bytes(),
56                source_id.as_bytes().as_slice(),
57            ]),
58        }
59    }
60}
61
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub struct CacheConfig {
64    pub max_documents: usize,
65    pub max_source_bytes: usize,
66}
67
68impl Default for CacheConfig {
69    fn default() -> Self {
70        Self {
71            max_documents: DEFAULT_MAX_DOCUMENTS,
72            max_source_bytes: DEFAULT_SOURCE_BYTES,
73        }
74    }
75}
76
77impl From<usize> for CacheConfig {
78    fn from(max_documents: usize) -> Self {
79        Self {
80            max_documents,
81            ..Self::default()
82        }
83    }
84}
85
86#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
87pub struct CacheUsage {
88    pub document_entries: usize,
89}
90
91#[derive(Debug, Clone, Default)]
92pub struct SyntaxStream {
93    hint: String,
94    source: String,
95    revision: u64,
96    theme_revision: Option<Fingerprint>,
97    highlights: Arc<DocumentHighlights>,
98    document: Option<IncrementalDocument>,
99    language: Option<&'static str>,
100    line_starts: Vec<usize>,
101    work: SyntaxWorkStats,
102}
103
104#[derive(Debug, Clone)]
105pub struct SyntaxStreamUpdate {
106    pub base_revision: u64,
107    pub revision: u64,
108    pub changed_lines: Range<usize>,
109    pub highlights: Arc<DocumentHighlights>,
110}
111
112impl SyntaxStream {
113    #[must_use]
114    pub fn new<'a>(hint: impl Into<LanguageHint<'a>>) -> Self {
115        Self {
116            hint: hint.into().as_str().to_owned(),
117            ..Self::default()
118        }
119    }
120
121    #[must_use]
122    pub fn source(&self) -> &str {
123        &self.source
124    }
125
126    #[must_use]
127    pub const fn revision(&self) -> u64 {
128        self.revision
129    }
130
131    #[must_use]
132    pub fn highlights(&self) -> &Arc<DocumentHighlights> {
133        &self.highlights
134    }
135
136    #[must_use]
137    pub const fn work_stats(&self) -> SyntaxWorkStats {
138        self.work
139    }
140
141    fn parse(
142        &mut self,
143        grammars: &mut Grammars,
144        previous_len: usize,
145    ) -> Result<usize, SyntaxError> {
146        let resolved = resolve_language(self.hint.as_str(), &self.source);
147        let language_changed = resolved != self.language;
148        if language_changed {
149            self.document = None;
150            self.language = resolved;
151        }
152        let Some(language) = resolved else {
153            return Ok(if language_changed { 0 } else { previous_len });
154        };
155        let document = match &mut self.document {
156            Some(document) => document,
157            slot => slot.insert(grammars.document(language)?.ok_or_else(|| {
158                SyntaxError::MissingGrammar {
159                    language: language.to_owned(),
160                }
161            })?),
162        };
163        document.append(
164            &self.source,
165            0,
166            MAX_INJECTION_DEPTH,
167            &mut AppendContext {
168                grammars,
169                stats: &mut self.work,
170                line_starts: &self.line_starts,
171            },
172        )
173    }
174
175    fn update(
176        &mut self,
177        grammars: &mut Grammars,
178        stats: &mut HighlightStats,
179        theme: &SyntaxTheme,
180        appended: &str,
181    ) -> Result<SyntaxStreamUpdate, SyntaxError> {
182        let base_revision = self.revision;
183        let theme_revision = theme.revision();
184        if appended.is_empty() && self.theme_revision == Some(theme_revision) {
185            let end = self.highlights.line_count();
186            return Ok(SyntaxStreamUpdate {
187                base_revision,
188                revision: self.revision,
189                changed_lines: end..end,
190                highlights: Arc::clone(&self.highlights),
191            });
192        }
193        let previous_len = self.source.len();
194        let previous_lines = self.line_starts.len();
195        self.source.push_str(appended);
196        if self.line_starts.is_empty() && !self.source.is_empty() {
197            self.line_starts.push(0);
198        }
199        self.line_starts.extend(
200            appended
201                .bytes()
202                .enumerate()
203                .filter_map(|(index, byte)| (byte == b'\n').then_some(previous_len + index + 1)),
204        );
205        let previous_work = self.work.parser_input_bytes;
206        let result = self.parse(grammars, previous_len);
207        stats.bytes += self.work.parser_input_bytes - previous_work;
208        let start = match result {
209            Ok(start) => start,
210            Err(error) => {
211                self.source.truncate(previous_len);
212                self.line_starts.truncate(previous_lines);
213                self.document = None;
214                return Err(error);
215            }
216        };
217        let start = if self.theme_revision == Some(theme_revision) {
218            start
219        } else {
220            0
221        };
222        let first_changed = self.project(theme, start);
223        if !appended.is_empty() {
224            self.revision = self.revision.wrapping_add(1);
225        }
226        self.theme_revision = Some(theme_revision);
227        Ok(SyntaxStreamUpdate {
228            base_revision,
229            revision: self.revision,
230            changed_lines: first_changed..self.highlights.line_count(),
231            highlights: Arc::clone(&self.highlights),
232        })
233    }
234
235    fn project(&mut self, theme: &SyntaxTheme, start: usize) -> usize {
236        let first_line = self
237            .line_starts
238            .partition_point(|&offset| offset <= start)
239            .saturating_sub(1);
240        let byte_start = self.line_starts.get(first_line).copied().unwrap_or(0);
241        let raw = self
242            .document
243            .as_ref()
244            .map_or_else(Vec::new, |document| document.spans_from(byte_start));
245        let spans = map_spans(theme, &self.source, raw);
246        let mut lines = self
247            .highlights
248            .lines
249            .take(first_line.min(self.highlights.line_count()));
250        let mut first_changed = lines.len();
251        let mut first_span = 0;
252        self.work.projected_bytes += self.source.len() - byte_start;
253        self.work.reused_lines += lines.len();
254        for (line, &from) in self.line_starts.iter().enumerate().skip(first_line) {
255            if from == self.source.len() {
256                break;
257            }
258            let mut to = self
259                .line_starts
260                .get(line + 1)
261                .copied()
262                .unwrap_or(self.source.len());
263            if to > from && self.source.as_bytes()[to - 1] == b'\n' {
264                to -= 1;
265            }
266            if to > from && self.source.as_bytes()[to - 1] == b'\r' {
267                to -= 1;
268            }
269            while spans
270                .get(first_span)
271                .is_some_and(|span| span.range.end <= from)
272            {
273                first_span += 1;
274            }
275            let projected: Vec<_> = spans[first_span..]
276                .iter()
277                .take_while(|span| span.range.start < to)
278                .filter_map(|span| {
279                    let start = span.range.start.max(from);
280                    let end = span.range.end.min(to);
281                    (start < end).then_some(HighlightSpan {
282                        range: start - from..end - from,
283                        foreground: span.foreground,
284                        font_style: span.font_style,
285                    })
286                })
287                .collect();
288            self.work.projected_lines += 1;
289            if let Some(before) = self
290                .highlights
291                .lines
292                .get(line)
293                .filter(|before| before.as_ref() == projected)
294            {
295                lines.push_back(Arc::clone(before));
296                self.work.reused_lines += 1;
297                if first_changed == line {
298                    first_changed += 1;
299                }
300            } else {
301                lines.push_back(if projected.is_empty() {
302                    empty_spans()
303                } else {
304                    projected.into()
305                });
306            }
307        }
308        self.highlights = Arc::new(DocumentHighlights { lines });
309        first_changed
310    }
311}
312
313#[derive(Debug, Clone, Default)]
314pub struct DocumentHighlights {
315    lines: Vector<Arc<[HighlightSpan]>>,
316}
317
318impl DocumentHighlights {
319    #[must_use]
320    pub fn line(&self, index: usize) -> Option<&[HighlightSpan]> {
321        self.lines.get(index).map(AsRef::as_ref)
322    }
323
324    #[must_use]
325    pub fn line_shared(&self, index: usize) -> Option<Arc<[HighlightSpan]>> {
326        self.lines.get(index).cloned()
327    }
328
329    #[must_use]
330    pub fn line_count(&self) -> usize {
331        self.lines.len()
332    }
333}
334
335pub struct SyntaxHighlighter {
336    config: CacheConfig,
337    documents: Option<LruCache<CacheKey, SyntaxStream>>,
338    stats: HighlightStats,
339    grammars: Grammars,
340}
341
342impl fmt::Debug for SyntaxHighlighter {
343    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
344        f.debug_struct("SyntaxHighlighter")
345            .field("config", &self.config)
346            .field(
347                "documents",
348                &self.documents.as_ref().map_or(0, LruCache::len),
349            )
350            .field("stats", &self.stats)
351            .finish_non_exhaustive()
352    }
353}
354
355impl Default for SyntaxHighlighter {
356    fn default() -> Self {
357        Self::new(CacheConfig::default())
358    }
359}
360
361impl SyntaxHighlighter {
362    #[must_use]
363    pub fn new(config: impl Into<CacheConfig>) -> Self {
364        let config = config.into();
365        Self {
366            config,
367            documents: NonZeroUsize::new(config.max_documents).map(LruCache::new),
368            stats: HighlightStats::default(),
369            grammars: Grammars::default(),
370        }
371    }
372
373    #[must_use]
374    pub const fn stats(&self) -> HighlightStats {
375        self.stats
376    }
377
378    pub fn reset_stats(&mut self) {
379        self.stats = HighlightStats::default();
380    }
381
382    pub fn take_stats(&mut self) -> HighlightStats {
383        mem::take(&mut self.stats)
384    }
385
386    #[must_use]
387    pub const fn config(&self) -> CacheConfig {
388        self.config
389    }
390
391    #[must_use]
392    pub fn with_theme<'a>(&'a mut self, theme: &'a SyntaxTheme) -> ThemedHighlighter<'a> {
393        ThemedHighlighter {
394            highlighter: self,
395            theme,
396        }
397    }
398
399    #[must_use]
400    pub fn cache_usage(&self) -> CacheUsage {
401        CacheUsage {
402            document_entries: self.documents.as_ref().map_or(0, LruCache::len),
403        }
404    }
405
406    pub fn clear_cache(&mut self) {
407        if let Some(documents) = &mut self.documents {
408            documents.clear();
409        }
410    }
411}
412
413pub struct ThemedHighlighter<'a> {
414    highlighter: &'a mut SyntaxHighlighter,
415    theme: &'a SyntaxTheme,
416}
417
418impl ThemedHighlighter<'_> {
419    pub fn highlight_document<'a, T: AsRef<str>>(
420        &mut self,
421        source_id: Fingerprint,
422        language: impl Into<LanguageHint<'a>>,
423        source: impl FnOnce() -> T,
424    ) -> Result<Arc<DocumentHighlights>, SyntaxError> {
425        let language = language.into();
426        let key = CacheKey::document(resolve_language(language, "").unwrap_or("auto"), source_id);
427        let SyntaxHighlighter {
428            config,
429            documents,
430            stats,
431            grammars,
432        } = &mut *self.highlighter;
433        if let Some(stream) = documents
434            .as_mut()
435            .and_then(|documents| documents.get_mut(&key))
436        {
437            stats.calls += 1;
438            stats.hits += 1;
439            return stream
440                .update(grammars, stats, self.theme, "")
441                .map(|update| update.highlights);
442        }
443        let text = source();
444        check_limit(config, text.as_ref().len())?;
445        stats.calls += 1;
446        stats.misses += 1;
447        let mut stream = SyntaxStream::new(language);
448        let highlights = stream
449            .update(grammars, stats, self.theme, text.as_ref())?
450            .highlights;
451        if let Some(documents) = documents
452            && documents.push(key, stream).is_some()
453        {
454            stats.evictions += 1;
455        }
456        Ok(highlights)
457    }
458
459    pub fn append(
460        &mut self,
461        stream: &mut SyntaxStream,
462        appended: &str,
463    ) -> Result<SyntaxStreamUpdate, SyntaxError> {
464        let SyntaxHighlighter {
465            config,
466            stats,
467            grammars,
468            ..
469        } = &mut *self.highlighter;
470        check_limit(config, stream.source.len().saturating_add(appended.len()))?;
471        stats.calls += 1;
472        stream.update(grammars, stats, self.theme, appended)
473    }
474}
475
476fn check_limit(config: &CacheConfig, attempted: usize) -> Result<(), SyntaxError> {
477    let limit = config.max_source_bytes.min(u32::MAX as usize);
478    if attempted > limit {
479        return Err(SyntaxError::InputLimit { limit, attempted });
480    }
481    Ok(())
482}
483
484fn map_spans(theme: &SyntaxTheme, source: &str, raw_spans: Vec<Span>) -> Vec<HighlightSpan> {
485    if raw_spans.is_empty() {
486        return Vec::new();
487    }
488    let tokens = spans_to_flat_tokens(source, raw_spans);
489    let mut spans: Vec<HighlightSpan> = Vec::with_capacity(tokens.len());
490    for token in tokens {
491        let start = token.start as usize;
492        let end = token.end as usize;
493        if start >= end
494            || end > source.len()
495            || !source.is_char_boundary(start)
496            || !source.is_char_boundary(end)
497        {
498            continue;
499        }
500        let Some(capture) = diff_capture_name(token.tag) else {
501            continue;
502        };
503        let Some(style) = theme.style(capture) else {
504            continue;
505        };
506        let span = HighlightSpan {
507            range: start..end,
508            foreground: style.foreground,
509            font_style: style.font_style,
510        };
511        if let Some(last) = spans.last_mut()
512            && last.range.end == span.range.start
513            && last.foreground == span.foreground
514            && last.font_style == span.font_style
515        {
516            last.range.end = span.range.end;
517        } else {
518            spans.push(span);
519        }
520    }
521    spans
522}
523
524fn diff_capture_name(tag: &str) -> Option<&'static str> {
525    Some(match tag_to_name(tag)? {
526        "title" => "markup.heading",
527        "strong" => "markup.bold",
528        "emphasis" => "markup.italic",
529        "link" => "markup.link",
530        "literal" => "markup.raw",
531        "strikethrough" => "markup.strikethrough",
532        name => name,
533    })
534}
clankerdiff_syntax/highlight.rs

clankerdiff_syntax/
highlight.rs
