splintr 0.15.0

//! Vocabulary normalization, merge-rank construction and dialect dispatch.

use rustc_hash::FxHashMap;

use super::loader::{
    build_merge_ranks, byte_level_pattern, find_special_token_id, normalize_wordpiece_vocab,
    unigram_prefix_space,
};
use super::{from_gguf_vocab, GgufVocab, GgufVocabError};
use crate::core::tokenizer::{GPT2_PATTERN, LLAMA3_PATTERN, QWEN2_PATTERN};

fn v(items: &[&str]) -> Vec<String> {
    items.iter().map(|s| (*s).to_owned()).collect()
}

// ── WordPiece vocab normalization ────────────────────────────────────────────

/// The `nomic-embed-text-v1.5` shape: `▁`-marked word-initial pieces, bare
/// continuations, bracketed specials.
#[test]
fn sentencepiece_marked_bert_vocab_is_converted_to_wordpiece() {
    let got = normalize_wordpiece_vocab(v(&[
        "[PAD]", "[CLS]", "[SEP]", "[UNK]", "▁the", "▁hello", "s", "ing", "▁!", "▁1",
    ]));
    assert_eq!(
        got,
        v(&["[PAD]", "[CLS]", "[SEP]", "[UNK]", "the", "hello", "##s", "##ing", "!", "1",]),
        "▁X must become X, bare X must become ##X, specials must be untouched"
    );
}

/// A vocab already in WordPiece convention must be returned byte-identical —
/// otherwise fixing one model's tokenizer would break every other BERT GGUF.
#[test]
fn already_wordpiece_vocab_is_left_untouched() {
    let original = v(&["[PAD]", "[CLS]", "the", "##s", "hello", "!"]);
    assert_eq!(normalize_wordpiece_vocab(original.clone()), original);
}

/// Mixed marking (some `▁`, some `##`) means the file is already using the
/// WordPiece continuation marker, so rewriting would corrupt it.
#[test]
fn mixed_marking_is_left_untouched() {
    let original = v(&["▁the", "##s", "hello"]);
    assert_eq!(normalize_wordpiece_vocab(original.clone()), original);
}

/// No `▁` anywhere → nothing to convert.
#[test]
fn unmarked_vocab_is_left_untouched() {
    let original = v(&["the", "hello", "world"]);
    assert_eq!(normalize_wordpiece_vocab(original.clone()), original);
}

// ── Byte-level BPE merge ranks ───────────────────────────────────────────────

fn rank(map: &FxHashMap<Vec<u8>, u32>, token: &str) -> u32 {
    match map.get(token.as_bytes()) {
        Some(rank) => *rank,
        None => panic!("{token:?} has no merge rank"),
    }
}

/// Merge priority comes from the `merges` list order, not from token id — the
/// two disagree in real vocabularies, and using ids silently changes every
/// tokenization.
#[test]
fn merge_priority_follows_list_order_not_token_id() {
    // Ids put "lo" before "he", but the merges list puts "he" first.
    let tokens = v(&["h", "e", "l", "o", "lo", "he", "hel", "hello"]);
    let ranks = build_merge_ranks(&v(&["h e", "he l", "l o", "hel lo"]), &tokens);

    assert!(
        rank(&ranks, "he") < rank(&ranks, "lo"),
        "\"he\" is earlier in the merges list, so it must merge first regardless \
         of \"lo\" having the lower token id"
    );
    assert!(rank(&ranks, "he") < rank(&ranks, "hel"));
    assert!(rank(&ranks, "hel") < rank(&ranks, "hello"));
}

/// Single characters are never a merge result, so they must all rank below every
/// merge — multi-byte UTF-8 has to coalesce before any real merge runs.
#[test]
fn the_base_alphabet_outranks_every_merge() {
    let tokens = v(&["a", "b", "c", "ab", "abc"]);
    let ranks = build_merge_ranks(&v(&["a b", "ab c"]), &tokens);

    let base_max = ["a", "b", "c"]
        .iter()
        .map(|t| rank(&ranks, t))
        .max()
        .unwrap_or_default();
    let merge_min = ["ab", "abc"]
        .iter()
        .map(|t| rank(&ranks, t))
        .min()
        .unwrap_or_default();
    assert!(
        base_max < merge_min,
        "every base-alphabet token must rank below every merge"
    );
}

/// Byte-level tokens spell real spaces as `Ġ`, so only the first space in a
/// merge entry is the separator — splitting on all of them would corrupt any
/// merge involving a space token.
#[test]
fn only_the_first_space_separates_a_merge_entry() {
    let tokens = v(&["Ġ", "a", "Ġa"]);
    let ranks = build_merge_ranks(&v(&["Ġ a"]), &tokens);
    assert!(
        ranks.contains_key("Ġa".as_bytes()),
        "the merge result must be the concatenation \"Ġa\""
    );
}

