devup_editor_syntax_highlight/

lib.rs

//! Tree-sitter-based syntax highlighting for code blocks.
//!
//! # Why this crate exists
//!
//! `devup-editor` is a Notion-style block editor. A `code` block
//! carries a `language` prop (e.g. `"rust"`, `"ts"`). When the block
//! renders, we need to turn its flat text content into a sequence of
//! styled tokens so the view layer can paint each slice with a colour.
//!
//! Doing this in Rust is deliberate. The WASM boundary for this
//! module is:
//!
//! ```text
//! (code: &str, language: &str)  →  Vec<Token { start_utf16, end_utf16, kind: &str }>
//! ```
//!
//! Both sides of that boundary are primitives, which is exactly the
//! case where porting to Rust pays for itself (see the
//! `primitive-boundary rule` in `AGENTS.md`). A full 200-line code
//! block serialised as a token list is a few KB of JSON — the parse
//! cost dominates the JSON round-trip cost by orders of magnitude.
//!
//! # Engine choice: `arborium`
//!
//! We use [`arborium`](https://docs.rs/arborium) rather than
//! `tree-sitter` + individual grammar crates for one reason only:
//! **`tree-sitter-*` grammar crates on crates.io as of 2026-01 still
//! fail `cargo build --target wasm32-unknown-unknown` with
//! `stdlib.h` not found.** The core `tree-sitter` crate was patched
//! (PR #4820, merged Sep 2025) but the fix has not propagated to
//! grammar crates. `arborium` ships its own `arborium-sysroot` and
//! re-generates every grammar against it, and its CI verifies the
//! output `.wasm` has no unsatisfied imports.
//!
//! See the librarian research report in the session history for the
//! full picture.
//!
//! # Supported language set
//!
//! We ship with the following grammars, selected to cover the
//! majority of code blocks a Notion-style editor encounters in the
//! wild:
//!
//! - `rust` — project's primary language
//! - `typescript` / `javascript` — React ecosystem consumers
//! - `python` — broad user base
//! - `json` — config-file ubiquity
//! - `bash` — README examples
//! - `html` / `css` — frontend context
//! - `markdown` — meta case (code blocks inside code blocks)
//! - `sql` — query pastes
//!
//! The measured WASM bundle cost is ~1.5 MB gzipped across all of
//! these grammars. See `.sisyphus/plans/tier-a-benchmark-decisions.md`
//! for the budget discussion. Adding another grammar is a one-line
//! change to `Cargo.toml` features, a one-line `normalize_language`
//! entry, and a one-arm addition to [`compile_grammar`]. Do it
//! deliberately — each new grammar adds roughly 50–300 KB gzipped.

use std::collections::HashMap;
use std::sync::{Arc, Mutex, OnceLock};

use arborium::advanced::{CompiledGrammar, GrammarConfig, ParseContext, Span};

/// A coarse token kind. The view layer maps each kind to a colour via
/// theme tokens. We deliberately flatten arborium's capture-name
/// taxonomy (which has ~50 distinct strings like `"keyword.function"`,
/// `"string.special.symbol"`, …) into a small enum because:
///
/// 1. A Notion-style code block doesn't need surgical precision. It
///    needs to be readable at a glance.
/// 2. Mapping every arborium capture to its own CSS variable would
///    explode `devup.json` with dozens of colour slots that nobody
///    tunes.
/// 3. Colour themes in the wild (One Dark, Dracula, GitHub) group
///    captures into ~8 slots anyway.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TokenKind {
    /// `fn`, `let`, `if`, `return`, language keywords.
    Keyword,
    /// Function names at definition or call sites.
    Function,
    /// Type names, struct names, trait names.
    Type,
    /// `42`, `0xff`, `3.14`.
    Number,
    /// `"hello"`, `'x'`, template literal bodies.
    String,
    /// `// line comments`, `/* block comments */`.
    Comment,
    /// User-facing identifiers that don't fit another slot
    /// (variables, parameters).
    Variable,
    /// Built-in constants like `true`, `false`, `null`, `None`.
    Constant,
    /// `+`, `=>`, `::`, punctuation that carries meaning.
    Operator,
}

