open-standards/wai/wai-rs/tests/dred_exact.rs

//! `live-conformance/dred/` (extensions/opus-dred.md §9) against the reference module: every file
//! regenerates byte for byte from its fixed seeds and the pinned weights, no other file is there
//! but README.md, and every payload vector and case reaches the outcome it states. And the
//! extension's text against the module: §8 lists exactly the codes the module reports, and §4.2
//! and §4.5 state the digests the module computes.
//!
//! The weights file is not in the repository: `WAI_DRED_WEIGHTS` names it
//! (`live-conformance/dred/fetch.sh` fetches it and checks its digest). Without it the corpus
//! cannot be checked, and the test fails rather than pass without checking it.

use std::path::PathBuf;

use wai::dred_exact::{corpus, hex32, lower, payload, weights, Reason};

fn wai() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..")
}

fn extension() -> String {
    std::fs::read_to_string(wai().join("extensions/opus-dred.md")).unwrap()
}

fn weights() -> Vec<u8> {
    let p = std::env::var("WAI_DRED_WEIGHTS").expect(
        "WAI_DRED_WEIGHTS must name the draft's weights file: bash wai/live-conformance/dred/fetch.sh <dir>, then WAI_DRED_WEIGHTS=<dir>/dred_unquantized_weights.bin",
    );
    std::fs::read(&p).unwrap_or_else(|e| panic!("{p}: {e}"))
}

#[test]
fn the_corpus_regenerates_and_every_vector_reaches_its_outcome() {
    corpus::verify(&weights(), &wai().join("live-conformance/dred")).unwrap();
}

/// §8's table lists exactly the codes the module reports, in its order.
#[test]
fn the_refusals_table_is_the_modules() {
    let text = extension();
    let sec = text.split("\n## 8.").nth(1).expect("§8").split("\n## ").next().unwrap();
    let listed: Vec<&str> = sec.lines().filter_map(|l| l.strip_prefix("| `")).filter_map(|l| l.split('`').next()).collect();
    let want: Vec<&str> = Reason::LISTED.iter().map(|r| r.code()).collect();
    assert_eq!(listed, want);
}

/// The digests the extension states are the module's: the draft's tables (§4.2), the tables'
/// constant and the two activation tables (§4.5).
#[test]
fn the_extension_states_the_modules_digests() {
    let text = extension();
    let tb = |t: &[i16]| {
        use sha2::{Digest, Sha256};
        let d: [u8; 32] = Sha256::digest(t.iter().flat_map(|x| x.to_le_bytes()).collect::<Vec<u8>>()).into();
        hex32(&d)
    };
    for want in [hex32(&payload::tables_digest()), tb(&lower::tanh_table()), tb(&lower::sigmoid_table()), lower::EXP_C0.to_string()] {
        assert!(text.contains(&want), "extensions/opus-dred.md states {want}");
    }
}

/// §4.8: both graphs of the pinned set load under `acc64` and are refused under `acc32`, and
/// every operator they use is in `wai-ops/1-i32`.
#[test]
fn the_graphs_meet_acc64_and_not_acc32() {
    let g = lower::lower(&weights::pinned(&weights()).unwrap()).unwrap();
    for (name, graph) in [("init", &g.init), ("step", &g.step)] {
        let b = graph.to_bytes();
        assert!(wai_ops::graph::load(&b, wai_ops::Profile::Acc64).is_ok(), "{name} loads under acc64");
        assert!(wai_ops::graph::load(&b, wai_ops::Profile::Acc32).is_err(), "{name} is refused under acc32");
        use wai_ops::graph::opcode::*;
        assert!(graph.ops.iter().all(|o| [CONV, LUT16, SPLIT, CONCAT, MUL, ADD, RESCALE].contains(&o.code)), "{name}: wai-ops/1-i32 operators");
    }
}