rustik_highlight/

blob.rs

//! Incremental highlighting over complete text blobs.
//!
//! This module sits above raw line tokenization and keeps the per-line parser
//! state needed by editor and pager workloads. Callers can ask for one line or
//! a visible range, and unchanged lines reuse cached start states and scope
//! spans instead of reparsing from the top of the file.

use std::ops::Range;

use crate::grammar::{Grammar, LineState, LineTokenizer, ScopeSpan};
use crate::theme::{StyleCache, StyleScratch, StyleSpan, Theme};
use crate::util::line_starts;

/// Scratch space for highlighting one line without reallocating per call.
///
/// A buffer holds the latest tokenization and styling results so callers can
/// reuse allocations while scanning many lines or repeatedly re-highlighting a
/// changing line.
#[derive(Debug, Default)]
pub struct LineBuffer {
    /// Current scope spans.
    pub scopes: Vec<ScopeSpan>,
    /// Current style spans.
    pub styles: Vec<StyleSpan>,
    /// Cached resolved styles for the active grammar/theme pair.
    style_cache: StyleCache,
    /// Per-line scratch storage used when converting scope spans into style spans.
    style_scratch: StyleScratch,
}

/// Borrowed tokenization result for a line visited by a range highlighter.
///
/// This is the zero-copy callback view used when callers need grammar scopes
/// but not themed styles. The borrowed text and spans remain valid only for
/// the duration of the callback.
#[derive(Debug)]
pub struct LineTokens<'text, 'spans> {
    /// Zero-based line index in the blob.
    pub line_index: usize,
    /// Byte range of this line in the blob, excluding line terminators.
    pub byte_range: Range<usize>,
    /// Line text, excluding line terminators.
    pub text: &'text str,
    /// Scope spans for this line.
    pub scopes: &'spans [ScopeSpan],
}

/// Borrowed tokenization and styling result for a line visited by a range highlighter.
///
/// This is the zero-copy callback view used when callers need both grammar
/// scopes and theme-resolved style spans. The borrowed text and spans remain
/// valid only for the duration of the callback.
#[derive(Debug)]
pub struct StyledLine<'text, 'spans> {
    /// Zero-based line index in the blob.
    pub line_index: usize,
    /// Byte range of this line in the blob, excluding line terminators.
    pub byte_range: Range<usize>,
    /// Line text, excluding line terminators.
    pub text: &'text str,
    /// Scope spans for this line.
    pub scopes: &'spans [ScopeSpan],
    /// Styled spans for this line.
    pub styles: &'spans [StyleSpan],
}

/// Owned tokenization result for APIs that collect highlighted lines.
///
/// Unlike [`LineTokens`], this type owns its scope spans so callers can keep
/// the result after range traversal has moved on to another line.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct OwnedLineTokens {
    /// Zero-based line index in the blob.
    pub line_index: usize,
    /// Byte range of this line in the blob, excluding line terminators.
    pub byte_range: Range<usize>,
    /// Scope spans for this line.
    pub scopes: Vec<ScopeSpan>,
}

/// Incremental highlighter for a complete text blob.
///
/// The highlighter caches parser states and per-line scope spans so callers can
/// ask for individual lines or ranges without reparsing the whole blob each
/// time. Editing callers should call [`BlobHighlighter::reset_text`] or
/// [`BlobHighlighter::invalidate_from`] to keep those caches aligned with the
/// backing text.
#[derive(Debug)]
pub struct BlobHighlighter<'text> {
    /// Reusable tokenizer for the backing grammar.
    tokenizer: LineTokenizer<'text>,
    /// Backing source text.
    text: &'text str,
    /// Byte offsets for each source line start.
    line_starts: Vec<usize>,
    /// Parser state at the start of each source line.
    states: Vec<Option<LineState>>,
    /// Cached scope spans for each source line.
    scopes: Vec<Option<Vec<ScopeSpan>>>,
    /// Reusable per-line scratch for styled range calls.
    buffer: LineBuffer,
}

impl LineBuffer {
    /// Clears all buffered spans while retaining allocations.
    ///
    /// Cached scope styles are intentionally retained: they are keyed by `(theme,
    /// grammar, scope_count)` and stays valid across line resets.
    pub fn clear(&mut self) {
        self.scopes.clear();
        self.styles.clear();
        self.style_scratch.clear_line();
    }