impl TokenKind {
    /// The string handed across the WASM boundary. Kept short and
    /// lowercase so the TS side can use them as-is for CSS class
    /// suffixes (`syntax-keyword`, `syntax-string`, …).
    #[must_use]
    pub const fn as_str(self) -> &'static str {
        match self {
            Self::Keyword => "keyword",
            Self::Function => "function",
            Self::Type => "type",
            Self::Number => "number",
            Self::String => "string",
            Self::Comment => "comment",
            Self::Variable => "variable",
            Self::Constant => "constant",
            Self::Operator => "operator",
        }
    }

    /// Map a raw arborium capture name (e.g. `"keyword.function"`,
    /// `"string.special.symbol"`) to our coarse enum. Unknown
    /// captures return `None` — the caller should skip them so we
    /// don't paint random tokens.
    ///
    /// The matching is prefix-based (`"keyword.function"` →
    /// `Keyword`) because arborium's capture names are dotted
    /// namespaces where the first segment is the most important.
    #[must_use]
    pub fn from_capture(capture: &str) -> Option<Self> {
        // Strip the leaf segment; match on the root. This handles both
        // bare captures like `"keyword"` and dotted ones like
        // `"keyword.function.macro"` uniformly.
        let root = capture.split('.').next().unwrap_or(capture);
        Some(match root {
            "keyword" | "conditional" | "repeat" | "include" | "storageclass" => Self::Keyword,
            "function" | "method" => Self::Function,
            "type" | "class" | "constructor" => Self::Type,
            "number" | "float" => Self::Number,
            "string" | "character" => Self::String,
            "comment" => Self::Comment,
            "variable" | "parameter" | "property" | "field" => Self::Variable,
            "constant" | "boolean" => Self::Constant,
            "operator" | "punctuation" => Self::Operator,
            _ => return None,
        })
    }
}

/// A highlighted range within the source, ready to cross the WASM
/// boundary. Offsets are **UTF-16 code units** to match the
/// convention used everywhere else in this codebase
/// (`TextSpan`, selection, `utf16_len`) — the JS side does
/// `string.slice(token.start, token.end)` directly with these
/// numbers.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Token {
    /// UTF-16 code-unit offset where the token starts (inclusive).
    pub start_utf16: u32,
    /// UTF-16 code-unit offset where the token ends (exclusive).
    pub end_utf16: u32,
    /// Short kebab-case kind, see [`TokenKind::as_str`].
    pub kind: &'static str,
}

/// Names we accept from the editor's `language` prop.
///
/// The matching is permissive and case-insensitive. Canonical
/// targets are in `normalize_language`. In particular:
/// - `rust` / `rs` → Rust
/// - `typescript` / `ts` → TypeScript
/// - `javascript` / `js` → JavaScript
/// - `jsx` / `tsx` → JavaScript (bundle-size fallback; the dedicated
///   tree-sitter-tsx grammar would add ~1 MB uncompressed, so we
///   accept losing JSX tag/attribute granularity and TypeScript
///   type syntax for `.tsx` snippets)
/// - `python` / `py`, `json`, `bash` / `sh` / `shell` / `zsh`,
///   `html` / `htm`, `css`, `markdown` / `md`, `sql` + dialect
///   aliases → their respective grammars
///
/// Returns `false` for unsupported languages — the caller should
/// fall back to rendering the code unstyled.
#[must_use]
pub fn is_supported_language(language: &str) -> bool {
    normalize_language(language).is_some()
}

/// Canonicalise a user-supplied language name to an arborium-backed
/// grammar, or `None` if we don't bundle a grammar for it.
fn normalize_language(language: &str) -> Option<&'static str> {
    let lang = language.trim().to_ascii_lowercase();
    match lang.as_str() {
        "rust" | "rs" => Some("rust"),
        "typescript" | "ts" => Some("typescript"),
        // `jsx` and `tsx` both map to the JavaScript grammar. The
        // dedicated TSX/JSX tree-sitter parsers would give slightly
        // better coverage of the `<Component />` fragment syntax but
        // they ship their own full parse tables (they do NOT inherit
        // the JS/TS parser), which would double the bundle size. The
        // tradeoff: `.jsx`/`.tsx` code blocks get solid
        // keyword/string/comment/identifier colouring, minus
        // JSX-specific element/attribute granularity. Revisit if
        // users ask for JSX tag-name colouring.
        "javascript" | "js" | "jsx" | "tsx" => Some("javascript"),
        "python" | "py" => Some("python"),
        "json" => Some("json"),
        // Aliases favoured by common Markdown fences / editor
        // mode-lines. `shell` and `sh` map to Bash because the
        // grammar handles both POSIX sh and Bash-specific syntax
        // gracefully, and dedicated sh grammars are rarely shipped.
        "bash" | "sh" | "shell" | "zsh" => Some("bash"),
        "html" | "htm" => Some("html"),
        "css" => Some("css"),
        "markdown" | "md" => Some("markdown"),
        // Various SQL dialect aliases collapse to the generic grammar.
        // The tree-sitter-sql grammar covers standard SQL well; dialect
        // specialisations (PostgreSQL-only features, T-SQL, etc.) would
        // need separate grammars which we deliberately don't ship.
        "sql" | "postgres" | "postgresql" | "psql" | "mysql" | "sqlite" => Some("sql"),
        _ => None,
    }
}

