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

//! A grain-free AV1 base refined by `wai.video.int_refine` (`extensions/platform-decode.md`
//! §6): the corpus `platform-decode-conformance/grain-free-av1/`, replayed through
//! `staged::StageSink` from its committed frames.
//!
//! Without a decoder (feature `staged`): every probe digest is the SHA-256 of its frame of
//! `base.i420`; the probe's codec string and size are the stream's; the stage-0 hashes are
//! the frames' `RGB8` units; every case gives its stated reports, units and clip hash; the
//! stream with film grain is refused as a base, and so is the grain-free stream with any
//! sequence header's grain flag set; and frames a platform altered by one sample give
//! other stage hashes. With feature `av1_base`, the reference decode through libdav1d gives
//! `base.i420` byte for byte, and a decode of the stream with film grain is refused before
//! it runs.

use std::path::PathBuf;

use serde_json::Value;
use wai::container::Wai;
use wai::staged::av1::{self, Av1Decoder};
use wai::staged::{tier, ChainConfig, HoldBound, ReceiptEvidence, StageSink};

fn dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../platform-decode-conformance/grain-free-av1")
}

fn read(f: &str) -> Vec<u8> {
    std::fs::read(dir().join(f)).unwrap_or_else(|e| panic!("{f}: {e}"))
}

fn json(f: &str) -> Value {
    serde_json::from_slice(&read(f)).unwrap_or_else(|e| panic!("{f}: {e}"))
}

fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{x:02x}")).collect()
}

fn decoder(t: &str) -> Av1Decoder {
    match t {
        "reference" => Av1Decoder::Reference,
        "probe-matched" => Av1Decoder::ProbeMatched,
        "unprobed" => Av1Decoder::Unprobed,
        other => panic!("decoder {other}"),
    }
}

const FILES: [&str; 5] = ["base.wai", "base-grain.wai", "layer1.wai", "layer2.wai", "layer1-over-grain.wai"];

fn name_of(h: &[u8; 32]) -> String {
    FILES.iter().find(|f| wai::jwp_hash::content_hash(&read(f)) == *h).map(|f| f.to_string()).unwrap_or_else(|| hex(h))
}

#[test]
fn the_probe_is_the_reference_frames() {
    let p = json("probe.json");
    let frames = read("base.i420");
    let base = av1::read_base(&Wai::from_bytes(&read("base.wai")).unwrap().payload).unwrap();
    assert_eq!(p["codec"], base.codec.as_str());
    assert_eq!((p["codedWidth"].as_u64(), p["codedHeight"].as_u64(), p["frames"].as_u64()), (Some(base.width.into()), Some(base.height.into()), Some(base.frames.into())));
    assert_eq!(p["film_grain_params_present"], false);
    let fl = base.frame_len();
    assert_eq!(frames.len(), fl * base.frames as usize);
    let digests: Vec<String> = frames.chunks_exact(fl).map(|f| hex(&av1::frame_sha256(f))).collect();
    assert_eq!(p["frame_sha256"], serde_json::json!(digests));
    let cat: Vec<u8> = frames.chunks_exact(fl).flat_map(av1::frame_sha256).collect();
    assert_eq!(p["clip_sha256"], hex(&wai::pin::sha256(&cat)).as_str());
    let c = json("cases.json");
    let stage0: Vec<String> = frames.chunks_exact(fl).map(|f| hex(&wai::pin::sha256(&av1::i420_to_rgb8(f, base.width, base.height)))).collect();
    assert_eq!(c["stage0_sha256"], serde_json::json!(stage0));
}

#[test]
fn every_case_replays() {
    let c = json("cases.json");
    let frames = read("base.i420");
    let cases = c["cases"].as_array().unwrap();
    assert_eq!(cases.len(), 7);
    for case in cases {
        let name = case["name"].as_str().unwrap();
        let e = &case["expect"];
        let mut sink = StageSink::new(ChainConfig { bound: HoldBound { max_envelopes: 16, max_age_ms: 1000 }, cost_limit: None, budget: 1 << 30, record_budget: 1 << 20, live: None });
        let env = read(case["base"].as_str().unwrap());
        let hash = wai::jwp_hash::content_hash(&env);
        match av1::base_from_i420(&env, &frames, decoder(case["decoder"].as_str().unwrap()), None) {
            Ok((info, units)) => {
                assert!(e["base_refusal"].is_null(), "{name}");
                assert!(sink.add_base(info, units).is_empty());
            }
            Err(r) => assert_eq!(e["base_refusal"], r.code(), "{name}"),
        }
        for (i, (s, want)) in case["steps"].as_array().unwrap().iter().zip(e["steps"].as_array().unwrap()).enumerate() {
            let s = s.as_str().unwrap();
            let rs = match s.strip_prefix("expire@") {
                Some(t) => sink.expire(t.parse().unwrap()),
                None => sink.offer(&read(s), 1 + i as u64, ReceiptEvidence::None).unwrap(),
            };
            let got: Vec<Value> = rs
                .iter()
                .map(|r| {
                    serde_json::json!({
                        "envelope": name_of(&r.content_hash),
                        "reason": r.envelope.map(|x| x.code()),
                        "held_for_base": r.held_for_base,
                        "applied": r.applied,
                        "held": r.held,
                        "refused": r.refused.iter().map(|(u, w)| serde_json::json!([u, w.code()])).collect::<Vec<_>>(),
                    })
                })
                .collect();
            assert_eq!(serde_json::json!(got), want["reports"], "{name} step {i}");
        }
        match sink.chain(&hash) {
            Some(ch) => {
                let units: Vec<Value> = (0..ch.n_units())
                    .map(|u| {
                        let s = ch.unit(u).unwrap();
                        serde_json::json!({ "stage": s.stage, "producer": name_of(&s.producer), "tier": tier::token(ch.tier(u)), "sha256": ch.unit_sha256(u).map(|h| hex(&h)) })
                    })
                    .collect();
                assert_eq!(serde_json::json!(units), e["units"], "{name}");
                assert_eq!(serde_json::json!(ch.clip_sha256().map(|h| hex(&h))), e["clip_sha256"], "{name}");
            }
            None => assert!(e["units"].is_null() && e["clip_sha256"].is_null(), "{name}"),
        }
    }
    // The two decoders at decode-equivalence reach the same stage hashes, and two layers
    // reach them in either order.
    let by = |n: &str| cases.iter().find(|c| c["name"] == n).unwrap()["expect"]["clip_sha256"].clone();
    assert!(by("two-layers").is_string());
    assert_eq!(by("two-layers"), by("two-layers-reference"));
    assert_eq!(by("two-layers"), by("two-layers-out-of-order"));
    assert_ne!(by("two-layers"), by("one-layer"));
}