    /// Tokenizes one line into this reusable buffer.
    pub fn tokenize<'a>(
        &'a mut self,
        grammar: &Grammar,
        state: &mut LineState,
        line: &str,
    ) -> &'a [ScopeSpan] {
        grammar.tokenize_line_into(state, line, &mut self.scopes);
        &self.scopes
    }

    /// Applies a theme to this buffer's current scope spans.
    pub fn style<'a>(
        &'a mut self,
        grammar: &Grammar,
        theme: &Theme,
        line: &str,
    ) -> &'a [StyleSpan] {
        self.style_cache.refresh(theme, grammar);
        theme.style_line_into(
            grammar,
            line,
            &self.scopes,
            &self.style_cache,
            &mut self.style_scratch,
            &mut self.styles,
        );
        &self.styles
    }

    /// Tokenizes and styles one line into this reusable buffer.
    pub fn highlight<'a>(
        &'a mut self,
        grammar: &Grammar,
        theme: &Theme,
        state: &mut LineState,
        line: &str,
    ) -> (&'a [ScopeSpan], &'a [StyleSpan]) {
        grammar.tokenize_line_into(state, line, &mut self.scopes);
        self.style(grammar, theme, line);
        (&self.scopes, &self.styles)
    }
}

impl<'text> BlobHighlighter<'text> {
    /// Creates a blob highlighter and caches only the initial empty state.
    pub fn new(grammar: &'text Grammar, text: &'text str) -> Self {
        let mut highlighter = Self {
            tokenizer: LineTokenizer::new(grammar),
            text,
            line_starts: Vec::new(),
            states: Vec::new(),
            scopes: Vec::new(),
            buffer: LineBuffer::default(),
        };
        highlighter.rebuild_caches();
        highlighter
    }

    /// Replaces the blob text and clears cached line states.
    pub fn reset_text(&mut self, text: &'text str) {
        self.text = text;
        self.rebuild_caches();
    }

    /// Rebuilds line metadata and clears cached highlighting results.
    fn rebuild_caches(&mut self) {
        self.line_starts = line_starts(self.text);
        self.states = vec![None; self.line_starts.len() + 1];
        self.states[0] = Some(LineState::default());
        self.scopes = vec![None; self.line_starts.len()];
    }

    /// Returns the number of highlightable lines in the blob.
    pub fn line_count(&self) -> usize {
        self.line_starts.len()
    }

