open-standards/aac/aac-rs/crates/aac-core/tests/langpacks.rs

//! The language packs (`aac-rs/data/lang`, which the app serves as `public/lang`), loaded the way
//! the app loads them — bytes in, model out — and prediction checked in
//! several scripts. Rebuild the packs with `tools/lang/build_packs.py`.

use std::collections::BTreeSet;
use std::path::PathBuf;
use std::sync::Arc;

use aac_core::i18n::supported_languages;
use aac_core::langpack::{PackError, PackIndex, PackQuality, Segmentation};
use aac_core::{
    CommunicationContext, HistoryEntry, LanguagePack, Letter, PredictionBasis, PredictionEngine,
    PredictionLanguageStatus, SuggestionScope,
};

/// Packs stay small enough to fetch on a phone connection.
const MAX_PACK_BYTES: usize = 300 * 1024;
/// And small enough in memory for a mid-range tablet running a wasm app.
const MAX_PACK_HEAP: usize = 2 * 1024 * 1024;

fn lang_dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../data/lang")
}

fn index() -> PackIndex {
    let json = std::fs::read_to_string(lang_dir().join("index.json"))
        .expect("data/lang/index.json — build it with tools/lang/build_packs.py");
    PackIndex::from_json(&json).expect("valid index")
}

fn raw(code: &str) -> Vec<u8> {
    let entry = index()
        .packs
        .get(code)
        .cloned()
        .unwrap_or_else(|| panic!("no pack for {code}"));
    std::fs::read(lang_dir().join(entry.file)).expect("pack file")
}

fn pack(code: &str) -> Arc<LanguagePack> {
    Arc::new(LanguagePack::from_bytes(&raw(code)).expect("pack loads"))
}

/// An engine proposing from the pack's whole word list, to test the pack.
fn engine(code: &str) -> PredictionEngine {
    let mut e = PredictionEngine::new();
    e.set_language(code, Some(pack(code)))
        .expect("pack matches language");
    e.set_suggestion_scope(SuggestionScope::WholeLanguage);
    e
}

fn said(engine: &mut PredictionEngine, text: &str) {
    engine.learn_from_utterance(&HistoryEntry {
        id: uuid::Uuid::nil(),
        text: text.to_string(),
        context: CommunicationContext::now(),
        energy_uj: 0,
        timestamp: chrono::DateTime::UNIX_EPOCH,
    });
}

fn texts(c: &[aac_core::PredictionCandidate]) -> Vec<&str> {
    c.iter().map(|c| c.text.as_str()).collect()
}

// ─── Loading ────────────────────────────────────────────────────────────────

#[test]
fn index_accounts_for_every_site_language() {
    let index = index();
    let site: BTreeSet<String> = supported_languages().into_iter().map(|l| l.code).collect();
    let packed: BTreeSet<String> = index.packs.keys().cloned().collect();
    let missing: BTreeSet<String> = index.missing.keys().cloned().collect();
    assert!(packed.is_disjoint(&missing));
    assert_eq!(
        &packed | &missing,
        site,
        "every site language has a pack or a stated reason"
    );
    assert_eq!(index.languages as usize, site.len());
    for (code, m) in &index.missing {
        assert!(!m.reason.is_empty(), "{code} has no reason");
    }
}

#[test]
fn every_pack_loads_round_trips_and_stays_small() {
    for (code, entry) in &index().packs {
        let bytes = std::fs::read(lang_dir().join(&entry.file)).unwrap();
        assert_eq!(bytes.len() as u64, entry.bytes, "{code}: size in index");
        assert!(
            bytes.len() <= MAX_PACK_BYTES,
            "{code}: {} bytes",
            bytes.len()
        );
        let p = LanguagePack::from_bytes(&bytes).unwrap_or_else(|e| panic!("{code}: {e}"));
        assert_eq!(p.code(), code);
        assert_eq!(p.quality(), entry.quality);
        assert_eq!(
            p.to_bytes(),
            bytes,
            "{code}: loader round trip is byte-identical"
        );
        let m = p.meta();
        assert_eq!(
            m.counts.words as usize,
            p.word_count(),
            "{code}: word count"
        );
        assert!(!m.attribution.is_empty() && !m.licence.is_empty() && !m.sources.is_empty());
        assert!(
            m.sources.iter().all(|s| s.url.contains("/tree/")),
            "{code}: sources pinned to a commit"
        );
        assert!(
            p.heap_bytes() <= MAX_PACK_HEAP,
            "{code}: {} bytes in memory",
            p.heap_bytes()
        );
        assert!(p.heap_bytes() <= 4 * bytes.len(), "{code}: memory estimate");
        // Words are sorted, so every word is found where it is.
        for i in [0, p.word_count() as u32 / 2, p.word_count() as u32 - 1] {
            assert_eq!(p.find(p.word(i)), Some(i), "{code}");
        }
    }
}