/// A merge naming a token that is not in the vocab must not displace or
/// renumber the ranks of tokens that are.
#[test]
fn merges_referencing_absent_tokens_do_not_disturb_the_rest() {
    let tokens = v(&["a", "b", "ab"]);
    let ranks = build_merge_ranks(&v(&["a b", "z z"]), &tokens);
    assert!(rank(&ranks, "a") < rank(&ranks, "ab"));
    assert!(rank(&ranks, "b") < rank(&ranks, "ab"));
}

// ── Pre-tokenizer selection ──────────────────────────────────────────────────

/// The `pre` name selects the split pattern, and the three families must not
/// collapse onto one another.
#[test]
fn pre_tokenizer_names_select_distinct_patterns() {
    let gpt2 = byte_level_pattern(None).expect("absent `pre` is llama.cpp's GPT-2 default");
    assert_eq!(
        byte_level_pattern(Some("default")).expect("default"),
        gpt2,
        "`default` is the same GPT-2 split as an absent key"
    );
    assert_eq!(
        byte_level_pattern(Some("jina-v2-code")).expect("jina"),
        gpt2
    );

    let qwen = byte_level_pattern(Some("qwen2")).expect("qwen2");
    let llama = byte_level_pattern(Some("llama-bpe")).expect("llama-bpe");
    assert_ne!(qwen, gpt2);
    assert_ne!(llama, gpt2);
    assert_ne!(qwen, llama);
}

/// Every `pre` name llama.cpp resolves to a single expression equal to
/// [`GPT2_PATTERN`] must resolve here too — the names come from
/// `LLAMA_VOCAB_PRE_TYPE_GPT2`/`MPT`/`OLMO`/`JAIS`/`TRILLION`/`GRANITE_DOCLING`,
/// which share one `regex_exprs` list.
#[test]
fn gpt2_family_pre_names_all_select_the_gpt2_pattern() {
    for name in [
        "gpt-2",
        "phi-2",
        "jina-es",
        "jina-de",
        "gigachat",
        "jina-v2-es",
        "jina-v2-de",
        "a.x-4.0",
        "mellum",
        "modern-bert",
        "jina-v1-en",
        "jina-v2-code",
        "roberta-bpe",
        "exaone4",
        "mpt",
        "olmo",
        "jais",
        "trillion",
        "granite-docling",
    ] {
        assert_eq!(
            byte_level_pattern(Some(name)).unwrap_or(&["<refused>"]),
            &[GPT2_PATTERN],
            "`{name}` names llama.cpp's GPT-2 split"
        );
    }
}

/// The `pre` names whose enum value shares the Qwen2 `regex_exprs` list —
/// `QWEN2`, `STABLELM2`, `HUNYUAN`, `SOLAR_OPEN` (one `case` label) and `GROK_2`
/// (a byte-identical copy of that same expression).
#[test]
fn qwen2_family_pre_names_all_select_the_qwen2_pattern() {
    for name in [
        "qwen2",
        "deepseek-r1-qwen",
        "kormo",
        "megrez",
        "stablelm2",
        "hunyuan",
        "solar-open",
        "grok-2",
    ] {
        assert_eq!(
            byte_level_pattern(Some(name)).unwrap_or(&["<refused>"]),
            &[QWEN2_PATTERN],
            "`{name}` names llama.cpp's Qwen2 split"
        );
    }
}

/// The `pre` names whose enum value carries llama.cpp's Llama-3 expression —
/// `LLAMA3`, `DBRX`/`SMAUG` (one `case` label) and `CHATGLM4`, whose three
/// single-expression lists are byte-identical to one another.
#[test]
fn llama3_family_pre_names_all_select_the_llama3_pattern() {
    for name in ["llama-bpe", "llama3", "dbrx", "smaug-bpe", "glm4"] {
        assert_eq!(
            byte_level_pattern(Some(name)).unwrap_or(&["<refused>"]),
            &[LLAMA3_PATTERN],
            "`{name}` names llama.cpp's Llama-3 split"
        );
    }
}

