ai-app/client-core/src/highlight/mod.rs

//! `code` read once, left to right, into the spans that carry a colour.
//! Ported from `app/.../Highlighter.kt`.
//!
//! One pass with a small state -- in a comment, in a string, or in ordinary
//! code -- rather than a locator per token kind over the whole text, which
//! is what the library this replaced did and is why it found comments
//! before it knew the language: a `#` inside a shell string, a `//` inside
//! a URL and a block-comment opener inside a shell glob each commented out
//! the rest of a line that was nothing of the sort.
//!
//! Every span is produced by advancing an index forward, so the result is
//! ordered, non-overlapping and inside the code by construction. Nothing
//! here panics: an unterminated string or comment runs to the end of the
//! code, which is also what it looks like while a fence is still being
//! written.
//!
//! **Indices are char offsets, not byte offsets** -- the scanner works over
//! `Vec<char>`, mirroring the Kotlin original's `Char`-indexed strings, so
//! [`span_text`] is how a caller (and every test here) turns a [`Span`]
//! back into the text it covers.

pub mod languages;
pub mod markdown;

pub use languages::{
    Attributes, BlockComment, Language, Quote, Rules, fence_language, file_language, rules_for,
};

/// What a span of code is, in the terms a palette has a colour for.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Kind {
    Keyword,
    String,
    Literal,
    Comment,
    Metadata,
    Punctuation,
    Mark,
}

/// A run of [`Kind`] in the code, as a half-open range of **char** indices.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Span {
    pub start: usize,
    pub end: usize,
    pub kind: Kind,
}

/// The text a [`Span`] covers, for a caller working in char indices (every
/// test in this module, and any UI that also holds `code` as `Vec<char>`).
pub fn span_text(code: &[char], span: &Span) -> String {
    code[span.start..span.end].iter().collect()
}

/// The spans `language` colours in `code` -- the one way to ask, whatever
/// the language turns out to be made of. `None` draws plain.
pub fn spans_of(code: &str, language: Language) -> Vec<Span> {
    if language == Language::Markdown {
        markdown::scan_markdown(code)
    } else {
        scan(code, &rules_for(language))
    }
}

/// `code` read into the spans [`Rules`] describes. Also reachable directly
/// for a caller that already has a [`Rules`] (there is currently only one:
/// [`spans_of`]), kept public because the Kotlin original exposed it the
/// same way.
pub fn scan(code: &str, rules: &Rules) -> Vec<Span> {
    Scanner::new(code, rules).run()
}

/// Characters coloured as punctuation, and as marks. Both sets are the ones
/// the library this replaced used.
const PUNCTUATION: &str = ",.:;";
const MARKS: &str = "()={}<>-+[]|&";

struct Scanner<'a> {
    code: Vec<char>,
    rules: &'a Rules,
    spans: Vec<Span>,
    at: usize,
}

impl<'a> Scanner<'a> {
    fn new(code: &str, rules: &'a Rules) -> Self {
        Self {
            code: code.chars().collect(),
            rules,
            spans: Vec::new(),
            at: 0,
        }
    }

    fn run(mut self) -> Vec<Span> {
        while self.at < self.code.len() {
            // Every branch that answers true has advanced `self.at`, so
            // this terminates.
            let consumed = self.block_comment()
                || self.line_comment()
                || self.raw_string()
                || self.character_or_lifetime()
                || self.string()
                || self.attribute()
                || self.number()
                || self.word()
                || self.single_character();
            if !consumed {
                self.at += 1;
            }
        }
        self.spans
    }

    fn emit(&mut self, start: usize, kind: Kind) {
        if self.at > start {
            self.spans.push(Span {
                start,
                end: self.at,
                kind,
            });
        }
    }

    fn starts(&self, token: &str) -> bool {
        starts_with_at(&self.code, self.at, token)
    }

    /// Whether a line comment token here opens one; see
    /// [`Rules::line_comments_at_word_start`].
    fn at_word_start(&self) -> bool {
        self.at == 0
            || self.code[self.at - 1].is_whitespace()
            || ";|&(".contains(self.code[self.at - 1])
    }

    /// Whether only whitespace stands between the start of this line and here.
    fn at_line_start(&self) -> bool {
        let mut back = self.at as isize - 1;
        while back >= 0 && self.code[back as usize] != '\n' {
            if !self.code[back as usize].is_whitespace() {
                return false;
            }
            back -= 1;
        }
        true
    }