/// Highlight `code` as `language` and return a de-duplicated,
/// left-to-right sorted list of tokens with UTF-16 offsets.
///
/// # Errors
///
/// Returns an empty vector (not an error) when:
/// - the language is not supported — callers render plain text
/// - the code is empty — no tokens to produce
///
/// Grammar-compilation or parse failures are recoverable: we return
/// an empty vector and the caller paints plain text. We never panic.
#[must_use]
pub fn highlight(code: &str, language: &str) -> Vec<Token> {
    if code.is_empty() {
        return Vec::new();
    }
    if code.len() > MAX_HIGHLIGHT_BYTES {
        // Paste-of-huge-file protection: tree-sitter happily parses
        // megabytes of code, but doing so synchronously on the main
        // thread (we are called from `useMemo` via WASM) would freeze
        // the editor for hundreds of ms. Callers render plain text.
        return Vec::new();
    }
    let Some(lang_name) = normalize_language(language) else {
        return Vec::new();
    };

    let Ok(grammar) = get_or_compile_grammar(lang_name) else {
        return Vec::new();
    };
    let Ok(mut ctx) = ParseContext::for_grammar(&grammar) else {
        return Vec::new();
    };

    let result = grammar.parse(&mut ctx, code);
    byte_spans_to_utf16_tokens(code, &result.spans)
}

/// Upper bound (in bytes) on the source we will attempt to highlight.
/// Above this size we fall back to plain text so pasting a large file
/// cannot freeze the UI thread.
///
/// 128 KiB is ~4,000 lines of typical code — comfortably above any
/// sensible code block you'd embed in a note, and well below the
/// range where tree-sitter latency becomes user-visible on mid-range
/// hardware.
///
/// The React integration enforces a tighter UTF-16 limit (64 Ki code
/// units — see `MAX_HIGHLIGHT_CHARS` in
/// `packages/react/src/plugins/code.tsx`) that always short-circuits
/// before this one for any realistic input: ASCII at the JS limit is
/// 64 KiB (half of this) and multi-byte text hits the JS limit even
/// sooner. This Rust-side bound is therefore a safety net for
/// non-React consumers (future CLI tools, server-side renderers) that
/// might call `highlight` directly with no upstream guard.
const MAX_HIGHLIGHT_BYTES: usize = 128 * 1024;

/// Global cache of compiled grammars keyed by canonical language
/// name. `CompiledGrammar::new` compiles the highlight query against
/// the parser tables, which is the expensive part; once compiled the
/// grammar is immutable and can be reused across every `highlight()`
/// call for that language.
///
/// `Mutex` over the map itself is enough — we only hold it briefly
/// (map lookup / insert), and the values are `Arc` so we hand out
/// cheap clones to callers that then drop the lock before parsing.
/// WASM is single-threaded so the lock is effectively uncontended,
/// but we still use `Mutex` rather than `RefCell` so native tests
/// (which are multi-threaded) don't panic on re-entry.
fn grammar_cache() -> &'static Mutex<HashMap<&'static str, Arc<CompiledGrammar>>> {
    static CACHE: OnceLock<Mutex<HashMap<&'static str, Arc<CompiledGrammar>>>> = OnceLock::new();
    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
}

/// Return a cached `CompiledGrammar` for the language, compiling on
/// first use. Subsequent calls for the same language are O(1) map
/// lookups.
fn get_or_compile_grammar(
    lang_name: &'static str,
) -> Result<Arc<CompiledGrammar>, arborium::advanced::GrammarError> {
    // Fast path: grammar already compiled.
    {
        let guard = grammar_cache().lock().expect("grammar cache poisoned");
        if let Some(existing) = guard.get(lang_name) {
            return Ok(Arc::clone(existing));
        }
    }
    // Slow path: compile, then insert. A parallel call for the same
    // language could race to compile twice on native; both results
    // are equivalent so we accept the occasional double-compile in
    // exchange for not holding the lock across compilation (which
    // could take tens of ms for a complex grammar).
    let compiled = compile_grammar(lang_name)?;
    let mut guard = grammar_cache().lock().expect("grammar cache poisoned");
    Ok(Arc::clone(guard.entry(lang_name).or_insert(compiled)))
}