/// The `pre` names whose enum value emits SEVERAL expressions, and the exact
/// list each one emits.
///
/// A list is applied pass by pass, so both its contents and its ORDER decide the
/// split; asserting the whole slice pins both. `falcon` cutting digit triples
/// before the GPT-2 split rather than after would silently produce different
/// ids.
#[test]
fn multi_pass_pre_names_select_their_full_expression_list() {
    let falcon = byte_level_pattern(Some("falcon")).expect("falcon");
    assert_eq!(falcon.len(), 3, "FALCON emits three expressions");
    assert_eq!(
        falcon[1], GPT2_PATTERN,
        "falcon's middle pass is llama.cpp's GPT-2 split"
    );
    assert_eq!(falcon[2], r"[0-9][0-9][0-9]");

    // One `case` label in llama.cpp covers all seven of these.
    let starcoder = byte_level_pattern(Some("starcoder")).expect("starcoder");
    assert_eq!(starcoder, &[r"\p{N}", GPT2_PATTERN]);
    for name in [
        "refact",
        "command-r",
        "smollm",
        "codeshell",
        "exaone",
        "minerva-7b",
    ] {
        assert_eq!(
            byte_level_pattern(Some(name)).unwrap_or(&["<refused>"]),
            starcoder,
            "`{name}` shares llama.cpp's STARCODER `case` label"
        );
    }

    let coder = byte_level_pattern(Some("deepseek-coder")).expect("deepseek-coder");
    let llm = byte_level_pattern(Some("deepseek-llm")).expect("deepseek-llm");
    assert_eq!(coder.len(), 5);
    assert_eq!(llm.len(), 6);
    // Both isolate line breaks FIRST, which is what keeps the later `$` in
    // deepseek-llm's `\s+$` from ever meeting a span containing a newline.
    assert_eq!(coder[0], r"[\r\n]");
    assert_eq!(llm[0], r"[\r\n]");
    // The two DeepSeek dialects differ on digits: runs versus single digits.
    assert_eq!(coder[4], r"\p{N}");
    assert_eq!(llm[5], r"\p{N}+");
    assert_ne!(coder, llm);
}

/// A multi-pass list is not the alternation of its expressions — the passes
/// compose, so no single-pattern name may resolve to the same thing.
#[test]
fn multi_pass_lists_are_never_a_single_expression() {
    for name in [
        "falcon",
        "starcoder",
        "refact",
        "command-r",
        "deepseek-coder",
        "deepseek-llm",
    ] {
        let list = byte_level_pattern(Some(name)).unwrap_or(&[]);
        assert!(
            list.len() > 1,
            "`{name}` is a sequence of passes, not one pattern"
        );
    }
}

/// Names llama.cpp resolves to an expression list splintr does not reproduce
/// verbatim must stay refused rather than be approximated by the nearest
/// pattern: a wrong split is invisible downstream.
#[test]
fn unreproduced_pre_names_stay_refused() {
    for name in [
        // multi-expression lists not transcribed here
        "deepseek-v3",
        "chameleon",
        "viking",
        "youtu",
        "superbpe",
        "afmoe",
        // one expression that is only a trigger for a hand-written scanner
        "kimi-k2",
        // single expression, but no byte-identical splintr constant
        "chatglm-bpe",
        "jais-2",
        "qwen35",
        "tekken",
        "gpt-4o",
        "llama4",
        "minimax-m2",
        "tiny_aya",
        "bailingmoe",
        "seed-coder",
        "exaone-moe",
        "poro-chat",
        "bloom",
        "gpt3-finnish",
    ] {
        assert!(
            matches!(
                byte_level_pattern(Some(name)),
                Err(GgufVocabError::UnsupportedPreTokenizer(ref got)) if got == name
            ),
            "`{name}` has no byte-identical splintr pattern and must be refused"
        );
    }
}

/// An unrecognised pre-tokenizer is refused, never defaulted: a wrong split is
/// invisible downstream because every id it produces is still in range.
#[test]
fn unknown_pre_tokenizer_is_refused_not_guessed() {
    assert!(matches!(
        byte_level_pattern(Some("some-future-pre")),
        Err(GgufVocabError::UnsupportedPreTokenizer(name)) if name == "some-future-pre"
    ));
}

// ── Flag resolution ──────────────────────────────────────────────────────────

/// `jina-embeddings-v3`'s shape: `add_space_prefix = false` with
/// `remove_extra_whitespaces = true` still marks the first word, because
/// llama.cpp's Unigram normalizer ORs the two flags.
#[test]
fn unigram_prefix_space_ors_the_two_flags() {
    let with = |space: Option<bool>, extra: Option<bool>| {
        unigram_prefix_space(&GgufVocab {
            add_space_prefix: space,
            remove_extra_whitespaces: extra,
            ..GgufVocab::default()
        })
    };
    assert!(with(None, None), "add_space_prefix defaults to true");
    assert!(
        with(Some(false), Some(true)),
        "remove_extra_whitespaces alone must still mark the first word"
    );
    assert!(with(Some(true), Some(false)));
    assert!(
        !with(Some(false), Some(false)),
        "neither flag set means the first word stays unmarked"
    );
    assert!(
        !with(Some(false), None),
        "remove_extra_whitespaces defaults to false"
    );
}

/// The vocabulary's own string is ground truth: a file whose `[UNK]` sits at a
/// different id than `unknown_token_id` claims would otherwise emit an id that
/// decodes to some other token.
#[test]
fn special_token_lookup_prefers_the_vocab_over_the_metadata() {
    let tokens = v(&["[PAD]", "[UNK]", "the"]);
    let vocab = GgufVocab {
        unknown_token_id: Some(99),
        ..GgufVocab::default()
    };
    assert_eq!(find_special_token_id(&tokens, &vocab, "[UNK]", 0), 1);
}