#[test]
fn damaged_packs_are_refused_whole() {
    let bytes = raw("en");
    let mut flipped = bytes.clone();
    let mid = flipped.len() / 2;
    flipped[mid] ^= 0x40;
    assert!(matches!(
        LanguagePack::from_bytes(&flipped),
        Err(PackError::Checksum { .. })
    ));
    assert!(LanguagePack::from_bytes(&bytes[..bytes.len() - 10]).is_err());
    assert!(LanguagePack::from_bytes(&bytes[..5]).is_err());
}

// ─── Language honesty (SPEC §4.1) ───────────────────────────────────────────

#[test]
fn status_says_what_predictions_rest_on() {
    let index = index();
    let en = PredictionLanguageStatus::from_index(&index, "en-US");
    assert_eq!(en.basis, PredictionBasis::Pack(PackQuality::MultiSource));
    assert!(
        en.attribution
            .as_deref()
            .is_some_and(|a| a.contains("wordfreq"))
    );

    // Swahili from one corpus, Zulu from a very small one.
    assert_eq!(
        PredictionLanguageStatus::from_index(&index, "sw").basis,
        PredictionBasis::Pack(PackQuality::SingleSource)
    );
    assert_eq!(
        PredictionLanguageStatus::from_index(&index, "zu").basis,
        PredictionBasis::Pack(PackQuality::Sparse)
    );

    // No pack: adaptive only, with the reason.
    for code in index.missing.keys() {
        let s = PredictionLanguageStatus::from_index(&index, code);
        assert_eq!(s.basis, PredictionBasis::AdaptiveOnly, "{code}");
        assert_eq!(s.notes.len(), 1, "{code}: reason carried");
    }

    // The engine reports the same once the pack is loaded, and learned words.
    let mut e = engine("es");
    said(&mut e, "quiero agua");
    let s = e.language_status();
    assert_eq!(s.code, "es");
    assert_eq!(s.basis, PredictionBasis::Pack(PackQuality::MultiSource));
    assert_eq!(s.learned_words, 2);
    assert_eq!(s.licence.as_deref(), Some("CC-BY-SA-4.0"));
}

#[test]
fn a_pack_for_another_language_is_refused() {
    let mut e = PredictionEngine::new();
    let err = e.set_language("es", Some(pack("en"))).unwrap_err();
    assert_eq!((err.expected.as_str(), err.found.as_str()), ("es", "en"));
    assert_eq!(e.language_status().basis, PredictionBasis::AdaptiveOnly);
}

// ─── Per-language smoke tests ───────────────────────────────────────────────

#[test]
fn english() {
    let e = engine("en");
    let done = e.complete_word("I want wat", 5);
    assert!(texts(&done).contains(&"water"), "{:?}", texts(&done));
    // A sentence opens with "I", written as English writes it.
    assert!(texts(&e.predict_next_word("", 10)).contains(&"I"));
    let next = e.predict_next_word("I want to", 8);
    assert_eq!(next.len(), 8);
    // Letters after "wat": 'e' (water) and 'c' (watch) lead.
    let letters = e.predict_letters("wat", 3);
    assert!(
        letters.iter().any(|l| l.letter == Letter::Char('e')),
        "{letters:?}"
    );
    // A prefix no word starts with is answered by the letter model alone.
    assert!(!e.predict_letters("xq", 3).is_empty());
    // The pinned/recent path and the word path share one list.
    let ctx = CommunicationContext::now();
    assert!(!e.predict("I want wat", &ctx, 8).is_empty());
}

#[test]
fn spanish() {
    let e = engine("es");
    let done = e.complete_word("yo quier", 3);
    assert_eq!(done[0].text, "quiero", "{:?}", texts(&done));
    let letters = e.predict_letters("quier", 2);
    assert!(
        letters.iter().any(|l| l.letter == Letter::Char('o')),
        "{letters:?}"
    );
}