    fn advance_to_end_of_line(&mut self) {
        while self.at < self.code.len() && self.code[self.at] != '\n' {
            self.at += 1;
        }
    }

    /// From an open bracket through the one that matches it, or to the end
    /// if none does.
    fn advance_to_matching_bracket(&mut self) {
        let mut depth = 0i32;
        while self.at < self.code.len() {
            match self.code[self.at] {
                '[' => depth += 1,
                ']' => depth -= 1,
                _ => {}
            }
            self.at += 1;
            if depth == 0 {
                return;
            }
        }
    }

    fn block_comment(&mut self) -> bool {
        let Some(comment) = self.rules.block_comment else {
            return false;
        };
        if !self.starts(comment.open) {
            return false;
        }
        let start = self.at;
        self.at += comment.open.chars().count();
        let mut depth = 1i32;
        while self.at < self.code.len() && depth > 0 {
            // The closer is tried first so that a language whose two
            // delimiters are the same string -- CoffeeScript's `###` --
            // closes rather than nesting forever.
            if self.starts(comment.close) {
                depth -= 1;
                self.at += comment.close.chars().count();
            } else if comment.nests && self.starts(comment.open) {
                depth += 1;
                self.at += comment.open.chars().count();
            } else {
                self.at += 1;
            }
        }
        self.emit(start, Kind::Comment);
        true
    }

    fn line_comment(&mut self) -> bool {
        if !self.rules.line_comments.iter().any(|c| self.starts(c)) {
            return false;
        }
        if self.rules.line_comments_at_word_start && !self.at_word_start() {
            return false;
        }
        let start = self.at;
        self.advance_to_end_of_line();
        self.emit(start, Kind::Comment);
        true
    }

    /// Rust and RON: `b`? `r` `#`* `"` ... `"` `#`*, with no escapes inside.
    fn raw_string(&mut self) -> bool {
        if !self.rules.raw_strings {
            return false;
        }
        let mut ahead = self.at;
        if self.code.get(ahead) == Some(&'b') {
            ahead += 1;
        }
        if self.code.get(ahead) != Some(&'r') {
            return false;
        }
        ahead += 1;
        let mut hashes = 0usize;
        while self.code.get(ahead) == Some(&'#') {
            ahead += 1;
            hashes += 1;
        }
        if self.code.get(ahead) != Some(&'"') {
            return false;
        }
        let start = self.at;
        let closer: String = std::iter::once('"')
            .chain(std::iter::repeat_n('#', hashes))
            .collect();
        let closer_chars: Vec<char> = closer.chars().collect();
        let closed = find_from(&self.code, ahead + 1, &closer_chars);
        self.at = match closed {
            Some(index) => index + closer_chars.len(),
            None => self.code.len(),
        };
        self.emit(start, Kind::String);
        true
    }

    /// See [`Rules::lifetimes`]: an apostrophe that is not a character
    /// literal opens nothing.
    fn character_or_lifetime(&mut self) -> bool {
        if !self.rules.lifetimes || self.code[self.at] != '\'' {
            return false;
        }
        let Some(&next) = self.code.get(self.at + 1) else {
            return false;
        };
        if next == '\\' || self.code.get(self.at + 2) == Some(&'\'') {
            self.quoted(Quote {
                open: "'",
                close: "'",
                escapes: true,
            });
        } else {
            self.at += 1;
        }
        true
    }

    fn string(&mut self) -> bool {
        // Longest opener wins, so Kotlin's `"""` is one delimiter rather
        // than an empty string followed by a quote.
        let mut quote: Option<Quote> = None;
        for candidate in &self.rules.quotes {
            let current_len = quote.map(|q| q.open.chars().count()).unwrap_or(0);
            if self.starts(candidate.open) && candidate.open.chars().count() > current_len {
                quote = Some(*candidate);
            }
        }
        let Some(quote) = quote else {
            return false;
        };
        self.quoted(quote);
        true
    }