/// With no matching string, the declared id is used — and with neither, the
/// caller's default.
#[test]
fn special_token_lookup_falls_back_to_metadata_then_default() {
    let tokens = v(&["a", "b"]);
    let declared = GgufVocab {
        unknown_token_id: Some(7),
        ..GgufVocab::default()
    };
    assert_eq!(find_special_token_id(&tokens, &declared, "[UNK]", 0), 7);
    assert_eq!(
        find_special_token_id(&tokens, &GgufVocab::default(), "[UNK]", 3),
        3
    );
}

// ── Dialect dispatch ─────────────────────────────────────────────────────────

/// A vocabulary whose algorithm we do not implement is refused rather than run
/// through whichever backend happens to accept its data.
#[test]
fn unsupported_model_is_refused() {
    let vocab = GgufVocab {
        model: "rwkv".to_owned(),
        tokens: v(&["a"]),
        ..GgufVocab::default()
    };
    assert!(matches!(
        from_gguf_vocab(vocab),
        Err(GgufVocabError::UnsupportedModel(name)) if name == "rwkv"
    ));
}

#[test]
fn empty_vocabulary_is_refused() {
    assert!(matches!(
        from_gguf_vocab(GgufVocab {
            model: "llama".to_owned(),
            ..GgufVocab::default()
        }),
        Err(GgufVocabError::EmptyVocab)
    ));
}

/// byte-level BPE *is* its merge list, so a `gpt2` vocabulary without one cannot
/// be reconstructed and must not be approximated from the vocabulary.
#[test]
fn gpt2_without_merges_is_refused() {
    let vocab = GgufVocab {
        model: "gpt2".to_owned(),
        tokens: v(&["a", "b", "ab"]),
        ..GgufVocab::default()
    };
    assert!(matches!(
        from_gguf_vocab(vocab),
        Err(GgufVocabError::MissingMerges)
    ));
}

fn llama_vocab() -> GgufVocab {
    GgufVocab {
        model: "llama".to_owned(),
        tokens: v(&["<unk>", "<s>", "</s>", "▁hello", "▁world"]),
        bos_token_id: Some(1),
        eos_token_id: Some(2),
        ..GgufVocab::default()
    }
}

/// llama.cpp's defaults for a SentencePiece BPE vocabulary: BOS is prepended,
/// EOS is not appended. The policy owns both — the backend was built with
/// neither id, so nothing else can insert them.
#[test]
fn llama_prepends_bos_and_omits_eos_by_default() {
    let tok = from_gguf_vocab(llama_vocab()).expect("builds");
    assert_eq!(tok.family(), "Spm");

    let ids = tok.encode("hello world");
    assert_eq!(ids.first(), Some(&1), "add_bos_token defaults to true");
    assert_ne!(ids.last(), Some(&2), "add_eos_token defaults to false");
    assert_eq!(
        tok.encode_raw("hello world").as_slice(),
        &ids[1..],
        "the boundary token must come from the policy, not the backend"
    );
    assert_eq!(tok.eos_token_id(), Some(2));
    assert!(tok.is_eos(2));
}

/// The file's own flags win over the defaults, in both directions.
#[test]
fn llama_honours_the_declared_boundary_flags() {
    let tok = from_gguf_vocab(GgufVocab {
        add_bos_token: Some(false),
        add_eos_token: Some(true),
        ..llama_vocab()
    })
    .expect("builds");

    let ids = tok.encode("hello");
    assert_ne!(ids.first(), Some(&1));
    assert_eq!(ids.last(), Some(&2));
}

/// A flag asking for a boundary token the file never gives an id for adds
/// nothing — there is no id to add.
#[test]
fn a_boundary_flag_without_an_id_adds_nothing() {
    let tok = from_gguf_vocab(GgufVocab {
        bos_token_id: None,
        ..llama_vocab()
    })
    .expect("builds");
    assert_eq!(tok.encode("hello"), tok.encode_raw("hello"));
}

/// `t5` is the one dialect whose defaults ask for both boundaries.
#[test]
fn t5_wraps_with_both_boundaries_by_default() {
    let tok = from_gguf_vocab(GgufVocab {
        model: "t5".to_owned(),
        tokens: v(&["<unk>", "<s>", "</s>", "▁hi"]),
        bos_token_id: Some(1),
        eos_token_id: Some(2),
        ..GgufVocab::default()
    })
    .expect("builds");

    assert_eq!(tok.family(), "Unigram");
    let ids = tok.encode("hi");
    assert_eq!(ids.first(), Some(&1));
    assert_eq!(ids.last(), Some(&2));
}

