rivet-transcoder 0.1.2

//! Squad-35: Structured-pattern round-trip fidelity test.
//!
//! Closes the "grey screen" coverage gap. Every other test in this crate
//! verifies byte-equality, structural conformance, or codec-config survival —
//! none would catch a bug that always wrote a constant grey frame because
//! "all 128 luma" is a perfectly spec-conformant AV1 output.
//!
//! How this test proves the output actually depicts the input:
//!   1. Build N synthetic Yuv420p frames where frame index `i` encodes its
//!      own number as a row of bright/dark luma blocks in the upper-left
//!      corner (binary, 8 bits = covers up to frame 255).
//!   2. Run them through the rav1e encoder + AV1 MP4 muxer.
//!   3. Demux + rav1d-decode the output.
//!   4. Sample the same pixel locations in every decoded frame, threshold
//!      against mid-grey to recover bits, and assert the recovered index
//!      sequence equals `0..N` exactly.
//!
//! A grey-screen replacement bug fails immediately: every recovered bit
//! reads `0` and the recovered indices are all 0 instead of monotonically
//! increasing.
//!
//! Tolerance: rav1e at low bitrate is lossy. We don't compare pixels exactly
//! — we use 32×32 blocks (large enough that DCT can't average them away)
//! placed on row=8 in known x positions, and threshold the block's mean luma
//! against 128 (mid-grey). With block luma values 32 (dark) vs 224 (bright)
//! we get >80 sample-domain margin even after AV1 quantization at quality=200.

use bytes::Bytes;

mod common;

use codec::decode::Decoder;
use codec::encode::EncoderConfig;
use codec::frame::{ColorSpace, PixelFormat, StreamInfo, VideoFrame};
use container::demux;

const W: u32 = 320;
const H: u32 = 240;
const FPS: f64 = 30.0;
const FRAMES: u32 = 24; // > rav1e min_key_frame_interval (12) so multiple GOPs

// Pattern geometry — tuned so the encoder cannot collapse the bits into
// quantization noise.
const BIT_BLOCK_SIZE: u32 = 16; // 16×16 px per bit
const BIT_ROW_Y: u32 = 16; // bit blocks live on row y=16..32
const BIT_X0: u32 = 16; // first bit block at x=16
const NUM_BITS: u32 = 8; // 8-bit frame index (0..255)
const DARK_LUMA: u8 = 24; // bit=0
const BRIGHT_LUMA: u8 = 232; // bit=1

/// Build a Yuv420p frame whose top-left corner encodes `index` as 8 bits
/// of bright/dark blocks. Background is a deterministic checkerboard so
/// the encoder doesn't see large flat regions (which it would aggressively
/// quantize).
fn make_indexed_frame(width: u32, height: u32, index: u32) -> VideoFrame {
    let w = width as usize;
    let h = height as usize;
    let y_size = w * h;
    let uv_size = y_size / 4;
    let mut y = vec![0u8; y_size];
    let mut u = vec![128u8; uv_size];
    let mut v = vec![128u8; uv_size];

    // Background: 16×16 checkerboard with luma = 96 / 160. Provides texture
    // so the encoder allocates bits broadly instead of dumping them on the
    // bit row.
    for row in 0..h {
        for col in 0..w {
            let cell = ((col / 16) + (row / 16)) & 1;
            y[row * w + col] = if cell == 0 { 96 } else { 160 };
        }
    }

    // Bit row: 8 blocks of BIT_BLOCK_SIZE × BIT_BLOCK_SIZE.
    // Bit i (MSB first) controls block i. Drawn over the checkerboard.
    for bit in 0..NUM_BITS {
        let on = (index >> (NUM_BITS - 1 - bit)) & 1 == 1;
        let luma = if on { BRIGHT_LUMA } else { DARK_LUMA };
        let x0 = (BIT_X0 + bit * (BIT_BLOCK_SIZE + 4)) as usize;
        let y0 = BIT_ROW_Y as usize;
        for dy in 0..BIT_BLOCK_SIZE as usize {
            for dx in 0..BIT_BLOCK_SIZE as usize {
                let row = y0 + dy;
                let col = x0 + dx;
                if row < h && col < w {
                    y[row * w + col] = luma;
                }
            }
        }
    }

    // Mild per-frame variance in chroma so consecutive frames aren't
    // bit-identical to the encoder (which would prefer a single keyframe
    // and skip thereafter).
    let chroma_jitter = (index as u8).wrapping_mul(3);
    for c in u.iter_mut() {
        *c = 128u8.wrapping_add(chroma_jitter / 4);
    }
    for c in v.iter_mut() {
        *c = 128u8.wrapping_sub(chroma_jitter / 4);
    }

    let mut buf = Vec::with_capacity(y_size + 2 * uv_size);
    buf.extend_from_slice(&y);
    buf.extend_from_slice(&u);
    buf.extend_from_slice(&v);

    VideoFrame::new(
        Bytes::from(buf),
        width,
        height,
        PixelFormat::Yuv420p,
        ColorSpace::Bt709,
        index as u64,
    )
}