#[test]
fn french_elision() {
    let e = engine("fr");
    assert!(e.language_pack().unwrap().meta().normalization.elision);
    // The particle keeps its apostrophe, and a particle just typed is finished.
    assert!(texts(&e.complete_word("j", 10)).contains(&"j'"));
    assert!(!e.predict_next_word("je veux l'", 5).is_empty());
}

#[test]
fn arabic_right_to_left_with_vowel_marks() {
    let e = engine("ar");
    let p = e.language_pack().unwrap();
    assert!(p.meta().normalization.remove_marks);
    // Typed with and without harakat, the prefix meets the same words.
    let plain = e.complete_word("شك", 5);
    let marked = e.complete_word("شُك", 5);
    assert_eq!(texts(&plain), texts(&marked));
    assert!(texts(&plain).contains(&"شكرا"), "{:?}", texts(&plain));
    // Every suggestion is Arabic script (logical order; the UI lays it out RTL).
    for c in &e.predict_next_word("أنا", 8) {
        assert!(
            c.text
                .chars()
                .all(|ch| ('\u{0600}'..='\u{06FF}').contains(&ch)),
            "{}",
            c.text
        );
    }
}

#[test]
fn hindi_devanagari() {
    let e = engine("hi");
    let done = e.complete_word("मुझे नही", 5);
    assert!(texts(&done).contains(&"नहीं"), "{:?}", texts(&done));
    // After क + virama the next letter is a consonant of a conjunct: य (क्या) leads.
    let letters = e.predict_letters("क्", 3);
    assert!(
        letters.iter().any(|l| l.letter == Letter::Char('य')),
        "{letters:?}"
    );
}

/// Chinese is written without spaces, so the engine finds words by longest
/// dictionary match against the pack (an approximation — the pack and the
/// status both say so). Typed text like 我想喝 is read as 我 / 想 / 喝, and
/// because 喝 is itself a word the engine predicts what follows it rather than
/// completing it: it cannot know whether 喝 is finished.
#[test]
fn chinese_dictionary_segmentation_caveat() {
    let e = engine("zh");
    let p = e.language_pack().unwrap();
    assert_eq!(p.meta().segmentation, Segmentation::Dictionary);
    assert!(p.meta().notes.iter().any(|n| n.contains("approximate")));
    let next = e.predict_next_word("我想", 8);
    assert!(!next.is_empty());
    assert!(next.iter().all(|c| {
        c.text
            .chars()
            .all(|ch| ('\u{4E00}'..='\u{9FFF}').contains(&ch))
    }));
    assert!(
        e.complete_word("我想喝", 5).is_empty(),
        "a known word is read as finished"
    );
    // Learning uses the same segmentation: said once, 喝水 follows 想.
    let mut e = engine("zh");
    said(&mut e, "我想喝水");
    said(&mut e, "我想喝水");
    assert_eq!(e.predict_next_word("我想", 3)[0].text, "喝水");
}

// ─── What may be proposed ───────────────────────────────────────────────────

/// Words of an English board, as a client passes them.
const BOARD: &[&str] = &[
    "I", "you", "want", "more", "help", "go", "stop", "like", "not", "yes", "no", "it", "that",
    "here", "all done", "eat", "drink", "play", "fun", "funny", "kind", "Mom",
];

fn board_words() -> BTreeSet<String> {
    BOARD.iter().flat_map(|w| w.split(' ')).map(str::to_lowercase).collect()
}

fn all_allowed(c: &[aac_core::PredictionCandidate], allowed: &BTreeSet<String>) -> bool {
    c.iter().all(|c| allowed.contains(&c.text.to_lowercase()))
}