    fn quoted(&mut self, quote: Quote) {
        let start = self.at;
        self.at += quote.open.chars().count();
        while self.at < self.code.len() {
            if quote.escapes && self.code[self.at] == '\\' && self.at + 1 < self.code.len() {
                self.at += 2;
                continue;
            }
            if self.starts(quote.close) {
                self.at += quote.close.chars().count();
                break;
            }
            self.at += 1;
        }
        self.at = self.at.min(self.code.len());
        self.emit(start, Kind::String);
    }

    fn attribute(&mut self) -> bool {
        let start = self.at;
        match self.rules.attributes {
            Attributes::None => return false,
            Attributes::AtWord => {
                if self.code[self.at] != '@' || !is_word_start(self.code.get(self.at + 1).copied())
                {
                    return false;
                }
                self.at += 1;
                while self.at < self.code.len() && is_word_part(self.code[self.at]) {
                    self.at += 1;
                }
            }
            Attributes::HashBracket => {
                if self.code[self.at] != '#' {
                    return false;
                }
                let mut ahead = self.at + 1;
                if self.code.get(ahead) == Some(&'!') {
                    ahead += 1;
                }
                if self.code.get(ahead) != Some(&'[') {
                    return false;
                }
                self.at = ahead;
                self.advance_to_matching_bracket();
            }
            Attributes::HashLine => {
                if self.code[self.at] != '#' || !self.at_line_start() {
                    return false;
                }
                self.advance_to_end_of_line();
            }
            Attributes::LineBracket => {
                if self.code[self.at] != '[' || !self.at_line_start() {
                    return false;
                }
                self.advance_to_matching_bracket();
            }
        }
        self.emit(start, Kind::Metadata);
        true
    }

    /// A number is a run starting with a digit and carrying on through
    /// letters, digits, `_` and `.` -- which covers `0xFF`, `1_000`, `1u32`
    /// and `3.14` without a grammar for any of them.
    fn number(&mut self) -> bool {
        if !self.code[self.at].is_ascii_digit() {
            return false;
        }
        let start = self.at;
        while self.at < self.code.len() {
            let c = self.code[self.at];
            if c.is_alphanumeric() || c == '_' || c == '.' {
                self.at += 1;
            } else {
                break;
            }
        }
        self.emit(start, Kind::Literal);
        true
    }

    fn word(&mut self) -> bool {
        if !is_word_start(Some(self.code[self.at])) {
            return false;
        }
        let start = self.at;
        while self.at < self.code.len() && is_word_part(self.code[self.at]) {
            self.at += 1;
        }
        let word: String = self.code[start..self.at].iter().collect();
        if self.rules.keywords.contains(word.as_str()) {
            self.emit(start, Kind::Keyword);
        }
        true
    }

    fn single_character(&mut self) -> bool {
        let kind = if PUNCTUATION.contains(self.code[self.at]) {
            Kind::Punctuation
        } else if MARKS.contains(self.code[self.at]) {
            Kind::Mark
        } else {
            return false;
        };
        self.at += 1;
        self.emit(self.at - 1, kind);
        true
    }
}

fn is_word_start(c: Option<char>) -> bool {
    matches!(c, Some(c) if c.is_alphabetic() || c == '_')
}

fn is_word_part(c: char) -> bool {
    c.is_alphanumeric() || c == '_'
}

/// Whether `code[at..]` starts with `token`, both read as chars.
fn starts_with_at(code: &[char], at: usize, token: &str) -> bool {
    let token: Vec<char> = token.chars().collect();
    if at + token.len() > code.len() {
        return false;
    }
    code[at..at + token.len()] == token[..]
}