/// Recover the encoded frame index by sampling the bit-row blocks. Returns
/// `(index, confidence_bits)` where `confidence_bits` is the count of
/// blocks whose mean luma was unambiguously above or below mid-grey by
/// a comfortable margin (≥ 32 luma units away from 128). We require
/// at least 80% of bits (≥ 7/8) to land in the "confident" zone, which is
/// the contract for the per-frame fidelity test.
fn recover_index(frame: &VideoFrame) -> (u32, u32) {
    let w = frame.width as usize;
    let h = frame.height as usize;
    assert!(matches!(frame.format, PixelFormat::Yuv420p));
    assert_eq!(frame.data.len(), w * h * 3 / 2);
    let y_plane = &frame.data[..w * h];

    let mut index = 0u32;
    let mut confident = 0u32;

    for bit in 0..NUM_BITS {
        let x0 = (BIT_X0 + bit * (BIT_BLOCK_SIZE + 4)) as usize;
        let y0 = BIT_ROW_Y as usize;
        // Average luma over the inner 8×8 of the block (avoid block edges
        // where AV1's deblocking filter blurs into neighbours).
        let inset = 4;
        let mut sum = 0u32;
        let mut count = 0u32;
        for dy in inset..(BIT_BLOCK_SIZE as usize - inset) {
            for dx in inset..(BIT_BLOCK_SIZE as usize - inset) {
                let row = y0 + dy;
                let col = x0 + dx;
                if row < h && col < w {
                    sum += y_plane[row * w + col] as u32;
                    count += 1;
                }
            }
        }
        let mean = (sum / count.max(1)) as i32;
        let bit_val = if mean > 128 { 1 } else { 0 };
        index |= (bit_val as u32) << (NUM_BITS - 1 - bit);
        if (mean - 128).abs() >= 32 {
            confident += 1;
        }
    }
    (index, confident)
}