#[test]
fn grain_anywhere_refuses_the_base() {
    let g = Wai::from_bytes(&read("base-grain.wai")).unwrap();
    assert_eq!(av1::read_base(&g.payload).unwrap_err().code, "av1_film_grain");
    // The grain-free stream with the flag set in its first sequence header: the one bit
    // whose flip makes the header read with film grain and nothing else changed.
    let w = Wai::from_bytes(&read("base.wai")).unwrap();
    let s = wai::platform_decode::av1_stream(&w.payload).unwrap();
    assert!(s.grain_free() && !s.sequences.is_empty());
    let tu = s.units[0];
    // The first unit starts after the framing's 20-byte header and its 4-byte length.
    let start = 24;
    let mut o = 0usize;
    let (body_at, body_len) = loop {
        let header = tu[o];
        let ext = usize::from(header & 0x04 != 0);
        let (mut size, mut k) = (0usize, 0usize);
        loop {
            let b = tu[o + 1 + ext + k];
            size |= usize::from(b & 0x7f) << (7 * k);
            k += 1;
            if b & 0x80 == 0 {
                break;
            }
        }
        let at = o + 1 + ext + k;
        if (header >> 3) & 0x0f == 1 {
            break (at, size);
        }
        o = at + size;
    };
    let mut found = false;
    for bit in 0..body_len * 8 {
        let mut p = w.payload.clone();
        p[start + body_at + bit / 8] ^= 0x80 >> (bit % 8);
        let seq = wai::platform_decode::Av1Sequence::parse(&p[start + body_at..start + body_at + body_len]);
        if let Ok(q) = seq {
            let orig = &s.sequences[0];
            if q.film_grain_params_present && q.codec_string() == orig.codec_string() && q.color_range == orig.color_range && q.chroma_sample_position == orig.chroma_sample_position {
                assert_eq!(av1::read_base(&p).unwrap_err().code, "av1_film_grain", "bit {bit}");
                found = true;
                break;
            }
        }
    }
    assert!(found, "a bit flip that sets only the grain flag");
}

#[test]
fn an_altered_sample_changes_the_stage_hashes() {
    let env = read("base.wai");
    let mut frames = read("base.i420");
    let (info, units) = av1::base_from_i420(&env, &frames, Av1Decoder::Reference, None).unwrap();
    let mut a = StageSink::new(ChainConfig { bound: HoldBound { max_envelopes: 4, max_age_ms: 10 }, cost_limit: None, budget: 1 << 30, record_budget: 1 << 20, live: None });
    a.add_base(info.clone(), units);
    frames[1234] ^= 1;
    let (_, units) = av1::base_from_i420(&env, &frames, Av1Decoder::Reference, None).unwrap();
    let mut b = StageSink::new(ChainConfig { bound: HoldBound { max_envelopes: 4, max_age_ms: 10 }, cost_limit: None, budget: 1 << 30, record_budget: 1 << 20, live: None });
    b.add_base(info.clone(), units);
    for l in ["layer1.wai", "layer2.wai"] {
        a.offer(&read(l), 1, ReceiptEvidence::None).unwrap();
        b.offer(&read(l), 1, ReceiptEvidence::None).unwrap();
    }
    let (ca, cb) = (a.chain(&info.content_hash).unwrap(), b.chain(&info.content_hash).unwrap());
    assert_eq!(ca.unit(0).unwrap().stage, 2);
    assert_ne!(ca.unit_sha256(0), cb.unit_sha256(0), "the layers refine what the decoder gave");
    assert_eq!(ca.unit_sha256(1), cb.unit_sha256(1));
}

#[cfg(feature = "av1_base")]
#[test]
fn the_reference_decode_gives_the_committed_frames() {
    assert_eq!(av1::decode_i420(&read("base.wai")).unwrap(), read("base.i420"), "libdav1d's frames are base.i420");
    let e = av1::decode_i420(&read("base-grain.wai")).unwrap_err();
    assert_eq!(e.code(), "av1_film_grain");
    let (info, units) = av1::decode_base(&read("base.wai"), None).unwrap();
    assert_eq!(info.tier, Some(wai::codecs::DeterminismTier::DecodeEquivalence));
    let c = json("cases.json");
    let got: Vec<String> = units.iter().map(|u| hex(&wai::pin::sha256(&u.samples.canonical_bytes()))).collect();
    assert_eq!(c["stage0_sha256"], serde_json::json!(got));
}