/// The first index at or after `from` where `code` contains `needle`, or
/// `None`.
fn find_from(code: &[char], from: usize, needle: &[char]) -> Option<usize> {
    if needle.is_empty() || from > code.len() {
        return None;
    }
    (from..=code.len().saturating_sub(needle.len())).find(|&i| code[i..i + needle.len()] == *needle)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn spans(code: &str, language: Language, kind: Kind) -> Vec<String> {
        let chars: Vec<char> = code.chars().collect();
        spans_of(code, language)
            .into_iter()
            .filter(|s| s.kind == kind)
            .map(|s| span_text(&chars, &s))
            .collect()
    }

    fn assert_spans(code: &str, language: Language, kind: Kind, expected: &[&str]) {
        assert_eq!(
            spans(code, language, kind),
            expected.to_vec(),
            "{kind:?} in: {code}"
        );
    }

    #[test]
    fn a_quoted_glob_is_one_string_not_a_comment() {
        assert_spans("x '*/a/*'", Language::Shell, Kind::String, &["'*/a/*'"]);
        assert_spans("x '*/a/*'", Language::Shell, Kind::Comment, &[]);
    }

    #[test]
    fn a_find_with_globs_has_no_comment_in_it() {
        let code = "find . -path '*/.git/*' -prune -o -name '*.kt' -print";
        assert_spans(
            code,
            Language::Shell,
            Kind::String,
            &["'*/.git/*'", "'*.kt'"],
        );
        assert_spans(code, Language::Shell, Kind::Comment, &[]);
    }

    #[test]
    fn a_url_does_not_comment_out_the_rest_of_a_shell_line() {
        let code = "curl https://example.com/x && echo done";
        assert_spans(code, Language::Shell, Kind::Comment, &[]);
        assert_spans(code, Language::Shell, Kind::Keyword, &["echo"]);
    }

    #[test]
    fn a_url_inside_a_kotlin_string_stays_a_string() {
        let code = "val url = \"https://example.com\"\nfun f() = 1";
        assert_spans(code, Language::Kotlin, Kind::Comment, &[]);
        assert_spans(
            code,
            Language::Kotlin,
            Kind::String,
            &["\"https://example.com\""],
        );
        assert_spans(code, Language::Kotlin, Kind::Keyword, &["val", "fun"]);
    }

    #[test]
    fn a_rust_attribute_is_metadata_and_the_struct_after_it_still_colours() {
        let code = "#[derive(Debug)]\nstruct A { b: u8 }";
        assert_spans(code, Language::Rust, Kind::Metadata, &["#[derive(Debug)]"]);
        assert_spans(code, Language::Rust, Kind::Comment, &[]);
        assert_spans(code, Language::Rust, Kind::Keyword, &["struct"]);
    }

    #[test]
    fn an_inner_rust_attribute_closes_at_its_own_bracket() {
        let code = "#![allow(dead_code)]\nfn f() {}";
        assert_spans(
            code,
            Language::Rust,
            Kind::Metadata,
            &["#![allow(dead_code)]"],
        );
        assert_spans(code, Language::Rust, Kind::Keyword, &["fn"]);
    }

    #[test]
    fn a_c_preprocessor_line_is_metadata_rather_than_a_comment() {
        let code = "#include <stdio.h>\nint main() { return 0; }";
        assert_spans(code, Language::C, Kind::Metadata, &["#include <stdio.h>"]);
        assert_spans(code, Language::C, Kind::Comment, &[]);
        assert_spans(code, Language::C, Kind::Keyword, &["int", "return"]);
    }

    #[test]
    fn a_kotlin_annotation_is_metadata() {
        assert_spans(
            "@Composable fun f() {}",
            Language::Kotlin,
            Kind::Metadata,
            &["@Composable"],
        );
    }

    #[test]
    fn a_hash_inside_a_kotlin_string_is_not_a_comment() {
        let code = "val c = \"#FF0000\"\nval d = 1";
        assert_spans(code, Language::Kotlin, Kind::Comment, &[]);
        assert_spans(code, Language::Kotlin, Kind::String, &["\"#FF0000\""]);
    }

    #[test]
    fn an_apostrophe_inside_a_kotlin_string_does_not_open_one() {
        let code = "val a = \"don't\"\nval b = \"x\"";
        assert_spans(
            code,
            Language::Kotlin,
            Kind::String,
            &["\"don't\"", "\"x\""],
        );
    }

    #[test]
    fn a_rust_lifetime_does_not_open_a_string_but_a_character_literal_does() {
        let code = "fn f<'a>(x: &'a str) { let c = 'x'; }";
        assert_spans(code, Language::Rust, Kind::String, &["'x'"]);
    }

    #[test]
    fn an_escaped_quote_is_inside_the_rust_character_literal() {
        assert_spans("let c = '\\'';", Language::Rust, Kind::String, &["'\\''"]);
    }

    #[test]
    fn a_rust_raw_string_keeps_its_inner_quotes() {
        let code = "let s = r#\"a \"quoted\" b\"#;";
        assert_spans(
            code,
            Language::Rust,
            Kind::String,
            &["r#\"a \"quoted\" b\"#"],
        );
    }

    #[test]
    fn a_kotlin_triple_quoted_string_is_one_string() {
        assert_spans(
            "val s = \"\"\"a \"b\" c\"\"\"",
            Language::Kotlin,
            Kind::String,
            &["\"\"\"a \"b\" c\"\"\""],
        );
    }

    #[test]
    fn a_shell_single_quoted_string_takes_no_escapes() {
        assert_spans("echo 'a\\' b", Language::Shell, Kind::String, &["'a\\'"]);
    }

    #[test]
    fn rust_and_kotlin_nest_block_comments() {
        let code = "/* a /* b */ c */ x";
        assert_spans(code, Language::Rust, Kind::Comment, &["/* a /* b */ c */"]);
        assert_spans(
            code,
            Language::Kotlin,
            Kind::Comment,
            &["/* a /* b */ c */"],
        );
    }

    #[test]
    fn c_ends_a_block_comment_at_the_first_close() {
        assert_spans(
            "/* a /* b */ c */ x",
            Language::C,
            Kind::Comment,
            &["/* a /* b */"],
        );
    }

    #[test]
    fn a_shell_comment_starts_only_at_a_word_boundary() {
        let code = "${#x} $# a#b # real";
        assert_spans(code, Language::Shell, Kind::Comment, &["# real"]);
    }

    #[test]
    fn a_hash_anywhere_is_a_python_comment() {
        assert_spans("x = 1 # note", Language::Python, Kind::Comment, &["# note"]);
    }

    #[test]
    fn a_toml_table_header_is_metadata_and_a_hash_in_a_value_is_not_a_comment() {
        let code = "[server]\ncolour = \"#FF0000\"\nport = 8080 # the real one";
        assert_spans(code, Language::Toml, Kind::Metadata, &["[server]"]);
        assert_spans(code, Language::Toml, Kind::String, &["\"#FF0000\""]);
        assert_spans(code, Language::Toml, Kind::Comment, &["# the real one"]);
        assert_spans(code, Language::Toml, Kind::Literal, &["8080"]);
    }

    #[test]
    fn a_ron_attribute_and_its_values_colour() {
        let code = "#![enable(implicit_some)]\n(count: 3, on: true)";
        assert_spans(
            code,
            Language::Ron,
            Kind::Metadata,
            &["#![enable(implicit_some)]"],
        );
        assert_spans(code, Language::Ron, Kind::Keyword, &["true"]);
        assert_spans(code, Language::Ron, Kind::Literal, &["3"]);
    }

    #[test]
    fn an_unknown_fence_language_is_none() {
        assert_eq!(fence_language(Some("brainfuck")), None);
    }

    #[test]
    fn every_language_the_fence_table_knows_has_a_scanner() {
        for language in Language::ALL {
            spans_of("x", language);
        }
    }

    /// The scanner must never panic and must never answer a span the code
    /// does not contain: the library this replaced answered a reversed
    /// range here, which crashed a card, and a fence still being written is
    /// an unterminated string or comment on every keystroke.
    #[test]
    fn spans_stay_inside_the_code_for_every_language_and_every_nasty_input() {
        let nasty = [
            "",
            "'",
            "\"",
            "\"unterminated",
            "/* unterminated",
            "###",
            "#",
            "#![",
            "[",
            "r#\"",
            "\\",
            "'''",
            "\"\"\"",
            "0x",
            "1.2.3",
            "a#b//c/*d*/'e\"f",
            "```",
            "*",
            "**",
            "~~",
            "> ",
            "- ",
            "1.",
            "[x](",
            "#######",
            "|",
            "|---|",
            "<",
            "<>",
            "http://",
            "a://",
            "\n\n   \n",
        ];
        for language in Language::ALL {
            for code in nasty {
                let chars: Vec<char> = code.chars().collect();
                let spans = spans_of(code, language);
                for s in &spans {
                    assert!(
                        s.start <= s.end && s.end <= chars.len(),
                        "{language:?} answered {s:?} for {code:?}"
                    );
                }
                let mut sorted = spans.clone();
                sorted.sort_by_key(|s| s.start);
                assert_eq!(
                    spans, sorted,
                    "{language:?} answered spans out of order for {code:?}"
                );
            }
        }
    }
}