#[test]
fn by_default_only_the_boards_and_the_persons_words_are_proposed() {
    let mut e = PredictionEngine::new();
    e.set_language("en", Some(pack("en"))).unwrap();
    assert_eq!(e.suggestion_scope(), SuggestionScope::Vocabulary);
    // No vocabulary given and nothing said: the pack alone proposes nothing,
    // but still predicts letters (a letter proposes no word).
    assert!(e.predict_next_word("I want", 8).is_empty());
    assert!(e.complete_word("fu", 8).is_empty());
    assert!(!e.predict_letters("wat", 3).is_empty());

    e.set_vocabulary(BOARD.iter().copied());
    let mut allowed = board_words();
    for text in ["", "I ", "I want", "you", "go "] {
        let next = e.predict_next_word(text, 8);
        assert!(!next.is_empty(), "{text:?}: the board's words are proposed");
        assert!(all_allowed(&next, &allowed), "{text:?}: {:?}", texts(&next));
    }
    for prefix in ["f", "fu", "b", "k", "ki", "s", "m"] {
        let done = e.complete_word(prefix, 8);
        assert!(all_allowed(&done, &allowed), "{prefix:?}: {:?}", texts(&done));
    }
    // Typed "fu", the whole list's second word is a swear word; "ki" offers
    // "killed". The board's words come instead, and the pack ranks them.
    assert_eq!(texts(&e.complete_word("fu", 8)), ["fun", "funny"]);
    assert_eq!(texts(&e.complete_word("ki", 8)), ["kind"]);
    let after_i = texts(&e.predict_next_word("I", 8)).into_iter().map(String::from).collect::<Vec<_>>();
    let mut alphabetical = after_i.clone();
    alphabetical.sort();
    assert_ne!(after_i, alphabetical, "ranked by the pack, not listed");
    // The pack's display form wins over a label's capital; "I" stays "I".
    assert!(texts(&e.complete_word("m", 8)).contains(&"mom"));
    assert!(texts(&e.predict_next_word("", 8)).contains(&"I"));

    // Nothing stops the person: a word they say is theirs, and is offered.
    said(&mut e, "the kitten is asleep");
    allowed.extend(["the", "kitten", "is", "asleep"].map(String::from));
    assert!(texts(&e.complete_word("ki", 8)).contains(&"kitten"));
    assert!(all_allowed(&e.complete_word("ki", 8), &allowed));
}

#[test]
fn the_whole_word_list_is_a_choice() {
    let mut e = PredictionEngine::new();
    e.set_language("en", Some(pack("en"))).unwrap();
    e.set_vocabulary(BOARD.iter().copied());
    let allowed = board_words();
    assert!(all_allowed(&e.complete_word("ki", 8), &allowed));
    e.set_suggestion_scope(SuggestionScope::WholeLanguage);
    let whole = e.complete_word("ki", 8);
    assert!(!all_allowed(&whole, &allowed), "the whole list adds words: {:?}", texts(&whole));
    assert!(texts(&whole).contains(&"kind"), "the board's words still come");
    e.set_suggestion_scope(SuggestionScope::Vocabulary);
    assert!(all_allowed(&e.complete_word("ki", 8), &allowed));
}

#[test]
fn hindi_news_words_are_not_proposed_by_default() {
    // Hindi sentences are mostly news: after की ("of"), the whole list offers
    // मौत ("death") and हत्या ("murder").
    let mut e = PredictionEngine::new();
    e.set_language("hi", Some(pack("hi"))).unwrap();
    e.set_vocabulary(["मैं", "तुम", "चाहिए", "और", "मदद", "खाना", "पानी", "की", "है", "बात"]);
    let next = e.predict_next_word("की", 8);
    assert!(!next.is_empty());
    for c in &next {
        assert!(!["मौत", "हत्या"].contains(&c.text.as_str()), "{:?}", texts(&next));
    }
    let done = e.complete_word("म", 8);
    assert!(done.iter().all(|c| ["मैं", "मदद"].contains(&c.text.as_str())), "{:?}", texts(&done));
}

#[test]
fn a_vocabulary_is_split_and_keyed_as_the_pack_keys_words() {
    // Chinese has no spaces: a phrase is split with the pack's dictionary.
    let mut e = PredictionEngine::new();
    e.set_language("zh", Some(pack("zh"))).unwrap();
    e.set_vocabulary(["我想喝水", "帮助"]);
    let next = texts(&e.predict_next_word("我想", 8)).into_iter().map(String::from).collect::<Vec<_>>();
    assert!(!next.is_empty());
    assert!(next.iter().all(|w| "我想喝水帮助".contains(w.as_str())), "{next:?}");
}

// ─── Where a completion goes ────────────────────────────────────────────────