    /// Returns a line by index, excluding line terminators.
    pub fn line(&self, line: usize) -> Option<&'text str> {
        let range = self.line_byte_range(line)?;
        Some(&self.text[range])
    }

    /// Returns a line's byte range, excluding line terminators.
    pub fn line_byte_range(&self, line: usize) -> Option<Range<usize>> {
        (line < self.line_count()).then(|| self.line_range_unchecked(line))
    }

    /// Returns whether a start state for `line` is already cached.
    pub fn is_state_cached(&self, line: usize) -> bool {
        self.states.get(line).is_some_and(Option::is_some)
    }

    /// Invalidates cached states after a changed line.
    pub fn invalidate_from(&mut self, line: usize) {
        let start = line.saturating_add(1).min(self.states.len());
        for state in &mut self.states[start..] {
            *state = None;
        }
        let start = line.min(self.scopes.len());
        for scopes in &mut self.scopes[start..] {
            *scopes = None;
        }
    }

    /// Ensures the parser state at the start of `line` is cached.
    pub fn ensure_state(&mut self, line: usize) -> Option<&LineState> {
        if line > self.line_count() {
            return None;
        }
        if self.states[line].is_none() {
            let state = self.compute_state(line)?;
            self.states[line] = Some(state);
        }
        self.states[line].as_ref()
    }

    /// Tokenizes one line into a caller-owned buffer.
    pub fn highlight_line_into(&mut self, line: usize, scopes: &mut Vec<ScopeSpan>) -> bool {
        if line >= self.line_count() {
            scopes.clear();
            return false;
        }
        let Some(mut state) = self.ensure_state(line).cloned() else {
            scopes.clear();
            return false;
        };
        let range = self.line_range_unchecked(line);
        let text = &self.text[range];

        tokenize_or_cache_into(
            &mut self.tokenizer,
            &mut self.states,
            &mut self.scopes,
            &mut state,
            line,
            text,
            scopes,
        );
        true
    }

    /// Tokenizes a line range and invokes `f` for each requested line.
    pub fn highlight_range<F>(&mut self, range: Range<usize>, mut f: F)
    where
        F: FnMut(LineTokens<'_, '_>),
    {
        let start = range.start.min(self.line_count());
        let end = range.end.min(self.line_count());
        if start >= end {
            return;
        }
        let Some(mut state) = self.ensure_state(start).cloned() else {
            return;
        };

        let mut scopes = Vec::new();
        for line_index in start..end {
            let byte_range = self.line_range_unchecked(line_index);
            let text = &self.text[byte_range.clone()];

            tokenize_or_cache_into(
                &mut self.tokenizer,
                &mut self.states,
                &mut self.scopes,
                &mut state,
                line_index,
                text,
                &mut scopes,
            );
            f(LineTokens {
                line_index,
                byte_range,
                text,
                scopes: &scopes,
            });
        }
    }

    /// Tokenizes and styles a line range and invokes `f` for each requested line.
    pub fn highlight_styled_range<F>(&mut self, theme: &Theme, range: Range<usize>, mut f: F)
    where
        F: FnMut(StyledLine<'_, '_>),
    {
        let start = range.start.min(self.line_count());
        let end = range.end.min(self.line_count());
        if start >= end {
            return;
        }
        let Some(mut state) = self.ensure_state(start).cloned() else {
            return;
        };

        let grammar = self.tokenizer.grammar;
        for line_index in start..end {
            let byte_range = self.line_range_unchecked(line_index);
            let text = &self.text[byte_range.clone()];

            tokenize_or_cache_into(
                &mut self.tokenizer,
                &mut self.states,
                &mut self.scopes,
                &mut state,
                line_index,
                text,
                &mut self.buffer.scopes,
            );
            self.buffer.style(grammar, theme, text);

            f(StyledLine {
                line_index,
                byte_range,
                text,
                scopes: &self.buffer.scopes,
                styles: &self.buffer.styles,
            });
        }
    }

    /// Tokenizes a line range into an owned output vector.
    pub fn highlight_range_into(&mut self, range: Range<usize>, output: &mut Vec<OwnedLineTokens>) {
        output.clear();

        self.highlight_range(range, |line| {
            output.push(OwnedLineTokens {
                line_index: line.line_index,
                byte_range: line.byte_range,
                scopes: line.scopes.to_vec(),
            });
        });
    }

    /// Returns newly allocated token data for a line range.
    pub fn highlighted_range(&mut self, range: Range<usize>) -> Vec<OwnedLineTokens> {
        let mut output = Vec::new();
        self.highlight_range_into(range, &mut output);
        output
    }

    /// Computes the parser state at the start of `line` from the nearest cached predecessor.
    fn compute_state(&mut self, line: usize) -> Option<LineState> {
        let mut index = line;
        while index > 0 && self.states[index].is_none() {
            index -= 1;
        }
        let mut state = self.states[index].clone().unwrap_or_default();
        let mut scratch = Vec::new();

        while index < line {
            let range = self.line_range_unchecked(index);
            let text = &self.text[range];

            self.tokenizer
                .tokenize_line_into(&mut state, text, &mut scratch);

            index += 1;
            self.states[index] = Some(state.clone());
        }
        Some(state)
    }

    /// Returns a line's byte range without checking that the line index exists.
    fn line_range_unchecked(&self, line: usize) -> Range<usize> {
        let start = self.line_starts[line];
        let mut end = self
            .line_starts
            .get(line + 1)
            .copied()
            .unwrap_or(self.text.len());
        let bytes = self.text.as_bytes();
        if end > start && bytes[end - 1] == b'\n' {
            end -= 1;
        }
        if end > start && bytes[end - 1] == b'\r' {
            end -= 1;
        }
        start..end
    }
}

/// Tokenizes a line or copies its cached scopes and following parser state.
fn tokenize_or_cache_into(
    tokenizer: &mut LineTokenizer<'_>,
    states: &mut [Option<LineState>],
    scopes_cache: &mut [Option<Vec<ScopeSpan>>],
    state: &mut LineState,
    line_index: usize,
    text: &str,
    out: &mut Vec<ScopeSpan>,
) {
    if let Some(cached) = scopes_cache[line_index].as_ref()
        && let Some(next_state) = states[line_index + 1].as_ref()
    {
        out.clone_from(cached);
        *state = next_state.clone();
    } else {
        tokenizer.tokenize_line_into(state, text, out);

        states[line_index + 1] = Some(state.clone());
        scopes_cache[line_index] = Some(out.clone());
    }
}