/// BERT states its boundaries as `[CLS]`/`[SEP]` in the vocabulary rather than
/// through `add_bos_token`/`add_eos_token` — but stating them is not placing
/// them, and a caller asking `encode` for "the sequence this model was trained
/// on" needs them placed. Both references agree: HuggingFace's
/// `all-MiniLM-L6-v2` `tokenizer.json` declares `[CLS] A [SEP]` (measured with
/// `tokenizers` 0.22.1: `"hello world"` → `[101, 7592, 2088, 102]`), and
/// llama.cpp's WPM path prepends CLS and appends SEP whenever `add_special` is
/// set. So the two containers of one checkpoint must answer identically, and
/// the `bos_token_id`/`eos_token_id` flags below must still contribute nothing.
///
/// `encode_raw` stays the bare content tokens — that is the surface
/// `examples/verify_gguf.rs` scores against llama.cpp's `add_special = false`
/// fixtures.
#[test]
fn bert_wraps_with_cls_sep_and_keeps_the_named_ids() {
    let tok = from_gguf_vocab(GgufVocab {
        model: "bert".to_owned(),
        tokens: v(&["[PAD]", "[UNK]", "[CLS]", "[SEP]", "the"]),
        // Set even though BERT ignores them: they must not leak into the ids.
        add_bos_token: Some(true),
        add_eos_token: Some(true),
        bos_token_id: Some(2),
        eos_token_id: Some(3),
        ..GgufVocab::default()
    })
    .expect("builds");

    assert_eq!(tok.family(), "WordPiece");
    assert_eq!(tok.encode_raw("the"), vec![4]);
    assert_eq!(
        tok.encode("the"),
        vec![2, 4, 3],
        "[CLS] A [SEP], as both HuggingFace and llama.cpp produce"
    );
    assert_eq!(
        tok.encode_pair("the", "the")
            .expect("bert defines a pair template"),
        vec![2, 4, 3, 4, 3],
        "[CLS] A [SEP] B [SEP] — the shape a reranker head was trained on"
    );
    assert_eq!(tok.policy().single_overhead(), 2);

    assert_eq!(tok.special_token_id("[CLS]"), Some(2));
    assert_eq!(tok.special_token_id("[SEP]"), Some(3));
    assert_eq!(tok.special_token_id("[UNK]"), Some(1));
    assert_eq!(tok.special_token_id("[PAD]"), Some(0));
}

/// The other half of the rule: a vocabulary that names no `[CLS]`/`[SEP]` gets
/// the identity policy. Boundary tokens are placed only when the file states
/// which ids they are — never invented from a position or a default.
#[test]
fn bert_without_cls_sep_keeps_the_identity_policy() {
    let tok = from_gguf_vocab(GgufVocab {
        model: "bert".to_owned(),
        tokens: v(&["[PAD]", "[UNK]", "the"]),
        add_bos_token: Some(true),
        add_eos_token: Some(true),
        ..GgufVocab::default()
    })
    .expect("builds");

    assert_eq!(tok.encode("the"), tok.encode_raw("the"));
    assert_eq!(tok.policy().single_overhead(), 0);
    assert_eq!(tok.special_token_id("[CLS]"), None);
}

/// Decode drops the specials the FILE declares, whatever they are spelled.
///
/// The two vocabularies below are identical — same `token_type` (3 == CONTROL),
/// same `bos`/`eos`/`unknown` ids, same content tokens — and differ only in how
/// their special tokens are named. A surface-string skip rule kept the first
/// one's `[CLS]`/`[SEP]` out of the text and leaked the second one's
/// `<s>`/`</s>` into it, though the file declares exactly the same ids as
/// special in both. Every other dialect (`t5`, `llama`) already skips by id.
#[test]
fn bert_decode_drops_declared_specials_whatever_they_are_named() {
    fn decoded(cls: &str, sep: &str, unk: &str, ids: &[u32]) -> String {
        from_gguf_vocab(GgufVocab {
            model: "bert".to_owned(),
            tokens: v(&[cls, sep, unk, "hello", "world", "##ing"]),
            token_type: Some(vec![3, 3, 3, 1, 1, 1]),
            bos_token_id: Some(0),
            eos_token_id: Some(1),
            unknown_token_id: Some(2),
            ..GgufVocab::default()
        })
        .expect("builds")
        .decode(ids)
        .expect("decodes")
    }

    assert_eq!(
        decoded("[CLS]", "[SEP]", "[UNK]", &[0, 3, 4, 1]),
        "hello world"
    );
    assert_eq!(
        decoded("<s>", "</s>", "<unk>", &[0, 3, 4, 1]),
        "hello world",
        "the file declares ids 0 and 1 special; their spelling is not the rule"
    );

    // The declared unknown id is dropped in both spellings too, matching what
    // the `t5`/`llama` backends do with theirs.
    assert_eq!(
        decoded("[CLS]", "[SEP]", "[UNK]", &[3, 2, 4]),
        "hello world"
    );
    assert_eq!(decoded("<s>", "</s>", "<unk>", &[3, 2, 4]), "hello world");
}