/// Build a [`CompiledGrammar`] for the given canonical language name.
/// Callers should use [`get_or_compile_grammar`] instead — this
/// variant is the uncached path, kept as a named helper so the big
/// language `match` stays readable.
fn compile_grammar(
    lang_name: &'static str,
) -> Result<Arc<CompiledGrammar>, arborium::advanced::GrammarError> {
    // arborium ships `HIGHLIGHTS_QUERY` as a `LazyLock<String>` for
    // the languages whose highlights build on another language's
    // (TypeScript inherits JavaScript, TSX inherits TypeScript). We
    // dereference + borrow to satisfy `GrammarConfig`'s `&str`
    // contract. For languages without inheritance (Rust, JavaScript)
    // it is a plain `&'static str`.
    let (language, highlights_query, injections_query, locals_query): (
        arborium::tree_sitter::Language,
        &str,
        &str,
        &str,
    ) = match lang_name {
        "rust" => (
            arborium::lang_rust::language().into(),
            arborium::lang_rust::HIGHLIGHTS_QUERY,
            arborium::lang_rust::INJECTIONS_QUERY,
            arborium::lang_rust::LOCALS_QUERY,
        ),
        "typescript" => (
            arborium::lang_typescript::language().into(),
            &arborium::lang_typescript::HIGHLIGHTS_QUERY,
            arborium::lang_typescript::INJECTIONS_QUERY,
            arborium::lang_typescript::LOCALS_QUERY,
        ),
        "javascript" => (
            arborium::lang_javascript::language().into(),
            arborium::lang_javascript::HIGHLIGHTS_QUERY,
            arborium::lang_javascript::INJECTIONS_QUERY,
            arborium::lang_javascript::LOCALS_QUERY,
        ),
        "python" => (
            arborium::lang_python::language().into(),
            arborium::lang_python::HIGHLIGHTS_QUERY,
            arborium::lang_python::INJECTIONS_QUERY,
            arborium::lang_python::LOCALS_QUERY,
        ),
        "json" => (
            arborium::lang_json::language().into(),
            arborium::lang_json::HIGHLIGHTS_QUERY,
            arborium::lang_json::INJECTIONS_QUERY,
            arborium::lang_json::LOCALS_QUERY,
        ),
        "bash" => (
            arborium::lang_bash::language().into(),
            arborium::lang_bash::HIGHLIGHTS_QUERY,
            arborium::lang_bash::INJECTIONS_QUERY,
            arborium::lang_bash::LOCALS_QUERY,
        ),
        "html" => (
            arborium::lang_html::language().into(),
            arborium::lang_html::HIGHLIGHTS_QUERY,
            arborium::lang_html::INJECTIONS_QUERY,
            arborium::lang_html::LOCALS_QUERY,
        ),
        "css" => (
            arborium::lang_css::language().into(),
            arborium::lang_css::HIGHLIGHTS_QUERY,
            arborium::lang_css::INJECTIONS_QUERY,
            arborium::lang_css::LOCALS_QUERY,
        ),
        "markdown" => (
            arborium::lang_markdown::language().into(),
            arborium::lang_markdown::HIGHLIGHTS_QUERY,
            arborium::lang_markdown::INJECTIONS_QUERY,
            arborium::lang_markdown::LOCALS_QUERY,
        ),
        "sql" => (
            arborium::lang_sql::language().into(),
            arborium::lang_sql::HIGHLIGHTS_QUERY,
            arborium::lang_sql::INJECTIONS_QUERY,
            arborium::lang_sql::LOCALS_QUERY,
        ),
        _ => {
            // normalize_language filters this out upstream, so this
            // arm is unreachable at runtime. A debug_assert keeps the
            // invariant honest in tests; in release we degrade to a
            // LanguageError rather than panicking.
            debug_assert!(false, "unreachable lang_name: {lang_name}");
            return Err(arborium::advanced::GrammarError::LanguageError);
        }
    };
    let config = GrammarConfig {
        language,
        highlights_query,
        injections_query,
        locals_query,
    };
    CompiledGrammar::new(config).map(Arc::new)
}