/// The text left once the chosen completion replaces the word being typed.
fn completed(e: &PredictionEngine, text: &str, word: &str) -> String {
    match e.completion_start(text) {
        Some(at) => format!("{}{word}", &text[..at]),
        None => format!("{text}{word}"),
    }
}

#[test]
fn a_completion_replaces_exactly_the_letters_being_typed() {
    let en = engine("en");
    assert_eq!(completed(&en, "I want fu", "fun"), "I want fun");
    assert_eq!(completed(&en, "(wa", "want"), "(want");
    assert_eq!(completed(&en, "don\u{2019}", "don't"), "don't");
    assert_eq!(en.completion_start("I want "), None, "a finished word: the next one follows");
    assert_eq!(en.completion_start("I want."), None);
    assert!(en.writes_spaces());

    let es = engine("es");
    assert_eq!(completed(&es, "¿Qu", "qué"), "¿qué");
    assert_eq!(completed(&es, "¡ho", "hola"), "¡hola");

    let fr = engine("fr");
    assert_eq!(completed(&fr, "je veux l'e", "eau"), "je veux l'eau");
    assert_eq!(completed(&fr, "je veux l\u{2019}e", "eau"), "je veux l\u{2019}eau");
    assert_eq!(fr.completion_start("je veux l'"), None, "a particle just typed is finished");

    // Arabic typed with vowel marks: the marks are part of the word typed.
    let ar = engine("ar");
    assert_eq!(completed(&ar, "أنا أُرِ", "أريد"), "أنا أريد");

    // Chinese: no spaces, and the completion replaces only the open segment.
    let zh = engine("zh");
    assert!(!zh.writes_spaces());
    let start = zh.completion_start("我想喝").map(|at| &"我想喝"[at..]);
    assert!(matches!(start, None | Some("喝")), "{start:?}");
}

// ─── Blending, learning, determinism ───────────────────────────────────────

#[test]
fn learns_only_from_what_was_said_and_takes_over() {
    let mut e = engine("en");
    let before = e.predict_next_word("I want my", 5);
    let bytes = e.memory_bytes();
    // Asking many times teaches nothing.
    for _ in 0..50 {
        let _ = e.predict_next_word("I want my tablet", 5);
        let _ = e.predict_letters("I want my tab", 5);
    }
    assert_eq!(e.memory_bytes(), bytes);
    assert_eq!(texts(&e.predict_next_word("I want my", 5)), texts(&before));

    // Saying it does: this person's next word overtakes the seed.
    for _ in 0..3 {
        said(&mut e, "I want my blue tablet");
    }
    assert_eq!(e.predict_next_word("I want my", 5)[0].text, "blue");
    assert_eq!(e.complete_word("I want my bl", 1)[0].text, "blue");
    assert_eq!(
        e.predict_letters("I want my bl", 1)[0].letter,
        Letter::Char('u')
    );
}

#[test]
fn ranking_is_deterministic() {
    let a = engine("de");
    let b = engine("de");
    for text in ["ich möchte", "ich möchte ", "gro", "", "wie geht es "] {
        let pa = a.predict_next_word(text, 10);
        let pb = b.predict_next_word(text, 10);
        let bits = |p: &[aac_core::PredictionCandidate]| {
            p.iter()
                .map(|c| (c.text.clone(), c.score.to_bits()))
                .collect::<Vec<_>>()
        };
        assert_eq!(bits(&pa), bits(&pb), "{text:?}");
        assert_eq!(
            bits(&pa),
            bits(&a.predict_next_word(text, 10)),
            "{text:?} repeated"
        );
        let la: Vec<_> = a
            .predict_letters(text, 10)
            .into_iter()
            .map(|l| (l.letter, l.score.to_bits()))
            .collect();
        let lb: Vec<_> = b
            .predict_letters(text, 10)
            .into_iter()
            .map(|l| (l.letter, l.score.to_bits()))
            .collect();
        assert_eq!(la, lb);
    }
    // German keeps ß (wordfreq folds it to ss; the builder restores it).
    assert!(texts(&a.complete_word("gro", 10)).contains(&"groß"));
}

#[test]
fn engine_memory_is_bounded() {
    let mut e = engine("zh"); // the largest pack
    for i in 0..200 {
        said(&mut e, &format!("我想喝水 {i}"));
    }
    assert!(e.memory_bytes() < 3 * 1024 * 1024, "{}", e.memory_bytes());
}