/// The `t5` arm must declare the file's specials to its backend the way the
/// `llama` arm does — it was the one arm that never called
/// `with_special_decode_ids`, so the ids it passes `None` for leaked into
/// `decode()`.
///
/// The Unigram backend skips its OWN bos/eos/unk fields, but two of the three
/// are not the file's: `build_unigram` passes `None` for BOS (boundaries belong
/// to the policy), and the backend resolves its unk by the spelling `<unk>` /
/// `<UNK>`, so a vocabulary naming its unknown piece anything else has no unk to
/// skip. This vocabulary is built to expose exactly those two — its BOS is
/// declared, and its unknown piece is spelled `<unknown>` — so both surfaces
/// used to survive into the decoded text.
///
/// The set is the same three metadata ids the `llama` arm chooses, and
/// deliberately not every `token_type == 3` (CONTROL) id: `▁hi` below is CONTROL
/// here purely to show that the broader rule is not the one in force, since a
/// CONTROL token the file never names as a boundary still decodes.
#[test]
fn t5_decode_drops_the_specials_the_file_declares() {
    let tok = from_gguf_vocab(GgufVocab {
        model: "t5".to_owned(),
        tokens: v(&["<s>", "</s>", "<unknown>", "▁hello", "▁world", "▁hi"]),
        token_type: Some(vec![3, 3, 3, 1, 1, 3]),
        bos_token_id: Some(0),
        eos_token_id: Some(1),
        unknown_token_id: Some(2),
        ..GgufVocab::default()
    })
    .expect("builds");

    assert_eq!(tok.family(), "Unigram");
    assert_eq!(
        tok.decode(&[0, 3, 4, 1]).expect("decodes"),
        "hello world",
        "the declared BOS/EOS must not reach the text"
    );
    assert_eq!(
        tok.decode(&[3, 2, 4]).expect("decodes"),
        "hello world",
        "the declared unknown id must not reach the text, whatever it is spelled"
    );
    assert_eq!(
        tok.decode(&[3, 5]).expect("decodes"),
        "hello hi",
        "a CONTROL token the file never names as a special still decodes"
    );
}

/// CONTROL-flagged tokens are the special tokens of a `gpt2` vocabulary, and
/// they must be reachable by name as well as matched in the input.
#[test]
fn gpt2_control_tokens_become_named_specials() {
    let tok = from_gguf_vocab(GgufVocab {
        model: "gpt2".to_owned(),
        tokens: v(&["a", "b", "ab", "<|endoftext|>"]),
        merges: Some(v(&["a b"])),
        token_type: Some(vec![1, 1, 1, 3]),
        eos_token_id: Some(3),
        ..GgufVocab::default()
    })
    .expect("builds");

    assert_eq!(tok.family(), "BPE");
    assert_eq!(tok.special_token_id("<|endoftext|>"), Some(3));
    assert_eq!(
        tok.special_token_id("ab"),
        None,
        "only CONTROL-flagged tokens are special"
    );
    assert_eq!(tok.eos_token_id(), Some(3));
    assert_eq!(
        tok.encode_raw("ab<|endoftext|>"),
        vec![2, 3],
        "a control token in the text stays whole"
    );
}

// ── Control tokens across every dialect ──────────────────────────────────────
//
// A chat template is assembled by splicing markers like `<start_of_turn>` into
// the prompt string. If the backend does not match them, they shatter into
// content pieces; if the policy does not name them, the caller cannot splice the
// id instead. Both failures are invisible — the ids stay in range and decode
// back to the original string — so every dialect is pinned here.

/// `llama` (SPM-BPE): the dialect that had no matcher at all, so a Gemma-style
/// chat marker was silently ground into fragments.
///
/// The vocabulary carries the pieces a real SPM file would (`hi` before `▁hi`),
/// so the gap after the marker has to complete a merge chain rather than land on
/// a single entry.
#[test]
fn llama_control_tokens_are_matched_and_named() {
    let tok = from_gguf_vocab(GgufVocab {
        tokens: v(&[
            "<unk>",
            "<s>",
            "</s>",
            "<start_of_turn>",
            "▁",
            "h",
            "i",
            "hi",
            "▁hi",
        ]),
        token_type: Some(vec![3, 3, 3, 3, 1, 1, 1, 1, 1]),
        ..llama_vocab()
    })
    .expect("builds");

    assert_eq!(tok.family(), "Spm");
    assert_eq!(tok.special_token_id("<start_of_turn>"), Some(3));
    assert_eq!(
        tok.encode_raw("<start_of_turn>hi"),
        vec![3, 8],
        "the marker is one id, and the text after it still merges to a whole word"
    );
    assert_eq!(
        tok.special_token_id("▁hi"),
        None,
        "only CONTROL-flagged tokens are special"
    );
}