/// Convert arborium's byte-indexed spans into UTF-16-indexed tokens,
/// de-duplicating overlaps and rejecting unknown capture names.
///
/// # Why the conversion is non-trivial
///
/// 1. Arborium emits spans in byte offsets; we need UTF-16 to match
///    the JS side's `string.length`-based offsets. One emoji (U+1F600)
///    is 4 bytes UTF-8 but 2 UTF-16 code units.
/// 2. Tree-sitter queries often produce **overlapping spans** at the
///    same range — e.g. `keyword` and `keyword.function` can both
///    match `fn`. We keep the highest-priority span per byte range
///    and drop duplicates, matching arborium's own HTML renderer's
///    behaviour.
/// 3. Many captures (like `punctuation.bracket` for `(` `)`) are
///    noise for a Notion-style renderer. We filter them out so the
///    output is visually calm.
fn byte_spans_to_utf16_tokens(code: &str, spans: &[Span]) -> Vec<Token> {
    // Deduplicate: for spans with the same (start, end) byte range,
    // keep the one with the highest `pattern_index`. This matches
    // arborium's `spans_to_html` convention where later patterns in
    // `highlights.scm` win.
    let mut best: Vec<&Span> = Vec::with_capacity(spans.len());
    for span in spans {
        // Skip zero-width spans arborium occasionally emits for
        // anchor-only captures.
        if span.end <= span.start {
            continue;
        }
        match best
            .iter()
            .position(|s| s.start == span.start && s.end == span.end)
        {
            Some(idx) => {
                if span.pattern_index > best[idx].pattern_index {
                    best[idx] = span;
                }
            }
            None => best.push(span),
        }
    }
    // Sort left-to-right. Arborium usually emits sorted, but we don't
    // want to rely on an internal-layout promise.
    best.sort_by_key(|s| (s.start, s.end));

    // Build a byte→UTF-16 offset map for the source once. For a
    // typical 200-line code block (say 6-8 KB) this is a cheap linear
    // scan; much cheaper than repeatedly calling `[..byte].encode_utf16()`
    // per span.
    let byte_to_utf16 = build_byte_to_utf16_map(code);

    let mut out: Vec<Token> = Vec::with_capacity(best.len());
    // Track the last range we accepted, so we can flatten overlaps
    // into a non-overlapping sequence. Tree-sitter commonly emits
    // two kinds of overlap:
    //   (a) **nested** — a child span fully contained in a parent
    //       span (`function.call` inside `call_expression`). The
    //       parent was emitted first; we drop the nested child.
    //   (b) **partial** — a later span starts inside the previous
    //       span but extends beyond it. This is rare but happens
    //       when grammar captures straddle a token boundary. Instead
    //       of dropping the whole span (which silently loses colour
    //       on the tail) we clamp its start to `prev_end_byte` and
    //       emit the remaining suffix.
    // The output is guaranteed sorted + non-overlapping — callers
    // (the React renderer) depend on this invariant.
    let mut prev_end_byte: u32 = 0;
    for span in best {
        let Some(kind) = TokenKind::from_capture(&span.capture) else {
            continue;
        };
        // Fully covered by a previous span → nested, drop it.
        if span.end <= prev_end_byte {
            continue;
        }
        // Clamp start forward on partial overlap so we don't emit
        // a range that collides with the previous token.
        let start_byte = span.start.max(prev_end_byte);
        let start_utf16 = byte_offset_to_utf16(&byte_to_utf16, start_byte);
        let end_utf16 = byte_offset_to_utf16(&byte_to_utf16, span.end);
        if end_utf16 <= start_utf16 {
            continue;
        }
        out.push(Token {
            start_utf16,
            end_utf16,
            kind: kind.as_str(),
        });
        prev_end_byte = span.end;
    }
    out
}