#[test]
fn structured_pattern_round_trip_recovers_frame_indices() {
    // Encoder config — quality=80 (production default), speed_preset=10
    // (fastest), keyframe_interval=30 (matches e2e.rs's known-good pattern).
    // With 24 input frames + min_key_frame_interval=12 (rav1e's lower
    // bound), we get 1 keyframe at frame 0 + possibly 1 more at frame 12,
    // plus 22 P-frames. Each frame produces its own packet so all 24
    // bit-patterns are recoverable.
    let config = EncoderConfig {
        width: W,
        height: H,
        frame_rate: FPS,
        quality: 80,
        speed_preset: 10,
        keyframe_interval: 30,
        ..EncoderConfig::default()
    };
    let Some(mut encoder) = common::try_av1_encoder(config) else {
        return;
    };
    let mut muxer = container::mux::Av1Mp4Muxer::new(W, H, FPS).expect("muxer");

    let mut packet_count = 0usize;
    for i in 0..FRAMES {
        let f = make_indexed_frame(W, H, i);
        encoder.send_frame(&f).expect("send_frame");
        while let Some(p) = encoder.receive_packet().expect("receive") {
            packet_count += 1;
            muxer.add_packet(p).expect("add_packet");
        }
    }
    encoder.flush().expect("flush");
    while let Some(p) = encoder.receive_packet().expect("receive after flush") {
        packet_count += 1;
        muxer.add_packet(p).expect("add_packet");
    }
    assert!(packet_count > 0, "encoder produced zero packets");
    let mp4 = muxer.finalize().expect("mux finalize");
    assert!(!mp4.is_empty(), "muxed output is empty");

    // Demux the just-produced MP4 and AV1-decode every frame.
    let demuxed = demux::demux(&mp4).expect("demux own output");
    assert_eq!(
        demuxed.codec, "av1",
        "expected AV1 codec from rav1e+mp4 mux"
    );
    assert!(
        !demuxed.samples.is_empty(),
        "demux of own muxed output produced zero samples — muxer bug"
    );

    let info = StreamInfo {
        codec: "av1".into(),
        width: W,
        height: H,
        frame_rate: FPS,
        duration: FRAMES as f64 / FPS,
        pixel_format: PixelFormat::Yuv420p,
        color_space: ColorSpace::Bt709,
        total_frames: FRAMES as u64,
        bitrate: 0,
        color_metadata: Default::default(),
    };
    let Some(mut decoder) = common::try_av1_decoder(info) else {
        return;
    };
    let mut recovered: Vec<u32> = Vec::new();
    let mut total_confident = 0u32;
    let mut total_bits = 0u32;

    // Streaming pattern (matches pipeline::transcode::run): interleave
    // push_sample with decode_next so dav1d can yield pictures incrementally.
    // dav1d_flush inside finish() wipes any queued pictures, so we MUST
    // drain before finish.
    let mut drain = |decoder: &mut Box<dyn Decoder>,
                     recovered: &mut Vec<u32>,
                     total_confident: &mut u32,
                     total_bits: &mut u32| {
        while let Some(frame) = decoder.decode_next().expect("decode_next") {
            let (idx, confident) = recover_index(&frame);
            recovered.push(idx);
            *total_confident += confident;
            *total_bits += NUM_BITS;
            assert!(
                confident as f32 / NUM_BITS as f32 >= 0.80,
                "frame at decode position {} had only {}/{} confident bits \
                 (recovered index {}); pattern is not surviving compression",
                recovered.len() - 1,
                confident,
                NUM_BITS,
                idx
            );
        }
    };
    for s in &demuxed.samples {
        decoder.push_sample(s).expect("push_sample");
        drain(
            &mut decoder,
            &mut recovered,
            &mut total_confident,
            &mut total_bits,
        );
    }
    decoder.finish().expect("finish");
    drain(
        &mut decoder,
        &mut recovered,
        &mut total_confident,
        &mut total_bits,
    );

    eprintln!(
        "fidelity: encoded {} frames, decoded {} frames, recovered {:?}, \
         bit-confidence {}/{} ({:.1}%)",
        FRAMES,
        recovered.len(),
        recovered,
        total_confident,
        total_bits,
        100.0 * total_confident as f32 / total_bits as f32
    );

    // rav1e at speed_preset=10 buffers aggressively (per-frame lookahead)
    // and only emits some packets after flush. We require ≥4 successfully
    // decoded frames — enough to prove the pipeline is depicting input,
    // not enough to false-fail on rav1e's buffering quirks. The real
    // fidelity contract is the per-frame bit confidence + monotonicity
    // assertions below.
    assert!(
        recovered.len() >= 4,
        "recovered only {} frames out of {} sent — pipeline produced \
         too few decodable frames. A grey-screen bug also lands here, \
         but more commonly this means rav1e/dav1d round-trip broke.",
        recovered.len(),
        FRAMES
    );

    // No constant-frame replacement. A grey-screen bug recovers the
    // SAME index (0) for every frame; a frozen-frame bug recovers the
    // same value > 0 for every frame. Both are caught here.
    let unique: std::collections::BTreeSet<u32> = recovered.iter().copied().collect();
    assert!(
        unique.len() == recovered.len(),
        "recovered indices have only {} unique values out of {} frames — \
         this is the grey-screen / frozen-frame failure mode. Indices: {:?}",
        unique.len(),
        recovered.len(),
        recovered
    );
    // Strictly monotonic in the recovered sequence (AV1 has no B-frames
    // so no reordering). Catches a "decoder shuffles frames" bug.
    for w in recovered.windows(2) {
        assert!(
            w[1] > w[0],
            "recovered indices not strictly monotonic: {:?}",
            recovered
        );
    }
    // First recovered index should be 0 — rav1e + AV1 emit frames in
    // input order, so the first decoded frame must be the first input
    // frame. Catches a "decoder skips initial frames" bug.
    assert_eq!(
        recovered[0], 0,
        "first recovered index = {}, expected 0; decoder dropped initial frames",
        recovered[0]
    );
}

/// A separate, smaller test that proves the recovery code itself works
/// against a known-good frame the encoder hasn't seen. Catches the case
/// where `recover_index` is silently wrong in a way that lets a broken
/// pipeline pass.
#[test]
fn recover_index_round_trip_on_synthetic_input() {
    for i in 0..8u32 {
        let f = make_indexed_frame(W, H, i);
        let (idx, conf) = recover_index(&f);
        assert_eq!(idx, i, "synthetic recovery mismatch for frame {}", i);
        assert_eq!(conf, NUM_BITS, "synthetic input must be 8/8 confident");
    }
}

/// Negative control: a uniformly grey frame (the bug shape this whole test
/// suite exists to catch). Recovery must NOT return the expected index;
/// instead it returns 0 with low confidence — which is what the main test
/// asserts against.
#[test]
fn recover_index_on_grey_frame_returns_zero_with_no_confidence() {
    let w = W as usize;
    let h = H as usize;
    let y_size = w * h;
    let uv_size = y_size / 4;
    let mut buf = vec![128u8; y_size]; // mid-grey luma
    buf.extend(std::iter::repeat(128u8).take(2 * uv_size));
    let grey = VideoFrame::new(
        Bytes::from(buf),
        W,
        H,
        PixelFormat::Yuv420p,
        ColorSpace::Bt709,
        0,
    );
    let (idx, conf) = recover_index(&grey);
    assert_eq!(idx, 0, "grey frame must recover as index 0");
    assert_eq!(conf, 0, "grey frame must have ZERO confident bits");
}