/// `t5` (Unigram): the map was empty, so nothing resolved by name.
#[test]
fn t5_control_tokens_are_matched_and_named() {
    let tok = from_gguf_vocab(GgufVocab {
        model: "t5".to_owned(),
        tokens: v(&["<unk>", "<s>", "</s>", "▁hi", "<start_of_turn>"]),
        scores: Some(vec![-10.0, -10.0, -10.0, -1.0, -10.0]),
        token_type: Some(vec![3, 3, 3, 1, 3]),
        bos_token_id: Some(1),
        eos_token_id: Some(2),
        ..GgufVocab::default()
    })
    .expect("builds");

    assert_eq!(tok.family(), "Unigram");
    assert_eq!(tok.special_token_id("<start_of_turn>"), Some(4));
    assert_eq!(tok.encode_raw("<start_of_turn>hi"), vec![4, 3]);
}

/// `bert` (WordPiece): the control map must be merged into the `[UNK]`/`[CLS]`/
/// `[SEP]` lookups, never replace them.
#[test]
fn bert_control_tokens_are_matched_without_losing_the_bracketed_ids() {
    let tok = from_gguf_vocab(GgufVocab {
        model: "bert".to_owned(),
        tokens: v(&["[PAD]", "[UNK]", "[CLS]", "[SEP]", "the", "<start_of_turn>"]),
        token_type: Some(vec![3, 3, 3, 3, 1, 3]),
        ..GgufVocab::default()
    })
    .expect("builds");

    assert_eq!(tok.family(), "WordPiece");
    assert_eq!(tok.special_token_id("<start_of_turn>"), Some(5));
    assert_eq!(tok.encode_raw("<start_of_turn>the"), vec![5, 4]);

    // The pre-existing lookups must survive the merge.
    assert_eq!(tok.special_token_id("[UNK]"), Some(1));
    assert_eq!(tok.special_token_id("[CLS]"), Some(2));
    assert_eq!(tok.special_token_id("[SEP]"), Some(3));
    assert_eq!(tok.special_token_id("[PAD]"), Some(0));
}

/// A file with no `token_type` array names no control tokens at all, so no
/// matcher is attached and tokenization is exactly what it was before.
#[test]
fn a_vocabulary_without_token_types_gets_no_specials() {
    let tok = from_gguf_vocab(llama_vocab()).expect("builds");
    assert_eq!(tok.special_token_id("<s>"), None);
    assert!(
        !tok.encode_raw("<s>hello").contains(&1),
        "nothing declared the token special, so it is ordinary text"
    );
}

// ── USER_DEFINED tokens (Gemma whitespace runs) ──────────────────────────────
//
// llama.cpp matches CONTROL *and* USER_DEFINED tokens as literal strings
// before merging even begins — neither ever reaches the merge loop. Gemma
// spells its whitespace-run pieces (`"  "`, `"   "`, ...) as USER_DEFINED, so
// selecting CONTROL alone leaves those runs unmatched: they fall through to
// the merge loop and come out as the single-space piece repeated instead of
// the trained multi-space token.

/// A USER_DEFINED-flagged multi-space piece must match the run verbatim, the
/// same way a CONTROL token does, and be resolvable by name.
#[test]
fn user_defined_whitespace_run_matches_as_one_token() {
    let tok = from_gguf_vocab(GgufVocab {
        tokens: v(&["<unk>", "<s>", "</s>", "▁", "  "]),
        token_type: Some(vec![3, 3, 3, 1, 4]),
        ..llama_vocab()
    })
    .expect("builds");

    assert_eq!(tok.family(), "Spm");
    assert_eq!(
        tok.encode_raw("  "),
        vec![4],
        "a USER_DEFINED whitespace run must match verbatim, not merge from repeated single-space pieces"
    );
    assert_eq!(tok.special_token_id("  "), Some(4));
}

/// Widening the selector to include USER_DEFINED must not have swapped it in
/// place of CONTROL: a CONTROL token still has to match and resolve by name.
#[test]
fn control_tokens_still_match_after_widening_to_user_defined() {
    let tok = from_gguf_vocab(GgufVocab {
        tokens: v(&["<unk>", "<s>", "</s>", "<start_of_turn>", "▁", "  "]),
        token_type: Some(vec![3, 3, 3, 3, 1, 4]),
        ..llama_vocab()
    })
    .expect("builds");

    assert_eq!(tok.encode_raw("<start_of_turn>"), vec![3]);
    assert_eq!(tok.special_token_id("<start_of_turn>"), Some(3));
}

// ── Precompiled charsmap ─────────────────────────────────────────────────────
//
// `tokenizer.ggml.precompiled_charsmap` is SentencePiece's normalization table.
// A Unigram vocabulary is built over its *normalized* pieces, so a character the
// table folds — tab, newline, NBSP, ZWJ, fullwidth punctuation — has no piece of
// its own and becomes `<unk>` when the table is not applied. The charsmap below
// is hand-built rather than borrowed from a real model: a real one is a quarter
// of a megabyte, and none of these tests may depend on ids the reference
// tokenizer produces.