/// Build a vector such that `map[b]` is the UTF-16 offset of byte
/// `b`. `map` has length `code.len() + 1` so we can look up the end
/// boundary.
///
/// This is O(n) over the source; doing the same work inline per-span
/// was measured as the hot path when benchmarking on large Korean
/// text code blocks during the Phase-3 WASM benchmarks.
fn build_byte_to_utf16_map(code: &str) -> Vec<u32> {
    let mut map = Vec::with_capacity(code.len() + 1);
    let mut utf16 = 0u32;
    let mut last_byte = 0usize;
    for (byte_idx, ch) in code.char_indices() {
        // Fill the gap for continuation bytes of the previous char —
        // they all resolve to the same UTF-16 offset as the char's
        // starting byte.
        while map.len() < byte_idx {
            map.push(utf16);
        }
        map.push(utf16);
        utf16 += u32::try_from(ch.len_utf16()).unwrap_or(0);
        last_byte = byte_idx + ch.len_utf8();
    }
    // Tail: pad out to code.len() + 1 with the final UTF-16 value so
    // `byte_offset_to_utf16(code.len())` returns the total length.
    while map.len() <= code.len() {
        map.push(utf16);
    }
    // Sanity: last_byte should equal code.len() for valid UTF-8.
    debug_assert!(last_byte <= code.len());
    map
}

fn byte_offset_to_utf16(map: &[u32], byte_offset: u32) -> u32 {
    let idx = (byte_offset as usize).min(map.len().saturating_sub(1));
    map.get(idx).copied().unwrap_or(0)
}

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

    #[test]
    fn is_supported_language_accepts_aliases() {
        // Core MVP languages
        assert!(is_supported_language("rust"));
        assert!(is_supported_language("rs"));
        assert!(is_supported_language("RUST")); // case-insensitive
        assert!(is_supported_language("ts"));
        assert!(is_supported_language("typescript"));
        assert!(is_supported_language("js"));
        assert!(is_supported_language("javascript"));
        assert!(is_supported_language("tsx"));
        assert!(is_supported_language("jsx"));
        // Extended language set
        assert!(is_supported_language("python"));
        assert!(is_supported_language("py"));
        assert!(is_supported_language("json"));
        assert!(is_supported_language("bash"));
        assert!(is_supported_language("sh"));
        assert!(is_supported_language("shell"));
        assert!(is_supported_language("zsh"));
        assert!(is_supported_language("html"));
        assert!(is_supported_language("htm"));
        assert!(is_supported_language("css"));
        assert!(is_supported_language("markdown"));
        assert!(is_supported_language("md"));
        assert!(is_supported_language("sql"));
        assert!(is_supported_language("postgres"));
        assert!(is_supported_language("mysql"));
        assert!(is_supported_language("sqlite"));
        // Still unsupported
        assert!(!is_supported_language("go"));
        assert!(!is_supported_language("c"));
        assert!(!is_supported_language(""));
    }

    #[test]
    fn highlight_empty_code_returns_empty() {
        assert_eq!(highlight("", "rust"), Vec::<Token>::new());
    }

    #[test]
    fn highlight_unsupported_language_returns_empty() {
        // `go` is NOT in our bundled grammar set; the Rust side must
        // return an empty vec so the React view falls back to plain
        // text.
        assert_eq!(
            highlight("package main\nfunc main() {}", "go"),
            Vec::<Token>::new()
        );
    }

    #[test]
    fn highlight_rust_fn_declaration_includes_keyword_and_function() {
        let tokens = highlight("fn main() {}", "rust");
        // We should have at least a `keyword` token for `fn` and a
        // `function` token for `main`. We don't assert exact count or
        // ordering of operators/punctuation because arborium's grammar
        // emits implementation-dependent extras.
        assert!(
            tokens.iter().any(|t| t.kind == "keyword"),
            "expected a keyword token, got {tokens:?}"
        );
        assert!(
            tokens.iter().any(|t| t.kind == "function"),
            "expected a function token, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_rust_string_literal_tagged_as_string() {
        let tokens = highlight(r#"let s = "hello";"#, "rust");
        assert!(
            tokens.iter().any(|t| t.kind == "string"),
            "expected a string token, got {tokens:?}"
        );
        assert!(
            tokens.iter().any(|t| t.kind == "keyword"),
            "expected a keyword token for `let`, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_rust_comment_tagged_as_comment() {
        let tokens = highlight("// hello\nfn a() {}", "rust");
        assert!(
            tokens.iter().any(|t| t.kind == "comment"),
            "expected a comment token, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_typescript_async_function() {
        let tokens = highlight("async function fetch() { return 1; }", "ts");
        assert!(
            tokens.iter().any(|t| t.kind == "keyword"),
            "expected a keyword token, got {tokens:?}"
        );
        assert!(
            tokens.iter().any(|t| t.kind == "function"),
            "expected a function token, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_javascript_const_and_number() {
        let tokens = highlight("const x = 42;", "js");
        assert!(tokens.iter().any(|t| t.kind == "keyword"));
        assert!(tokens.iter().any(|t| t.kind == "number"));
    }

    #[test]
    fn partial_overlap_clamps_start_instead_of_dropping() {
        // Synthesize a span list that mimics tree-sitter emitting a
        // child whose end exceeds its parent's end (partial overlap,
        // not full nesting). `byte_spans_to_utf16_tokens` should NOT
        // drop the suffix — it should emit a clamped token for the
        // tail region.
        let code = "abcdefghij"; // 10 ASCII bytes, same UTF-16 offsets
        let spans = vec![
            Span {
                start: 0,
                end: 4,
                capture: "keyword".to_string(),
                pattern_index: 0,
            },
            // Partially overlaps the previous span: starts at 2,
            // ends at 7. Previously the whole span was skipped; now
            // we emit a `string` token for [4, 7).
            Span {
                start: 2,
                end: 7,
                capture: "string".to_string(),
                pattern_index: 1,
            },
        ];
        let out = byte_spans_to_utf16_tokens(code, &spans);
        assert_eq!(
            out,
            vec![
                Token {
                    start_utf16: 0,
                    end_utf16: 4,
                    kind: "keyword",
                },
                Token {
                    start_utf16: 4,
                    end_utf16: 7,
                    kind: "string",
                },
            ],
        );
    }

    #[test]
    fn full_nesting_still_drops_inner_span() {
        // Invariant: a span fully contained in a previous one (nested
        // capture like `function.call` inside `call_expression`)
        // stays dropped — we do not emit duplicate overlapping tokens.
        let code = "abcdefghij";
        let spans = vec![
            Span {
                start: 0,
                end: 10,
                capture: "function".to_string(),
                pattern_index: 0,
            },
            Span {
                start: 3,
                end: 6,
                capture: "keyword".to_string(),
                pattern_index: 1,
            },
        ];
        let out = byte_spans_to_utf16_tokens(code, &spans);
        assert_eq!(out.len(), 1, "nested child should be dropped: {out:?}");
        assert_eq!(out[0].kind, "function");
    }

    #[test]
    fn highlight_tokens_are_sorted_and_non_overlapping() {
        let tokens = highlight("fn a() { let x = 42; }", "rust");
        let mut prev_end = 0u32;
        for t in &tokens {
            assert!(
                t.start_utf16 >= prev_end,
                "tokens overlap: prev_end={prev_end} start={} in {tokens:?}",
                t.start_utf16
            );
            assert!(t.end_utf16 > t.start_utf16, "zero-width token: {t:?}");
            prev_end = t.end_utf16;
        }
    }

    #[test]
    fn byte_offsets_are_converted_to_utf16_for_non_ascii() {
        // "한글" is 2 UTF-16 units, "🚀" is a surrogate pair (2 UTF-16 units).
        // The code below has a Korean identifier comment + an emoji
        // comment. We just verify that when arborium reports a byte
        // offset past the multi-byte chars, the utf16 offset makes
        // sense (strictly less than the byte offset).
        let source = "// 한글 🚀\nfn main() {}";
        let tokens = highlight(source, "rust");
        let total_utf16 =
            u32::try_from(source.encode_utf16().count()).expect("test source fits in u32");
        for t in &tokens {
            assert!(
                t.end_utf16 <= total_utf16,
                "token end {} exceeds total utf16 {total_utf16}: {t:?}",
                t.end_utf16
            );
        }
    }

    #[test]
    fn token_kind_from_capture_maps_dotted_captures() {
        assert_eq!(
            TokenKind::from_capture("keyword.function"),
            Some(TokenKind::Keyword)
        );
        assert_eq!(
            TokenKind::from_capture("string.special.symbol"),
            Some(TokenKind::String)
        );
        assert_eq!(
            TokenKind::from_capture("function.call"),
            Some(TokenKind::Function)
        );
        assert_eq!(TokenKind::from_capture("bogus.capture"), None);
    }

    // ---------------------------------------------------------------
    // Per-language smoke tests for the extended set
    //
    // Each test checks a single strong signal unique to the language
    // (e.g. `def` keyword in Python, `SELECT` keyword in SQL). We do
    // not assert exhaustively — tree-sitter grammars evolve, and a
    // single robust signal is enough to prove the wiring is correct.
    // ---------------------------------------------------------------

    #[test]
    fn highlight_python_def_and_string() {
        let tokens = highlight("def greet(name):\n    return \"hi\"", "python");
        assert!(
            tokens.iter().any(|t| t.kind == "keyword"),
            "expected a keyword token for `def`/`return`, got {tokens:?}"
        );
        assert!(
            tokens.iter().any(|t| t.kind == "string"),
            "expected a string token, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_json_basic_shape() {
        // JSON grammar emits property-name captures as strings; we
        // assert both the "key" strings and the value string are
        // tagged. JSON has no keywords in the classic sense (true /
        // false are booleans → constant).
        let tokens = highlight(r#"{"name": "alice", "age": 30}"#, "json");
        assert!(
            tokens.iter().any(|t| t.kind == "string"),
            "expected string tokens, got {tokens:?}"
        );
        assert!(
            tokens.iter().any(|t| t.kind == "number"),
            "expected a number token for `30`, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_bash_comment_and_keyword() {
        let tokens = highlight("#!/usr/bin/env bash\nif true; then echo hi; fi", "bash");
        assert!(
            tokens.iter().any(|t| t.kind == "comment"),
            "expected a comment token for shebang, got {tokens:?}"
        );
        assert!(
            tokens.iter().any(|t| t.kind == "keyword"),
            "expected a keyword token for `if`/`then`/`fi`, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_bash_accepts_sh_alias() {
        // The `sh` / `shell` / `zsh` aliases should all land on the
        // bash grammar and produce tokens. This guards the
        // normalize_language routing.
        let tokens = highlight("echo hello", "sh");
        assert!(
            !tokens.is_empty(),
            "expected `sh` alias to route to bash grammar and produce tokens"
        );
    }

    #[test]
    fn highlight_html_tag_and_attribute() {
        // The HTML grammar's highlights emit tag names and attribute
        // names as captures. We only assert "some token came out" —
        // the exact capture names vary across tree-sitter-html
        // versions. The cross-cutting `function` / `keyword` / `type`
        // bucket our coarse enum uses absorbs those variations.
        let tokens = highlight(r#"<div class="box">hello</div>"#, "html");
        assert!(!tokens.is_empty(), "expected html tokens, got {tokens:?}");
    }

    #[test]
    fn highlight_css_selector_and_property() {
        let tokens = highlight(".card { color: red; font-size: 14px; }", "css");
        assert!(!tokens.is_empty(), "expected css tokens, got {tokens:?}");
        // A number token is the most stable signal in CSS: 14 from
        // `14px`. Colour keywords vary (some grammars tag `red` as
        // constant, others as variable), so we don't pin that.
        assert!(
            tokens.iter().any(|t| t.kind == "number"),
            "expected a number token, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_markdown_produces_tokens() {
        // Markdown's "highlighting" is structural (heading markers,
        // list bullets, emphasis). The exact captures differ across
        // tree-sitter-markdown versions, so we only assert that
        // *something* came out — empty result would mean routing
        // broke or the grammar failed to parse.
        let tokens = highlight("# Title\n\nSome **bold** text.\n", "markdown");
        assert!(
            !tokens.is_empty(),
            "expected markdown tokens, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_sql_select_keyword() {
        let tokens = highlight("SELECT id, name FROM users WHERE age > 18;", "sql");
        // The SQL grammar's highlights.scm is unusual: it tags
        // numeric literals under the `string` capture root rather
        // than `number`. This is a grammar-authoring choice we can't
        // override without rewriting queries. The important assertion
        // is that SELECT/FROM/WHERE come out as keywords.
        assert!(
            tokens.iter().any(|t| t.kind == "keyword"),
            "expected a keyword token for `SELECT`/`FROM`/`WHERE`, got {tokens:?}"
        );
        assert!(
            tokens.iter().any(|t| t.kind == "type"),
            "expected a type token for the table name `users`, got {tokens:?}"
        );
    }

    #[test]
    fn highlight_sql_alias_routes_to_sql_grammar() {
        // `postgres` / `mysql` / `sqlite` all collapse to the generic
        // SQL grammar. A SELECT statement is common to all dialects
        // and must tokenise.
        for alias in ["postgres", "postgresql", "mysql", "sqlite"] {
            let tokens = highlight("SELECT 1;", alias);
            assert!(
                !tokens.is_empty(),
                "expected `{alias}` alias to produce tokens"
            );
        }
    }
}