/// A genuine, minimal SentencePiece charsmap carrying exactly one rule:
/// TAB → SPACE.
///
/// The format is a darts-clone double-array — `u32` trie size, that many bytes
/// of `u32` LE units, then the null-terminated replacement strings — and the
/// three non-zero units encode the single-byte walk the matcher performs:
///
/// | unit | meaning |
/// |---|---|
/// | `trie[0]` | root, offset 1, so traversal starts at node `0 ^ 1 = 1` |
/// | `trie[8]` | label `0x09` (TAB), has-leaf, offset 1: reached as `1 ^ 0x09`, leading to `8 ^ 1 = 9` |
/// | `trie[9]` | leaf, value 0 — offset 0 into the replacement table |
///
/// Every other unit is zero, so no other byte matches a label and the walk stops
/// after one step, which is the matcher's "no rule here, copy the character
/// verbatim" path.
fn tab_to_space_charsmap() -> Vec<u8> {
    let mut trie = [0u32; 16];
    trie[0] = 1 << 10;
    trie[8] = (1 << 10) | 0x100 | 0x09;
    // trie[9] stays 0: a leaf whose value is offset 0, where `" "` sits below.

    let mut blob = Vec::with_capacity(4 + trie.len() * 4 + 2);
    blob.extend_from_slice(&((trie.len() * 4) as u32).to_le_bytes());
    for unit in trie {
        blob.extend_from_slice(&unit.to_le_bytes());
    }
    // Replacement table: null-terminated strings, concatenated. Offset 0 is `" "`.
    blob.extend_from_slice(b" \0");
    blob
}

/// A Unigram vocabulary that distinguishes normalized from unnormalized input:
/// it has a piece for a space-separated `a` and `b`, and no piece for a tab.
fn t5_charsmap_vocab(charsmap: Option<Vec<u8>>) -> GgufVocab {
    GgufVocab {
        model: "t5".to_owned(),
        tokens: v(&["<unk>", "<s>", "</s>", "▁a", "▁b", "a", "b"]),
        scores: Some(vec![0.0, 0.0, 0.0, -1.0, -1.0, -5.0, -5.0]),
        bos_token_id: Some(1),
        eos_token_id: Some(2),
        precompiled_charsmap: charsmap,
        ..GgufVocab::default()
    }
}

/// The declared table is applied before segmentation, so a tab tokenizes as the
/// space it normalizes to.
#[test]
fn t5_applies_the_declared_charsmap() {
    let tok = from_gguf_vocab(t5_charsmap_vocab(Some(tab_to_space_charsmap()))).expect("builds");

    assert_eq!(tok.family(), "Unigram");
    assert_eq!(
        tok.encode_raw("a\tb"),
        tok.encode_raw("a b"),
        "the charsmap's TAB -> SPACE rule must run before pre-tokenization"
    );
    assert!(
        !tok.encode_raw("a\tb").contains(&0),
        "a normalized tab must not reach Viterbi as an uncovered character"
    );
}

/// The control: without the table the same input degrades to `<unk>`, which is
/// what makes the test above a test of the charsmap rather than of the vocab.
#[test]
fn t5_without_a_charsmap_leaves_the_tab_unnormalized() {
    let tok = from_gguf_vocab(t5_charsmap_vocab(None)).expect("builds");

    assert_ne!(tok.encode_raw("a\tb"), tok.encode_raw("a b"));
    assert!(
        tok.encode_raw("a\tb").contains(&0),
        "no vocabulary piece covers a raw tab, so it must fall back to <unk>"
    );
}

/// A blob that is not a usable table normalizes nothing, and must not refuse a
/// vocabulary that is otherwise complete — nor panic on the first input. Four
/// zero bytes are the awkward case: a well-formed header declaring an empty
/// trie, which has no root unit for the matcher to start from.
#[test]
fn an_unusable_charsmap_falls_back_to_no_normalization() {
    let tok = from_gguf_vocab(t5_charsmap_vocab(Some(vec![0, 0, 0, 0]))).expect("builds");
    let plain = from_gguf_vocab(t5_charsmap_vocab(None)).expect("builds");

    assert_eq!(tok.encode_raw("a\tb"), plain.encode_raw("a\tb"));
}

/// The charsmap is a Unigram rule. llama.cpp applies this table in its `ugm`
/// tokenizer only, so a `llama` (SentencePiece BPE) vocabulary that carries one
/// must tokenize exactly as it does without.
#[test]
fn the_charsmap_is_a_unigram_rule_only() {
    let with = from_gguf_vocab(GgufVocab {
        precompiled_charsmap: Some(tab_to_space_charsmap()),
        ..llama_vocab()
    })
    .expect("builds");
    let without = from_gguf_vocab(llama_vocab()).expect("builds");

    assert_eq!(
        with.encode_raw("hello\tworld"),
        without.encode_raw("hello\tworld")
    );
}