From 9c8de5468bb342e25e8ec5070b700bbf0130e13d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 30 Sep 2026 06:14:18 -0400 Subject: [PATCH] M2.2/M2.3: CAVLC macroblock decoding, intra and inter prediction, DPB - CAVLC residual parsing (coeff_token / total_zeros / run_before tables taken from the spec text, level decoding with escape codes), macroblock layer for I/P/B slices (mb_type, sub_mb_type, mb_pred, sub_mb_pred, me(v) cbp, mb_qp_delta), I_PCM - intra prediction: Intra4x4 (9 modes), Intra8x8 with reference filtering, Intra16x16, chroma; constrained_intra_pred - dequantisation with flat/custom scaling matrices; 4x4, 8x8, luma DC Hadamard and 2x2 chroma DC inverse transforms - inter prediction: median / directional mv prediction, P_Skip, B spatial and temporal direct (direct_8x8_inference), B_Skip, 6-tap luma and bilinear chroma interpolation, explicit and implicit weighted prediction - reference lists (init + modification), DPB with sliding window, all MMCOs, frame_num gaps, bumping / max_num_reorder_frames output order, IDR and no_output_of_prior_pics - public Decoder/Picture API (Annex-B and avcC input, pts passes through reordering) Deblocking (M2.5) follows in the next commit, so the bit-exact fixtures in this commit are the no-deblock ones. CABAC slices return Unsupported. Co-Authored-By: Claude Opus 5.5 --- crates/h264/examples/h264dec.rs | 35 + crates/h264/src/cavlc.rs | 281 +++++++ crates/h264/src/cavlc_tables.rs | 117 +++ crates/h264/src/decoder.rs | 412 +++++++++ crates/h264/src/dpb.rs | 409 +++++++++ crates/h264/src/inter.rs | 303 +++++++ crates/h264/src/intra.rs | 462 ++++++++++ crates/h264/src/lib.rs | 66 +- crates/h264/src/mbtypes.rs | 235 ++++++ crates/h264/src/picture.rs | 144 ++++ crates/h264/src/slicedec.rs | 1340 ++++++++++++++++++++++++++++++ crates/h264/src/transform.rs | 228 +++++ crates/h264/tests/common/mod.rs | 87 ++ crates/h264/tests/conformance.rs | 20 + 14 files changed, 4137 insertions(+), 2 deletions(-) create mode 100644 crates/h264/examples/h264dec.rs create mode 100644 crates/h264/src/cavlc.rs create mode 100644 crates/h264/src/cavlc_tables.rs create mode 100644 crates/h264/src/decoder.rs create mode 100644 crates/h264/src/dpb.rs create mode 100644 crates/h264/src/inter.rs create mode 100644 crates/h264/src/intra.rs create mode 100644 crates/h264/src/mbtypes.rs create mode 100644 crates/h264/src/picture.rs create mode 100644 crates/h264/src/slicedec.rs create mode 100644 crates/h264/src/transform.rs create mode 100644 crates/h264/tests/conformance.rs diff --git a/crates/h264/examples/h264dec.rs b/crates/h264/examples/h264dec.rs new file mode 100644 index 00000000..052b8a1b --- /dev/null +++ b/crates/h264/examples/h264dec.rs @@ -0,0 +1,35 @@ +//! Decode an Annex-B H.264 file to raw yuv420p: `h264dec in.h264 [out.yuv]`. +//! Prints timing information to stderr. + +use filmcraft_h264::Decoder; +use std::io::Write; + +fn main() { + let args: Vec = std::env::args().collect(); + if args.len() < 2 { + eprintln!("usage: h264dec in.h264 [out.yuv]"); + std::process::exit(2); + } + let data = std::fs::read(&args[1]).expect("read input"); + let mut out = args.get(2).map(|p| std::io::BufWriter::new(std::fs::File::create(p).expect("create output"))); + let mut dec = Decoder::new(); + let t0 = std::time::Instant::now(); + let mut n = 0usize; + let mut write = |pics: Vec| { + for p in pics { + n += 1; + if let Some(o) = out.as_mut() { + o.write_all(&p.y).unwrap(); + o.write_all(&p.u).unwrap(); + o.write_all(&p.v).unwrap(); + } + } + }; + match dec.decode(&data, 0) { + Ok(p) => write(p), + Err(e) => eprintln!("decode error: {e}"), + } + write(dec.flush()); + let dt = t0.elapsed().as_secs_f64(); + eprintln!("{n} frames in {dt:.3}s ({:.1} fps)", n as f64 / dt); +} diff --git a/crates/h264/src/cavlc.rs b/crates/h264/src/cavlc.rs new file mode 100644 index 00000000..28ebdaa8 --- /dev/null +++ b/crates/h264/src/cavlc.rs @@ -0,0 +1,281 @@ +//! CAVLC residual block parsing (9.2). + +use crate::cavlc_tables::*; +use crate::error::{Result, ensure}; +use filmcraft_bitstream::BitReader; +use std::sync::OnceLock; + +/// Prefix-code lookup keyed by (number of leading zeros, next 5 bits after the first one bit). +pub(crate) struct VlcTable { + /// entry = (len << 8) | value, 0 = invalid. + entries: Box<[u16; 17 * 32]>, +} + +impl VlcTable { + pub fn build(codes: &[(&str, u8)]) -> Self { + let mut entries = Box::new([0u16; 17 * 32]); + for &(code, value) in codes { + if code.is_empty() { + continue; + } + let len = code.len(); + let lz = code.bytes().take_while(|&b| b == b'0').count(); + let val = ((len as u16) << 8) | value as u16; + if lz == len { + // all-zero code: any continuation with at least `len` zeros + for l in lz..=16 { + for s in 0..32 { + entries[l * 32 + s] = val; + } + } + continue; + } + let suffix = &code[lz + 1..]; + assert!(suffix.len() <= 5, "suffix too long in {code}"); + let sbits = u32::from_str_radix(if suffix.is_empty() { "0" } else { suffix }, 2).unwrap(); + let free = 5 - suffix.len(); + for ext in 0..(1u32 << free) { + let s = (sbits << free) | ext; + entries[lz * 32 + s as usize] = val; + } + } + Self { entries } + } + + #[inline] + pub fn decode(&self, r: &mut BitReader) -> Result { + let bits = r.peek(32); + let lz = bits.leading_zeros().min(16); + let s = if lz >= 16 { 0 } else { ((bits << lz) << 1) >> 27 }; + let e = self.entries[(lz * 32 + s) as usize]; + ensure!(e != 0, "invalid VLC code"); + r.skip((e >> 8) as usize)?; + Ok(e as u8) + } +} + +struct Tables { + /// coeff_token per nC class; value = (TrailingOnes << 5) | TotalCoeff + coeff_token: [VlcTable; 6], + total_zeros_4x4: [VlcTable; 15], + total_zeros_2x2: [VlcTable; 3], + total_zeros_2x4: [VlcTable; 7], + run_before: [VlcTable; 7], +} + +fn tables() -> &'static Tables { + static T: OnceLock = OnceLock::new(); + T.get_or_init(|| { + let coeff_token = std::array::from_fn(|col| { + let codes: Vec<(&str, u8)> = COEFF_TOKEN.iter().map(|&(t1, tc, ref c)| (c[col], (t1 << 5) | tc)).collect(); + VlcTable::build(&codes) + }); + fn simple(row: &[&'static str; N]) -> VlcTable { + let codes: Vec<(&str, u8)> = row.iter().enumerate().map(|(i, &c)| (c, i as u8)).collect(); + VlcTable::build(&codes) + } + Tables { + coeff_token, + total_zeros_4x4: std::array::from_fn(|i| simple(&TOTAL_ZEROS_4X4[i])), + total_zeros_2x2: std::array::from_fn(|i| simple(&TOTAL_ZEROS_2X2[i])), + total_zeros_2x4: std::array::from_fn(|i| simple(&TOTAL_ZEROS_2X4[i])), + run_before: std::array::from_fn(|i| simple(&RUN_BEFORE[i])), + } + }) +} + +/// Force table construction (so the first slice does not pay for it). +pub(crate) fn init_tables() { + let _ = tables(); +} + +/// Parse residual_block_cavlc(). Coefficients are written to `coeff[start_idx..=end_idx]` in scan order +/// (entries outside the coded range are left untouched; the caller zeroes them). +/// `nc` is the nC value (-1 for 4:2:0 chroma DC, -2 for 4:2:2 chroma DC). Returns TotalCoeff. +pub(crate) fn residual_block(r: &mut BitReader, coeff: &mut [i32], start_idx: usize, end_idx: usize, max_num_coeff: usize, nc: i32) -> Result { + let t = tables(); + let col = match nc { + -1 => 4, + -2 => 5, + 0..=1 => 0, + 2..=3 => 1, + 4..=7 => 2, + _ => 3, + }; + let tok = t.coeff_token[col].decode(r)?; + let total_coeff = (tok & 31) as usize; + let trailing_ones = (tok >> 5) as usize; + if total_coeff == 0 { + return Ok(0); + } + ensure!(total_coeff <= max_num_coeff, "TotalCoeff exceeds maxNumCoeff"); + let mut level = [0i32; 16]; + let mut suffix_length: u32 = if total_coeff > 10 && trailing_ones < 3 { 1 } else { 0 }; + for (i, lv) in level.iter_mut().enumerate().take(total_coeff) { + if i < trailing_ones { + *lv = if r.read_bit()? { -1 } else { 1 }; + continue; + } + // level_prefix + let bits = r.peek(32); + let level_prefix = bits.leading_zeros(); + ensure!(level_prefix <= 25, "level_prefix too large"); + r.skip(level_prefix as usize + 1)?; + let level_suffix_size = if level_prefix == 14 && suffix_length == 0 { + 4 + } else if level_prefix >= 15 { + level_prefix - 3 + } else { + suffix_length + }; + let level_suffix = if level_suffix_size > 0 { r.read_bits(level_suffix_size)? as i32 } else { 0 }; + let mut level_code = ((level_prefix.min(15) as i32) << suffix_length) + level_suffix; + if level_prefix >= 15 && suffix_length == 0 { + level_code += 15; + } + if level_prefix >= 16 { + level_code += (1 << (level_prefix - 3)) - 4096; + } + if i == trailing_ones && trailing_ones < 3 { + level_code += 2; + } + *lv = if level_code % 2 == 0 { (level_code + 2) >> 1 } else { (-level_code - 1) >> 1 }; + if suffix_length == 0 { + suffix_length = 1; + } + if lv.abs() > (3 << (suffix_length - 1)) && suffix_length < 6 { + suffix_length += 1; + } + } + let mut zeros_left = if total_coeff < end_idx - start_idx + 1 { + let tz_idx = total_coeff - 1; + let tz = match max_num_coeff { + 4 => t.total_zeros_2x2[tz_idx].decode(r)?, + 8 => t.total_zeros_2x4[tz_idx].decode(r)?, + _ => t.total_zeros_4x4[tz_idx].decode(r)?, + } as usize; + ensure!(tz + total_coeff <= end_idx - start_idx + 1, "total_zeros out of range"); + tz + } else { + 0 + }; + // runs and placement: highest-frequency coefficient first + let mut pos = start_idx + total_coeff + zeros_left - 1; + for (i, &lv) in level.iter().enumerate().take(total_coeff) { + coeff[pos] = lv; + if i + 1 == total_coeff { + break; + } + let run = if zeros_left > 0 { + let run = t.run_before[zeros_left.min(7) - 1].decode(r)? as usize; + ensure!(run <= zeros_left, "run_before out of range"); + zeros_left -= run; + run + } else { + 0 + }; + pos -= 1 + run; + } + Ok(total_coeff as u8) +} + +#[cfg(test)] +mod tests { + use super::*; + use filmcraft_bitstream::BitWriter; + + fn kraft(codes: &[&str]) -> f64 { + codes.iter().filter(|c| !c.is_empty()).map(|c| 0.5f64.powi(c.len() as i32)).sum() + } + + fn prefix_free(codes: &[&str]) -> bool { + let c: Vec<&&str> = codes.iter().filter(|c| !c.is_empty()).collect(); + for (i, a) in c.iter().enumerate() { + for (j, b) in c.iter().enumerate() { + if i != j && b.starts_with(**a) { + return false; + } + } + } + true + } + + #[test] + fn coeff_token_tables_are_prefix_codes() { + for col in 0..6 { + let codes: Vec<&str> = COEFF_TOKEN.iter().map(|e| e.2[col]).collect(); + assert!(prefix_free(&codes), "column {col}"); + let k = kraft(&codes); + assert!(k <= 1.0 + 1e-12, "column {col} kraft {k}"); + } + // nC>=8 is a complete 6-bit fixed code except the unused 000011 pattern + let codes: Vec<&str> = COEFF_TOKEN.iter().map(|e| e.2[3]).collect(); + assert!(codes.iter().all(|c| c.len() == 6)); + } + + #[test] + fn total_zeros_and_run_before_are_complete_prefix_codes() { + for (i, row) in TOTAL_ZEROS_4X4.iter().enumerate() { + assert!(prefix_free(row)); + // tzVlcIndex 1 leaves the all-zero 9-bit code unused; all others are complete + let k = kraft(row); + if i == 0 { + assert!((k - (1.0 - 1.0 / 512.0)).abs() < 1e-12, "{row:?}"); + } else { + assert!((k - 1.0).abs() < 1e-12, "{row:?}"); + } + } + for row in TOTAL_ZEROS_2X2.iter() { + assert!(prefix_free(row)); + assert!((kraft(row) - 1.0).abs() < 1e-12); + } + for row in TOTAL_ZEROS_2X4.iter() { + assert!(prefix_free(row)); + assert!((kraft(row) - 1.0).abs() < 1e-12); + } + for row in RUN_BEFORE.iter() { + assert!(prefix_free(row)); + } + } + + fn write_code(w: &mut BitWriter, code: &str) { + for b in code.bytes() { + w.write_bit(b == b'1'); + } + } + + #[test] + fn decodes_spec_style_example() { + // Block: 0 3 -1 0 | 0 -1 1 0 | 1 0 0 0 | 0 0 0 0 (zig-zag order), nC = 0. + // TotalCoeff 5, TrailingOnes 3: coeff_token 0000 100. + // trailing ones (reverse order): +1, -1, -1 -> sign bits 0 1 1 + // level -1? No: remaining levels (reverse): 1 (idx 4?) ... use the classic example: + // coefficients in scan order: [0, 3, -1, 0, 0, -1, 1, 0, 1, 0, ...] + // reverse order non-zero: 1 (pos 8), 1 (pos 6), -1 (pos 5), -1 (pos 2), 3 (pos 1) + // TotalCoeff 5, TrailingOnes 3 (1, 1, -1). + let mut w = BitWriter::new(); + write_code(&mut w, "0000100"); // coeff_token TC=5 T1=3, nC 0..2 + write_code(&mut w, "0"); // +1 + write_code(&mut w, "0"); // +1 + write_code(&mut w, "1"); // -1 + // level -1 with suffixLength 0: levelCode = 1 (T1 == 3 so no +2) -> prefix 1 => "01" + write_code(&mut w, "01"); + // level 3 with suffixLength 1: levelCode = 4 -> prefix 2, suffix 0 => "001" "0" + write_code(&mut w, "0010"); + // total_zeros = 4 (positions 0, 3, 4, 7) with tzVlcIndex 5 -> "110" + write_code(&mut w, "110"); + // run_before: zerosLeft 4 at pos 8: run 1 -> "10"; zerosLeft 3 at pos 6: run 0 -> "11"; + // zerosLeft 3 at pos 5: run 2 -> "01"; zerosLeft 1 at pos 2: run 0 -> "1" + write_code(&mut w, "10"); + write_code(&mut w, "11"); + write_code(&mut w, "01"); + write_code(&mut w, "1"); + w.rbsp_trailing(); + let b = w.finish(); + let mut r = BitReader::new(&b); + let mut c = [0i32; 16]; + let tc = residual_block(&mut r, &mut c, 0, 15, 16, 0).unwrap(); + assert_eq!(tc, 5); + assert_eq!(&c[..10], &[0, 3, -1, 0, 0, -1, 1, 0, 1, 0]); + } +} diff --git a/crates/h264/src/cavlc_tables.rs b/crates/h264/src/cavlc_tables.rs new file mode 100644 index 00000000..ba509e2f --- /dev/null +++ b/crates/h264/src/cavlc_tables.rs @@ -0,0 +1,117 @@ +// Generated from ITU-T H.264 (06/2026) Tables 9-5, 9-7..9-10 by a table-extraction script. +// Codes are bit strings; "" marks an unused entry. + +/// (TrailingOnes, TotalCoeff, [nC 0..2, 2..4, 4..8, 8.., -1, -2]) +#[rustfmt::skip] +pub(crate) static COEFF_TOKEN: [(u8, u8, [&str; 6]); 62] = [ + (0, 0, ["1", "11", "1111", "000011", "01", "1"]), + (0, 1, ["000101", "001011", "001111", "000000", "000111", "0001111"]), + (1, 1, ["01", "10", "1110", "000001", "1", "01"]), + (0, 2, ["00000111", "000111", "001011", "000100", "000100", "0001110"]), + (1, 2, ["000100", "00111", "01111", "000101", "000110", "0001101"]), + (2, 2, ["001", "011", "1101", "000110", "001", "001"]), + (0, 3, ["000000111", "0000111", "001000", "001000", "000011", "000000111"]), + (1, 3, ["00000110", "001010", "01100", "001001", "0000011", "0001100"]), + (2, 3, ["0000101", "001001", "01110", "001010", "0000010", "0001011"]), + (3, 3, ["00011", "0101", "1100", "001011", "000101", "00001"]), + (0, 4, ["0000000111", "00000111", "0001111", "001100", "000010", "000000110"]), + (1, 4, ["000000110", "000110", "01010", "001101", "00000011", "000000101"]), + (2, 4, ["00000101", "000101", "01011", "001110", "00000010", "0001010"]), + (3, 4, ["000011", "0100", "1011", "001111", "0000000", "000001"]), + (0, 5, ["00000000111", "00000100", "0001011", "010000", "", "0000000111"]), + (1, 5, ["0000000110", "0000110", "01000", "010001", "", "0000000110"]), + (2, 5, ["000000101", "0000101", "01001", "010010", "", "000000100"]), + (3, 5, ["0000100", "00110", "1010", "010011", "", "0001001"]), + (0, 6, ["0000000001111", "000000111", "0001001", "010100", "", "00000000111"]), + (1, 6, ["00000000110", "00000110", "001110", "010101", "", "00000000110"]), + (2, 6, ["0000000101", "00000101", "001101", "010110", "", "0000000101"]), + (3, 6, ["00000100", "001000", "1001", "010111", "", "0001000"]), + (0, 7, ["0000000001011", "00000001111", "0001000", "011000", "", "000000000111"]), + (1, 7, ["0000000001110", "000000110", "001010", "011001", "", "000000000110"]), + (2, 7, ["00000000101", "000000101", "001001", "011010", "", "00000000101"]), + (3, 7, ["000000100", "000100", "1000", "011011", "", "0000000100"]), + (0, 8, ["0000000001000", "00000001011", "00001111", "011100", "", "0000000000111"]), + (1, 8, ["0000000001010", "00000001110", "0001110", "011101", "", "000000000101"]), + (2, 8, ["0000000001101", "00000001101", "0001101", "011110", "", "000000000100"]), + (3, 8, ["0000000100", "0000100", "01101", "011111", "", "00000000100"]), + (0, 9, ["00000000001111", "000000001111", "00001011", "100000", "", ""]), + (1, 9, ["00000000001110", "00000001010", "00001110", "100001", "", ""]), + (2, 9, ["0000000001001", "00000001001", "0001010", "100010", "", ""]), + (3, 9, ["00000000100", "000000100", "001100", "100011", "", ""]), + (0, 10, ["00000000001011", "000000001011", "000001111", "100100", "", ""]), + (1, 10, ["00000000001010", "000000001110", "00001010", "100101", "", ""]), + (2, 10, ["00000000001101", "000000001101", "00001101", "100110", "", ""]), + (3, 10, ["0000000001100", "00000001100", "0001100", "100111", "", ""]), + (0, 11, ["000000000001111", "000000001000", "000001011", "101000", "", ""]), + (1, 11, ["000000000001110", "000000001010", "000001110", "101001", "", ""]), + (2, 11, ["00000000001001", "000000001001", "00001001", "101010", "", ""]), + (3, 11, ["00000000001100", "00000001000", "00001100", "101011", "", ""]), + (0, 12, ["000000000001011", "0000000001111", "000001000", "101100", "", ""]), + (1, 12, ["000000000001010", "0000000001110", "000001010", "101101", "", ""]), + (2, 12, ["000000000001101", "0000000001101", "000001101", "101110", "", ""]), + (3, 12, ["00000000001000", "000000001100", "00001000", "101111", "", ""]), + (0, 13, ["0000000000001111", "0000000001011", "0000001101", "110000", "", ""]), + (1, 13, ["000000000000001", "0000000001010", "000000111", "110001", "", ""]), + (2, 13, ["000000000001001", "0000000001001", "000001001", "110010", "", ""]), + (3, 13, ["000000000001100", "0000000001100", "000001100", "110011", "", ""]), + (0, 14, ["0000000000001011", "0000000000111", "0000001001", "110100", "", ""]), + (1, 14, ["0000000000001110", "00000000001011", "0000001100", "110101", "", ""]), + (2, 14, ["0000000000001101", "0000000000110", "0000001011", "110110", "", ""]), + (3, 14, ["000000000001000", "0000000001000", "0000001010", "110111", "", ""]), + (0, 15, ["0000000000000111", "00000000001001", "0000000101", "111000", "", ""]), + (1, 15, ["0000000000001010", "00000000001000", "0000001000", "111001", "", ""]), + (2, 15, ["0000000000001001", "00000000001010", "0000000111", "111010", "", ""]), + (3, 15, ["0000000000001100", "0000000000001", "0000000110", "111011", "", ""]), + (0, 16, ["0000000000000100", "00000000000111", "0000000001", "111100", "", ""]), + (1, 16, ["0000000000000110", "00000000000110", "0000000100", "111101", "", ""]), + (2, 16, ["0000000000000101", "00000000000101", "0000000011", "111110", "", ""]), + (3, 16, ["0000000000001000", "00000000000100", "0000000010", "111111", "", ""]), +]; +/// total_zeros for 4x4 blocks: [tzVlcIndex-1][total_zeros] +#[rustfmt::skip] +pub(crate) static TOTAL_ZEROS_4X4: [[&str; 16]; 15] = [ + ["1", "011", "010", "0011", "0010", "00011", "00010", "000011", "000010", "0000011", "0000010", "00000011", "00000010", "000000011", "000000010", "000000001"], + ["111", "110", "101", "100", "011", "0101", "0100", "0011", "0010", "00011", "00010", "000011", "000010", "000001", "000000", ""], + ["0101", "111", "110", "101", "0100", "0011", "100", "011", "0010", "00011", "00010", "000001", "00001", "000000", "", ""], + ["00011", "111", "0101", "0100", "110", "101", "100", "0011", "011", "0010", "00010", "00001", "00000", "", "", ""], + ["0101", "0100", "0011", "111", "110", "101", "100", "011", "0010", "00001", "0001", "00000", "", "", "", ""], + ["000001", "00001", "111", "110", "101", "100", "011", "010", "0001", "001", "000000", "", "", "", "", ""], + ["000001", "00001", "101", "100", "011", "11", "010", "0001", "001", "000000", "", "", "", "", "", ""], + ["000001", "0001", "00001", "011", "11", "10", "010", "001", "000000", "", "", "", "", "", "", ""], + ["000001", "000000", "0001", "11", "10", "001", "01", "00001", "", "", "", "", "", "", "", ""], + ["00001", "00000", "001", "11", "10", "01", "0001", "", "", "", "", "", "", "", "", ""], + ["0000", "0001", "001", "010", "1", "011", "", "", "", "", "", "", "", "", "", ""], + ["0000", "0001", "01", "1", "001", "", "", "", "", "", "", "", "", "", "", ""], + ["000", "001", "1", "01", "", "", "", "", "", "", "", "", "", "", "", ""], + ["00", "01", "1", "", "", "", "", "", "", "", "", "", "", "", "", ""], + ["0", "1", "", "", "", "", "", "", "", "", "", "", "", "", "", ""], +]; +/// total_zeros for chroma DC 2x2: [tzVlcIndex-1][total_zeros] +#[rustfmt::skip] +pub(crate) static TOTAL_ZEROS_2X2: [[&str; 4]; 3] = [ + ["1", "01", "001", "000"], + ["1", "01", "00", ""], + ["1", "0", "", ""], +]; +/// total_zeros for chroma DC 2x4: [tzVlcIndex-1][total_zeros] +#[rustfmt::skip] +pub(crate) static TOTAL_ZEROS_2X4: [[&str; 8]; 7] = [ + ["1", "010", "011", "0010", "0011", "0001", "00001", "00000"], + ["000", "01", "001", "100", "101", "110", "111", ""], + ["000", "001", "01", "10", "110", "111", "", ""], + ["110", "00", "01", "10", "111", "", "", ""], + ["00", "01", "10", "11", "", "", "", ""], + ["00", "01", "1", "", "", "", "", ""], + ["0", "1", "", "", "", "", "", ""], +]; +/// run_before: [min(zerosLeft,7)-1][run_before] +#[rustfmt::skip] +pub(crate) static RUN_BEFORE: [[&str; 15]; 7] = [ + ["1", "0", "", "", "", "", "", "", "", "", "", "", "", "", ""], + ["1", "01", "00", "", "", "", "", "", "", "", "", "", "", "", ""], + ["11", "10", "01", "00", "", "", "", "", "", "", "", "", "", "", ""], + ["11", "10", "01", "001", "000", "", "", "", "", "", "", "", "", "", ""], + ["11", "10", "011", "010", "001", "000", "", "", "", "", "", "", "", "", ""], + ["11", "000", "001", "011", "010", "101", "100", "", "", "", "", "", "", "", ""], + ["111", "110", "101", "100", "011", "010", "001", "0001", "00001", "000001", "0000001", "00000001", "000000001", "0000000001", "00000000001"], +]; diff --git a/crates/h264/src/decoder.rs b/crates/h264/src/decoder.rs new file mode 100644 index 00000000..b9ed236b --- /dev/null +++ b/crates/h264/src/decoder.rs @@ -0,0 +1,412 @@ +//! Top-level decoder: NAL dispatch, picture boundaries, POC, DPB, output. + +use crate::dpb::{Dpb, Output, OutputMeta}; +use crate::error::{Error, Result, ensure, unsupported}; +use crate::params::{Pps, Sps}; +use crate::picture::{Frame, MotionField, Planes}; +use crate::slice::{NalHeader, Poc, PocState, SliceHeader, nal_type}; +use crate::slicedec::{PicState, SliceDecoder}; +use crate::transform::LevelScale; +use crate::{ColorInfo, Picture}; +use filmcraft_bitstream::{BitReader, annexb_nals, length_prefixed_nals, unescape_rbsp}; +use std::sync::Arc; + +struct CurPic { + pic: PicState, + first: SliceHeader, + sps: Arc, + poc: Poc, + pts: i64, + key: bool, + has_mmco5: bool, +} + +/// H.264 decoder. +pub struct Decoder { + spss: Vec>>, + ppss: Vec>>, + nal_length_size: Option, + dpb: Dpb, + poc_state: PocState, + prev_ref_frame_num: u32, + cur: Option, + next_id: u32, + active_sps: Option>, + ls_cache: Vec<(Arc, Arc)>, + out: Vec, +} + +impl Default for Decoder { + fn default() -> Self { + Self::new() + } +} + +/// Is `sh` the first slice of a new primary coded picture relative to `prev` (7.4.1.2.4)? +fn is_new_picture(prev: &SliceHeader, sh: &SliceHeader, sps: &Sps) -> bool { + if sh.frame_num != prev.frame_num + || sh.pps_id != prev.pps_id + || sh.field_pic != prev.field_pic + || sh.bottom_field != prev.bottom_field + || (sh.nal_ref_idc == 0) != (prev.nal_ref_idc == 0) + || sh.idr != prev.idr + || (sh.idr && prev.idr && sh.idr_pic_id != prev.idr_pic_id) + { + return true; + } + if sps.pic_order_cnt_type == 0 && (sh.pic_order_cnt_lsb != prev.pic_order_cnt_lsb || sh.delta_pic_order_cnt_bottom != prev.delta_pic_order_cnt_bottom) { + return true; + } + if sps.pic_order_cnt_type == 1 && sh.delta_pic_order_cnt != prev.delta_pic_order_cnt { + return true; + } + false +} + +impl Decoder { + pub fn new() -> Self { + crate::cavlc::init_tables(); + Decoder { + spss: vec![None; 32], + ppss: vec![None; 256], + nal_length_size: None, + dpb: Dpb::new(), + poc_state: PocState::default(), + prev_ref_frame_num: 0, + cur: None, + next_id: 1, + active_sps: None, + ls_cache: Vec::new(), + out: Vec::new(), + } + } + + /// Configure from an `avcC` (AVCDecoderConfigurationRecord) box payload. + pub fn from_avcc(avcc: &[u8]) -> Result { + let mut d = Decoder::new(); + ensure!(avcc.len() >= 7, "avcC too short"); + ensure!(avcc[0] == 1, "unsupported avcC version {}", avcc[0]); + d.nal_length_size = Some((avcc[4] & 3) as usize + 1); + let mut pos = 5; + let read_sets = |pos: &mut usize, count: usize, d: &mut Decoder| -> Result<()> { + for _ in 0..count { + ensure!(*pos + 2 <= avcc.len(), "avcC truncated"); + let len = u16::from_be_bytes([avcc[*pos], avcc[*pos + 1]]) as usize; + *pos += 2; + ensure!(*pos + len <= avcc.len(), "avcC truncated"); + d.handle_nal(&avcc[*pos..*pos + len], 0)?; + *pos += len; + } + Ok(()) + }; + let nsps = (avcc[pos] & 0x1f) as usize; + pos += 1; + read_sets(&mut pos, nsps, &mut d)?; + ensure!(pos < avcc.len(), "avcC truncated"); + let npps = avcc[pos] as usize; + pos += 1; + read_sets(&mut pos, npps, &mut d)?; + Ok(d) + } + + /// Length of NAL length prefixes (from avcC), or None for Annex-B input. + pub fn nal_length_size(&self) -> Option { + self.nal_length_size + } + + /// Decode one access unit. Returns pictures that became ready for output, in output order. + pub fn decode(&mut self, data: &[u8], pts: i64) -> Result> { + let nals = match self.nal_length_size { + Some(n) => length_prefixed_nals(data, n)?, + None => annexb_nals(data), + }; + let mut result = Ok(()); + for nal in nals { + if let Err(e) = self.handle_nal(nal, pts) { + result = Err(e); + break; + } + } + // Finish a complete picture at the end of the access unit. + if self.cur.as_ref().is_some_and(|c| c.pic.decoded_mbs >= c.pic.mbs.len()) { + self.finish_picture(); + } + result?; + Ok(std::mem::take(&mut self.out)) + } + + /// Output all remaining pictures. + pub fn flush(&mut self) -> Vec { + if self.cur.is_some() { + self.finish_picture(); + } + let mut outs = Vec::new(); + self.dpb.flush(&mut outs); + self.emit(outs); + std::mem::take(&mut self.out) + } + + fn handle_nal(&mut self, nal: &[u8], pts: i64) -> Result<()> { + if nal.is_empty() { + return Ok(()); + } + let hdr = NalHeader::parse(nal[0])?; + match hdr.nal_unit_type { + nal_type::SPS => { + let sps = Sps::parse(&unescape_rbsp(&nal[1..]))?; + let id = sps.id as usize; + self.spss[id] = Some(Arc::new(sps)); + } + nal_type::PPS => { + let rbsp = unescape_rbsp(&nal[1..]); + let pps = Pps::parse(&rbsp, &self.spss_plain())?; + let id = pps.id as usize; + self.ppss[id] = Some(Arc::new(pps)); + } + nal_type::SLICE | nal_type::IDR => self.handle_slice(nal, hdr, pts)?, + nal_type::SLICE_DPA | nal_type::SLICE_DPB | nal_type::SLICE_DPC => { + return unsupported("data partitioning (Extended profile)"); + } + nal_type::END_SEQ | nal_type::END_STREAM if self.cur.is_some() => self.finish_picture(), + // SEI, AUD, filler, SPS extension, prefix NAL, subset SPS, auxiliary and extension slices: skipped. + _ => {} + } + Ok(()) + } + + fn spss_plain(&self) -> Vec> { + self.spss.iter().map(|s| s.as_ref().map(|s| (**s).clone())).collect() + } + + fn level_scale(&mut self, pps: &Arc) -> Arc { + if let Some((_, ls)) = self.ls_cache.iter().find(|(p, _)| Arc::ptr_eq(p, pps)) { + return ls.clone(); + } + let ls: Arc = Arc::from(LevelScale::new(&pps.scaling)); + self.ls_cache.retain(|(p, _)| self.ppss.iter().flatten().any(|q| Arc::ptr_eq(p, q))); + self.ls_cache.push((pps.clone(), ls.clone())); + ls + } + + fn handle_slice(&mut self, nal: &[u8], hdr: NalHeader, pts: i64) -> Result<()> { + let rbsp = unescape_rbsp(&nal[1..]); + let ppss = &self.ppss; + let spss = &self.spss; + let mut found: Option<(Arc, Arc)> = None; + let (sh, _, _) = SliceHeader::parse(&rbsp, hdr, |id| { + let pps = ppss[id as usize].as_ref().ok_or_else(|| Error::MissingParameterSet(format!("PPS {id}")))?; + let sps = spss[pps.sps_id as usize].as_ref().ok_or_else(|| Error::MissingParameterSet(format!("SPS {}", pps.sps_id)))?; + found = Some((pps.clone(), sps.clone())); + Ok((&**pps, &**sps)) + })?; + let (pps, sps) = found.expect("lookup succeeded"); + sps.check_supported()?; + if pps.num_slice_groups > 1 { + return unsupported("slice groups (FMO)"); + } + if sh.field_pic { + return unsupported("field pictures"); + } + if matches!(sh.slice_type, crate::slice::SliceType::Sp | crate::slice::SliceType::Si) { + return unsupported("SP/SI slices"); + } + if sh.redundant_pic_cnt > 0 { + return Ok(()); // redundant slices are ignored + } + let new_pic = match &self.cur { + None => true, + Some(c) => is_new_picture(&c.first, &sh, &sps) || (sh.first_mb_in_slice == 0 && c.pic.decoded_mbs > 0) || !Arc::ptr_eq(&c.sps, &sps), + }; + if new_pic { + if self.cur.is_some() { + self.finish_picture(); + } + self.start_picture(&sh, &sps, pts)?; + } + let ls = self.level_scale(&pps); + let cur = self.cur.as_mut().expect("picture started"); + if sh.has_mmco5() { + cur.has_mmco5 = true; + } + let refs = self.dpb.build_ref_lists(&sh, cur.poc.frame(), sps.max_frame_num())?; + if !sh.slice_type.is_intra() { + ensure!(!refs[0].is_empty(), "no reference pictures available for inter slice"); + if sh.slice_type.is_b() { + ensure!(!refs[1].is_empty(), "empty RefPicList1 in B slice"); + } + } + let mut sd = SliceDecoder::new(&sh, &pps, &sps, &mut cur.pic, &refs, &ls)?; + if pps.entropy_coding_mode { + return unsupported("CABAC entropy coding"); + } else { + let mut r = BitReader::new(&rbsp); + r.seek_bits(sh.header_bits); + sd.decode_cavlc(&mut r)?; + } + Ok(()) + } + + fn start_picture(&mut self, sh: &SliceHeader, sps: &Arc, pts: i64) -> Result<()> { + // Activate SPS; a resolution change flushes the DPB. + let changed = match &self.active_sps { + None => true, + Some(a) => a.width() != sps.width() || a.height() != sps.height() || a.max_dpb_frames() != sps.max_dpb_frames(), + }; + if changed && self.active_sps.is_some() { + let mut outs = Vec::new(); + self.dpb.flush(&mut outs); + self.emit(outs); + self.dpb.entries.clear(); + } + self.active_sps = Some(sps.clone()); + self.dpb.capacity = sps.max_dpb_frames(); + self.dpb.max_reorder = sps.max_num_reorder_frames().min(self.dpb.capacity); + let mb_w = sps.pic_width_in_mbs as usize; + let mb_h = sps.frame_height_in_mbs() as usize; + if sh.idr { + let mut outs = Vec::new(); + self.dpb.idr(sh.no_output_of_prior_pics, &mut outs); + self.emit(outs); + self.prev_ref_frame_num = 0; + } else { + if self.dpb.entries.iter().all(|e| e.mark == crate::dpb::RefMark::Unused) && !sh.slice_type.is_intra() { + // Stream starts without an IDR: nothing to reference; decoding will conceal. + } + let max = sps.max_frame_num(); + if sh.frame_num != self.prev_ref_frame_num && sh.frame_num != (self.prev_ref_frame_num + 1) % max { + // frame_num gap (8.2.5.2) + let (w, h) = (mb_w * 16, mb_h * 16); + let meta = Arc::new(output_meta(sps, pts, false)); + let next_id = &mut self.next_id; + let poc_state = &mut self.poc_state; + let prev_last = self.dpb.entries.iter().filter(|e| !e.non_existing).max_by_key(|e| e.frame.id).map(|e| e.frame.clone()); + let mut make = |_fnum: u32| { + let id = *next_id; + *next_id += 1; + let planes = match &prev_last { + Some(f) => f.planes.clone(), + None => Planes::gray(w, h), + }; + let frame = Arc::new(Frame { id, poc: 0, planes, motion: MotionField::default() }); + (frame, meta.clone()) + }; + let mut upd = |fnum: u32| poc_state.update_gap_frame(fnum, sps); + self.dpb.fill_frame_num_gap(self.prev_ref_frame_num, sh.frame_num, max, sps.max_num_ref_frames as usize, &mut make, &mut upd); + self.prev_ref_frame_num = (sh.frame_num + max - 1) % max; + } + } + let poc = self.poc_state.compute(sh, sps); + let pic = PicState::new(mb_w, mb_h, poc.frame()); + self.cur = Some(CurPic { pic, first: sh.clone(), sps: sps.clone(), poc, pts, key: sh.idr, has_mmco5: false }); + Ok(()) + } +} + +fn output_meta(sps: &Sps, pts: i64, key: bool) -> OutputMeta { + { + let vui = sps.vui.clone().unwrap_or_default(); + OutputMeta { + pts, + key, + crop: sps.crop_rect(), + full_range: vui.full_range, + colour_primaries: vui.colour_primaries, + transfer_characteristics: vui.transfer_characteristics, + matrix_coefficients: vui.matrix_coefficients, + sar: vui.sar, + } + } +} + +impl Decoder { + fn finish_picture(&mut self) { + let Some(cur) = self.cur.take() else { return }; + // motion field for co-located use + let n = cur.pic.mbs.len(); + let mut motion = MotionField { + mv: [vec![[0; 2]; n * 16], vec![[0; 2]; n * 16]], + ref_idx: [vec![-1; n * 4], vec![-1; n * 4]], + ref_id: [vec![u32::MAX; n * 4], vec![u32::MAX; n * 4]], + intra: vec![true; n], + }; + for (i, st) in cur.pic.mbs.iter().enumerate() { + if st.slice_num == u32::MAX { + continue; + } + motion.intra[i] = st.kind.is_intra(); + let sl = &cur.pic.slices[st.slice_num as usize]; + for l in 0..2 { + motion.mv[l][i * 16..i * 16 + 16].copy_from_slice(&st.mv[l]); + for b in 0..4 { + let r = st.ref_idx[l][b]; + motion.ref_idx[l][i * 4 + b] = r; + if r >= 0 { + motion.ref_id[l][i * 4 + b] = sl.ref_ids[l].get(r as usize).copied().unwrap_or(u32::MAX); + } + } + } + } + let mut poc = cur.poc.frame(); + let sh = cur.first.clone(); + if cur.has_mmco5 { + // tempPicOrderCnt adjustment (8.2.1): the picture's POC becomes relative to itself. + poc -= cur.poc.frame(); + } + let id = self.next_id; + self.next_id += 1; + let frame = Arc::new(Frame { id, poc, planes: cur.pic.planes, motion }); + let meta = Arc::new(output_meta(&cur.sps, cur.pts, cur.key)); + let sh_eff = sh.clone(); + self.poc_state.update(&sh_eff, &cur.poc); + if sh.nal_ref_idc != 0 { + self.prev_ref_frame_num = if cur.has_mmco5 { 0 } else { sh.frame_num }; + } + let mut outs = Vec::new(); + self.dpb.store_picture(&sh_eff, frame, poc, cur.sps.max_frame_num(), cur.sps.max_num_ref_frames as usize, meta, &mut outs); + self.emit(outs); + } + + fn emit(&mut self, outs: Vec) { + for o in outs { + self.out.push(make_picture(&o)); + } + } +} + +fn make_picture(o: &Output) -> Picture { + let p = &o.frame.planes; + let (cx, cy, cw, ch) = o.meta.crop; + let (cx, cy, cw, ch) = (cx as usize, cy as usize, cw as usize, ch as usize); + let mut y = Vec::with_capacity(cw * ch); + for r in cy..cy + ch { + y.extend_from_slice(&p.y[r * p.width + cx..r * p.width + cx + cw]); + } + let (ccx, ccy, ccw, cch) = (cx / 2, cy / 2, cw.div_ceil(2), ch.div_ceil(2)); + let mut u = Vec::with_capacity(ccw * cch); + let mut v = Vec::with_capacity(ccw * cch); + for r in ccy..ccy + cch { + u.extend_from_slice(&p.cb[r * p.cwidth + ccx..r * p.cwidth + ccx + ccw]); + v.extend_from_slice(&p.cr[r * p.cwidth + ccx..r * p.cwidth + ccx + ccw]); + } + Picture { + width: cw as u32, + height: ch as u32, + chroma_width: ccw as u32, + chroma_height: cch as u32, + y, + u, + v, + y_stride: cw, + uv_stride: ccw, + pts: o.meta.pts, + poc: o.poc, + key: o.meta.key, + color: ColorInfo { + full_range: o.meta.full_range, + primaries: o.meta.colour_primaries, + transfer: o.meta.transfer_characteristics, + matrix: o.meta.matrix_coefficients, + }, + sar: o.meta.sar, + } +} diff --git a/crates/h264/src/dpb.rs b/crates/h264/src/dpb.rs new file mode 100644 index 00000000..2603c77b --- /dev/null +++ b/crates/h264/src/dpb.rs @@ -0,0 +1,409 @@ +//! Decoded picture buffer: reference marking (8.2.5), reference list construction (8.2.4) and output +//! ordering / bumping (C.4). + +use crate::error::{Result, ensure}; +use crate::picture::{FrameRef, RefPic}; +use crate::slice::{Mmco, SliceHeader, SliceType}; +use std::sync::Arc; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RefMark { + Unused, + Short, + Long, +} + +/// Output-related metadata attached to a decoded frame. +#[derive(Clone, Debug)] +pub struct OutputMeta { + pub pts: i64, + pub key: bool, + pub crop: (u32, u32, u32, u32), + pub full_range: bool, + pub colour_primaries: u8, + pub transfer_characteristics: u8, + pub matrix_coefficients: u8, + pub sar: (u16, u16), +} + +pub struct DpbEntry { + pub frame: FrameRef, + pub poc: i32, + pub frame_num: u32, + pub long_term_frame_idx: u32, + pub mark: RefMark, + pub needed_for_output: bool, + pub non_existing: bool, + pub meta: Arc, +} + +pub struct Dpb { + pub entries: Vec, + /// None = "no long-term frame indices". + pub max_long_term_frame_idx: Option, + pub capacity: usize, + pub max_reorder: usize, +} + +/// A picture leaving the DPB for output. +pub struct Output { + pub frame: FrameRef, + pub poc: i32, + pub meta: Arc, +} + +impl Dpb { + pub fn new() -> Self { + Dpb { entries: Vec::new(), max_long_term_frame_idx: None, capacity: 16, max_reorder: 16 } + } + + fn frame_num_wrap(e: &DpbEntry, cur_frame_num: u32, max_frame_num: u32) -> i32 { + if e.frame_num > cur_frame_num { e.frame_num as i32 - max_frame_num as i32 } else { e.frame_num as i32 } + } + + fn num_refs(&self) -> usize { + self.entries.iter().filter(|e| e.mark != RefMark::Unused).count() + } + + /// Remove entries not needed for output and not used for reference. + fn prune(&mut self) { + self.entries.retain(|e| e.needed_for_output || e.mark != RefMark::Unused); + } + + /// Output the picture with the smallest POC (C.4.5.3 "bumping"). Returns false if none waits. + fn bump(&mut self, out: &mut Vec) -> bool { + let Some(i) = self.entries.iter().enumerate().filter(|(_, e)| e.needed_for_output).min_by_key(|(_, e)| e.poc).map(|(i, _)| i) else { + return false; + }; + let e = &mut self.entries[i]; + e.needed_for_output = false; + out.push(Output { frame: e.frame.clone(), poc: e.poc, meta: e.meta.clone() }); + self.prune(); + true + } + + /// Output everything in POC order and clear non-reference entries. + pub fn flush(&mut self, out: &mut Vec) { + while self.bump(out) {} + self.prune(); + } + + /// IDR: mark all references unused; output (or drop) prior pictures. + pub fn idr(&mut self, no_output_of_prior_pics: bool, out: &mut Vec) { + for e in &mut self.entries { + e.mark = RefMark::Unused; + } + if no_output_of_prior_pics { + self.entries.clear(); + } else { + self.flush(out); + } + self.max_long_term_frame_idx = None; + } + + /// Sliding window marking (8.2.5.3). + fn sliding_window(&mut self, cur_frame_num: u32, max_frame_num: u32, max_num_ref_frames: usize) { + let limit = max_num_ref_frames.max(1); + while self.num_refs() >= limit { + let victim = self + .entries + .iter() + .enumerate() + .filter(|(_, e)| e.mark == RefMark::Short) + .min_by_key(|(_, e)| Self::frame_num_wrap(e, cur_frame_num, max_frame_num)) + .map(|(i, _)| i); + match victim { + Some(i) => self.entries[i].mark = RefMark::Unused, + None => break, + } + } + self.prune(); + } + + /// Insert "non-existing" frames for a frame_num gap (8.2.5.2). + #[allow(clippy::too_many_arguments)] + pub fn fill_frame_num_gap( + &mut self, + prev_ref_frame_num: u32, + frame_num: u32, + max_frame_num: u32, + max_num_ref_frames: usize, + make_frame: &mut dyn FnMut(u32) -> (FrameRef, Arc), + poc_state: &mut dyn FnMut(u32), + ) { + let mut unused = (prev_ref_frame_num + 1) % max_frame_num; + let mut guard = 0; + while unused != frame_num && guard < max_frame_num { + self.sliding_window(unused, max_frame_num, max_num_ref_frames); + let (frame, meta) = make_frame(unused); + poc_state(unused); + self.entries.push(DpbEntry { + poc: frame.poc, + frame, + frame_num: unused, + long_term_frame_idx: 0, + mark: RefMark::Short, + needed_for_output: false, + non_existing: true, + meta, + }); + unused = (unused + 1) % max_frame_num; + guard += 1; + } + } + + /// Initial reference picture lists (8.2.4.2) + modification (8.2.4.3) for a frame slice. + pub fn build_ref_lists(&self, sh: &SliceHeader, cur_poc: i32, max_frame_num: u32) -> Result<[Vec; 2]> { + let mut lists: [Vec; 2] = [Vec::new(), Vec::new()]; + if sh.slice_type.is_intra() { + return Ok(lists); + } + let cur_fn = sh.frame_num; + let pic_num = |e: &DpbEntry| Self::frame_num_wrap(e, cur_fn, max_frame_num); + let short: Vec<&DpbEntry> = self.entries.iter().filter(|e| e.mark == RefMark::Short).collect(); + let mut long: Vec<&DpbEntry> = self.entries.iter().filter(|e| e.mark == RefMark::Long).collect(); + long.sort_by_key(|e| e.long_term_frame_idx); + let to_ref = |e: &&DpbEntry| RefPic { frame: e.frame.clone(), long_term: e.mark == RefMark::Long }; + let mut entry_lists: [Vec<&DpbEntry>; 2] = [Vec::new(), Vec::new()]; + if sh.slice_type != SliceType::B { + let mut s = short.clone(); + s.sort_by_key(|e| std::cmp::Reverse(pic_num(e))); + entry_lists[0] = s.into_iter().chain(long.iter().copied()).collect(); + } else { + let mut before: Vec<&DpbEntry> = short.iter().copied().filter(|e| e.poc < cur_poc).collect(); + let mut after: Vec<&DpbEntry> = short.iter().copied().filter(|e| e.poc > cur_poc).collect(); + before.sort_by_key(|e| std::cmp::Reverse(e.poc)); + after.sort_by_key(|e| e.poc); + entry_lists[0] = before.iter().chain(after.iter()).chain(long.iter()).copied().collect(); + entry_lists[1] = after.iter().chain(before.iter()).chain(long.iter()).copied().collect(); + if entry_lists[1].len() > 1 && entry_lists[0].len() == entry_lists[1].len() { + let same = entry_lists[0].iter().zip(entry_lists[1].iter()).all(|(a, b)| std::ptr::eq(*a, *b)); + if same { + entry_lists[1].swap(0, 1); + } + } + } + let nlists = if sh.slice_type == SliceType::B { 2 } else { 1 }; + for l in 0..nlists { + let n = sh.num_ref_idx_active[l] as usize; + let mut list: Vec> = entry_lists[l].iter().copied().map(Some).collect(); + list.truncate(n); + list.resize(n, None); + if !sh.ref_pic_list_mod[l].is_empty() { + let max_pic_num = max_frame_num as i32; + let curr_pic_num = cur_fn as i32; + let mut pred = curr_pic_num; + let mut ref_idx = 0usize; + for m in &sh.ref_pic_list_mod[l] { + ensure!(ref_idx < n, "too many ref_pic_list_modification entries"); + list.push(None); // temporarily one longer + let (target, is_long): (Option<&DpbEntry>, bool) = match m.idc { + 0 | 1 => { + let d = m.value as i32 + 1; + let no_wrap = if m.idc == 0 { + if pred - d < 0 { pred - d + max_pic_num } else { pred - d } + } else if pred + d >= max_pic_num { + pred + d - max_pic_num + } else { + pred + d + }; + pred = no_wrap; + let pn = if no_wrap > curr_pic_num { no_wrap - max_pic_num } else { no_wrap }; + (short.iter().copied().find(|e| pic_num(e) == pn), false) + } + 2 => (long.iter().copied().find(|e| e.long_term_frame_idx == m.value), true), + _ => continue, + }; + let Some(target) = target else { + // Missing picture: keep the list unchanged (error concealment). + list.pop(); + ref_idx += 1; + continue; + }; + for c in (ref_idx + 1..=n).rev() { + list[c] = list[c - 1]; + } + list[ref_idx] = Some(target); + ref_idx += 1; + let mut nidx = ref_idx; + for c in ref_idx..=n { + let keep = match list[c] { + None => true, + Some(e) => { + if is_long { + !(e.mark == RefMark::Long && e.long_term_frame_idx == target.long_term_frame_idx) + } else { + !(e.mark == RefMark::Short && std::ptr::eq(e, target)) + } + } + }; + if keep { + list[nidx] = list[c]; + nidx += 1; + } + } + list.truncate(n); + } + } + // Fill missing entries (should not happen in conforming streams) with the first valid entry. + let fallback = list.iter().flatten().next().copied().or_else(|| self.entries.last()); + let out: Vec = list.iter().filter_map(|e| e.or(fallback)).map(|e| to_ref(&e)).collect(); + lists[l] = out; + } + Ok(lists) + } + + /// Adaptive memory control (8.2.5.4). Returns true if MMCO 5 was executed. + fn apply_mmcos(&mut self, mmcos: &[Mmco], cur_frame_num: u32, max_frame_num: u32, cur_long: &mut Option) -> bool { + let mut had5 = false; + let curr_pic_num = cur_frame_num as i32; + for m in mmcos { + match m.op { + 1 => { + let pn = curr_pic_num - (m.difference_of_pic_nums_minus1 as i32 + 1); + for e in &mut self.entries { + if e.mark == RefMark::Short && Self::frame_num_wrap(e, cur_frame_num, max_frame_num) == pn { + e.mark = RefMark::Unused; + } + } + } + 2 => { + for e in &mut self.entries { + if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_pic_num { + e.mark = RefMark::Unused; + } + } + } + 3 => { + let pn = curr_pic_num - (m.difference_of_pic_nums_minus1 as i32 + 1); + for e in &mut self.entries { + if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_frame_idx { + e.mark = RefMark::Unused; + } + } + for e in &mut self.entries { + if e.mark == RefMark::Short && Self::frame_num_wrap(e, cur_frame_num, max_frame_num) == pn { + e.mark = RefMark::Long; + e.long_term_frame_idx = m.long_term_frame_idx; + } + } + } + 4 => { + let max = m.max_long_term_frame_idx_plus1; + for e in &mut self.entries { + if e.mark == RefMark::Long && (max == 0 || e.long_term_frame_idx > max - 1) { + e.mark = RefMark::Unused; + } + } + self.max_long_term_frame_idx = if max == 0 { None } else { Some(max - 1) }; + } + 5 => { + for e in &mut self.entries { + e.mark = RefMark::Unused; + } + self.max_long_term_frame_idx = None; + had5 = true; + } + 6 => { + for e in &mut self.entries { + if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_frame_idx { + e.mark = RefMark::Unused; + } + } + *cur_long = Some(m.long_term_frame_idx); + } + _ => {} + } + } + self.prune(); + had5 + } + + /// Reference marking for the just-decoded picture and insertion into the DPB with output bumping. + /// Returns pictures ready for output. + #[allow(clippy::too_many_arguments)] + pub fn store_picture( + &mut self, + sh: &SliceHeader, + frame: FrameRef, + poc: i32, + max_frame_num: u32, + max_num_ref_frames: usize, + meta: Arc, + out: &mut Vec, + ) { + let mut mark = RefMark::Unused; + let mut long_idx = 0; + let mut had5 = false; + if sh.nal_ref_idc != 0 { + if sh.idr { + if sh.long_term_reference { + mark = RefMark::Long; + self.max_long_term_frame_idx = Some(0); + } else { + mark = RefMark::Short; + self.max_long_term_frame_idx = None; + } + } else { + let mut cur_long = None; + if sh.adaptive_ref_pic_marking { + had5 = self.apply_mmcos(&sh.mmcos, sh.frame_num, max_frame_num, &mut cur_long); + } else { + self.sliding_window(sh.frame_num, max_frame_num, max_num_ref_frames); + } + if let Some(i) = cur_long { + mark = RefMark::Long; + long_idx = i; + } else { + mark = RefMark::Short; + // make room if MMCOs left the DPB over-full (non-conforming streams) + if self.num_refs() >= max_num_ref_frames.max(1) && sh.adaptive_ref_pic_marking { + self.sliding_window(sh.frame_num, max_frame_num, max_num_ref_frames); + } + } + } + } + if had5 { + // All earlier pictures are output before the current one (C.4.4). + self.flush(out); + } + let frame_num = if had5 { 0 } else { sh.frame_num }; + // Bumping until there is space (C.4.5.1 / C.4.5.2). + if mark == RefMark::Unused { + loop { + let waiting = self.entries.len(); + if waiting < self.capacity { + break; + } + // non-reference picture with the smallest POC is output directly + let min_poc = self.entries.iter().filter(|e| e.needed_for_output).map(|e| e.poc).min(); + if min_poc.is_none_or(|m| poc < m) { + out.push(Output { frame, poc, meta }); + return; + } + if !self.bump(out) { + break; + } + } + } else { + while self.entries.len() >= self.capacity { + if !self.bump(out) { + break; + } + } + } + self.entries.push(DpbEntry { frame, poc, frame_num, long_term_frame_idx: long_idx, mark, needed_for_output: true, non_existing: false, meta }); + // Low-delay output: honour max_num_reorder_frames. + while self.entries.iter().filter(|e| e.needed_for_output).count() > self.max_reorder { + if !self.bump(out) { + break; + } + } + } +} + +impl Default for Dpb { + fn default() -> Self { + Self::new() + } +} diff --git a/crates/h264/src/inter.rs b/crates/h264/src/inter.rs new file mode 100644 index 00000000..fb73f2d9 --- /dev/null +++ b/crates/h264/src/inter.rs @@ -0,0 +1,303 @@ +//! Inter prediction sample interpolation (8.4.2.2) and weighted sample prediction (8.4.2.3). + +/// A reference plane. +#[derive(Clone, Copy)] +pub struct PlaneRef<'a> { + pub data: &'a [u8], + pub width: usize, + pub height: usize, + pub stride: usize, +} + +const WIN: usize = 16 + 5; + +/// Copy the (bw+5)x(bh+5) window whose top-left is (x-2, y-2) into `win` (stride WIN), clamping +/// coordinates to the picture. +#[inline] +fn fetch_window(p: PlaneRef, x: i32, y: i32, bw: usize, bh: usize, win: &mut [u8; WIN * WIN]) { + let x0 = x - 2; + let y0 = y - 2; + let ww = bw + 5; + let wh = bh + 5; + let maxx = p.width as i32 - 1; + let maxy = p.height as i32 - 1; + if x0 >= 0 && y0 >= 0 && x0 + ww as i32 - 1 <= maxx && y0 + wh as i32 - 1 <= maxy { + for r in 0..wh { + let src = (y0 as usize + r) * p.stride + x0 as usize; + win[r * WIN..r * WIN + ww].copy_from_slice(&p.data[src..src + ww]); + } + } else { + for r in 0..wh { + let yy = (y0 + r as i32).clamp(0, maxy) as usize; + let row = &p.data[yy * p.stride..yy * p.stride + p.width]; + for c in 0..ww { + let xx = (x0 + c as i32).clamp(0, maxx) as usize; + win[r * WIN + c] = row[xx]; + } + } + } +} + +#[inline(always)] +fn tap6(a: i32, b: i32, c: i32, d: i32, e: i32, f: i32) -> i32 { + a - 5 * b + 20 * c + 20 * d - 5 * e + f +} + +#[inline(always)] +fn clip1(v: i32) -> u8 { + v.clamp(0, 255) as u8 +} + +/// Luma sample interpolation for a bw x bh block. (x, y) is the integer sample position +/// (block position + (mv >> 2)); (fx, fy) the quarter-sample fraction. Output stride is `os`. +#[allow(clippy::too_many_arguments)] +pub fn mc_luma(p: PlaneRef, x: i32, y: i32, fx: u32, fy: u32, bw: usize, bh: usize, out: &mut [u8], os: usize) { + let mut win = [0u8; WIN * WIN]; + fetch_window(p, x, y, bw, bh, &mut win); + // sample G at (r, c) in block coordinates is win[(r + 2) * WIN + c + 2] + let g = |r: usize, c: usize| win[(r + 2) * WIN + c + 2] as i32; + // unclipped horizontal half-sample b1 at block row r (-2..bh+3 via offset) and column c + let b1 = |r: isize, c: usize| { + let base = ((r + 2) as usize) * WIN + c; + tap6(win[base] as i32, win[base + 1] as i32, win[base + 2] as i32, win[base + 3] as i32, win[base + 4] as i32, win[base + 5] as i32) + }; + let h1 = |r: usize, c: isize| { + let col = (c + 2) as usize; + tap6( + win[r * WIN + col] as i32, + win[(r + 1) * WIN + col] as i32, + win[(r + 2) * WIN + col] as i32, + win[(r + 3) * WIN + col] as i32, + win[(r + 4) * WIN + col] as i32, + win[(r + 5) * WIN + col] as i32, + ) + }; + match (fx, fy) { + (0, 0) => { + for r in 0..bh { + out[r * os..r * os + bw].copy_from_slice(&win[(r + 2) * WIN + 2..(r + 2) * WIN + 2 + bw]); + } + } + (_, 0) => { + // a, b, c + for r in 0..bh { + for c in 0..bw { + let b = clip1((b1(r as isize, c) + 16) >> 5) as i32; + out[r * os + c] = match fx { + 1 => ((g(r, c) + b + 1) >> 1) as u8, + 2 => b as u8, + _ => ((g(r, c + 1) + b + 1) >> 1) as u8, + }; + } + } + } + (0, _) => { + // d, h, n + for r in 0..bh { + for c in 0..bw { + let h = clip1((h1(r, c as isize) + 16) >> 5) as i32; + out[r * os + c] = match fy { + 1 => ((g(r, c) + h + 1) >> 1) as u8, + 2 => h as u8, + _ => ((g(r + 1, c) + h + 1) >> 1) as u8, + }; + } + } + } + (2, _) | (_, 2) => { + // j-based: j needs b1 for rows -2..bh+3 + let mut bcol = [0i32; WIN * WIN]; + for r in 0..bh + 5 { + for c in 0..bw { + bcol[r * WIN + c] = b1(r as isize - 2, c); + } + } + for r in 0..bh { + for c in 0..bw { + let j1 = tap6( + bcol[r * WIN + c], + bcol[(r + 1) * WIN + c], + bcol[(r + 2) * WIN + c], + bcol[(r + 3) * WIN + c], + bcol[(r + 4) * WIN + c], + bcol[(r + 5) * WIN + c], + ); + let j = clip1((j1 + 512) >> 10) as i32; + let v = match (fx, fy) { + (2, 2) => j, + (2, 1) => (clip1((bcol[(r + 2) * WIN + c] + 16) >> 5) as i32 + j + 1) >> 1, // f + (2, 3) => (clip1((bcol[(r + 3) * WIN + c] + 16) >> 5) as i32 + j + 1) >> 1, // q + (1, 2) => (clip1((h1(r, c as isize) + 16) >> 5) as i32 + j + 1) >> 1, // i + _ => (clip1((h1(r, c as isize + 1) + 16) >> 5) as i32 + j + 1) >> 1, // k + }; + out[r * os + c] = v as u8; + } + } + } + _ => { + // e, g, p, r: average of a horizontal half-sample (b or s) and a vertical one (h or m) + for r in 0..bh { + for c in 0..bw { + let hr = if fy == 1 { r as isize } else { r as isize + 1 }; + let vc = if fx == 1 { c as isize } else { c as isize + 1 }; + let bh_ = clip1((b1(hr, c) + 16) >> 5) as i32; + let vv = clip1((h1(r, vc) + 16) >> 5) as i32; + out[r * os + c] = ((bh_ + vv + 1) >> 1) as u8; + } + } + } + } +} + +/// Chroma sample interpolation (8.4.2.2.2) for a bw x bh block at integer position (x, y) with +/// eighth-sample fraction (fx, fy). +#[allow(clippy::too_many_arguments)] +pub fn mc_chroma(p: PlaneRef, x: i32, y: i32, fx: u32, fy: u32, bw: usize, bh: usize, out: &mut [u8], os: usize) { + let maxx = p.width as i32 - 1; + let maxy = p.height as i32 - 1; + let fx = fx as i32; + let fy = fy as i32; + let w00 = (8 - fx) * (8 - fy); + let w10 = fx * (8 - fy); + let w01 = (8 - fx) * fy; + let w11 = fx * fy; + let inside = x >= 0 && y >= 0 && x + bw as i32 <= maxx && y + bh as i32 <= maxy; + if inside { + let (x, y) = (x as usize, y as usize); + for r in 0..bh { + let r0 = &p.data[(y + r) * p.stride + x..]; + let r1 = &p.data[(y + r + 1) * p.stride + x..]; + for c in 0..bw { + let v = w00 * r0[c] as i32 + w10 * r0[c + 1] as i32 + w01 * r1[c] as i32 + w11 * r1[c + 1] as i32; + out[r * os + c] = ((v + 32) >> 6) as u8; + } + } + } else { + let at = |xx: i32, yy: i32| p.data[yy.clamp(0, maxy) as usize * p.stride + xx.clamp(0, maxx) as usize] as i32; + for r in 0..bh as i32 { + for c in 0..bw as i32 { + let v = w00 * at(x + c, y + r) + w10 * at(x + c + 1, y + r) + w01 * at(x + c, y + r + 1) + w11 * at(x + c + 1, y + r + 1); + out[(r as usize) * os + c as usize] = ((v + 32) >> 6) as u8; + } + } + } +} + +/// Weights for one prediction direction or pair. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Weight { + /// Default: copy (single list) or rounded average (bi). + Default, + /// Explicit/implicit weighting: log2 denominator, weights and offsets for L0 and L1. + Weighted { log_wd: i32, w0: i32, w1: i32, o0: i32, o1: i32 }, +} + +/// Combine prediction(s) into `dst` (stride `ds`). `p0`/`p1` have stride `ps`. +#[allow(clippy::too_many_arguments)] +pub fn weighted_store(dst: &mut [u8], ds: usize, p0: Option<&[u8]>, p1: Option<&[u8]>, ps: usize, bw: usize, bh: usize, w: Weight) { + match (p0, p1, w) { + (Some(a), Some(b), Weight::Default) => { + for r in 0..bh { + let d = &mut dst[r * ds..r * ds + bw]; + let ra = &a[r * ps..r * ps + bw]; + let rb = &b[r * ps..r * ps + bw]; + for ((d, &x), &y) in d.iter_mut().zip(ra).zip(rb) { + *d = ((x as u32 + y as u32 + 1) >> 1) as u8; + } + } + } + (Some(a), Some(b), Weight::Weighted { log_wd, w0, w1, o0, o1 }) => { + let round = 1 << log_wd; + let off = (o0 + o1 + 1) >> 1; + for r in 0..bh { + for c in 0..bw { + let v = ((a[r * ps + c] as i32 * w0 + b[r * ps + c] as i32 * w1 + round) >> (log_wd + 1)) + off; + dst[r * ds + c] = clip1(v); + } + } + } + (Some(a), None, w) | (None, Some(a), w) => { + let (wt, o) = match (w, p0.is_some()) { + (Weight::Default, _) => { + for r in 0..bh { + dst[r * ds..r * ds + bw].copy_from_slice(&a[r * ps..r * ps + bw]); + } + return; + } + (Weight::Weighted { w0, o0, .. }, true) => (w0, o0), + (Weight::Weighted { w1, o1, .. }, false) => (w1, o1), + }; + let log_wd = match w { + Weight::Weighted { log_wd, .. } => log_wd, + Weight::Default => 0, + }; + for r in 0..bh { + for c in 0..bw { + let s = a[r * ps + c] as i32; + let v = if log_wd >= 1 { ((s * wt + (1 << (log_wd - 1))) >> log_wd) + o } else { s * wt + o }; + dst[r * ds + c] = clip1(v); + } + } + } + (None, None, _) => {} + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn plane(w: usize, h: usize, f: impl Fn(usize, usize) -> u8) -> Vec { + (0..w * h).map(|i| f(i % w, i / w)).collect() + } + + #[test] + fn full_sample_copy_and_clamp() { + let d = plane(16, 16, |x, y| (x + 16 * y) as u8); + let p = PlaneRef { data: &d, width: 16, height: 16, stride: 16 }; + let mut out = [0u8; 16]; + mc_luma(p, 2, 3, 0, 0, 4, 4, &mut out, 4); + assert_eq!(&out[..4], &[50, 51, 52, 53]); + // clamped: far outside to the top-left gives sample (0,0) + mc_luma(p, -40, -40, 0, 0, 4, 4, &mut out, 4); + assert!(out.iter().all(|&v| v == 0)); + } + + #[test] + fn half_sample_constant_is_constant() { + let d = vec![100u8; 32 * 32]; + let p = PlaneRef { data: &d, width: 32, height: 32, stride: 32 }; + for fx in 0..4 { + for fy in 0..4 { + let mut out = [0u8; 64]; + mc_luma(p, 8, 8, fx, fy, 8, 8, &mut out, 8); + assert!(out.iter().all(|&v| v == 100), "{fx},{fy}"); + mc_chroma(p, 8, 8, fx * 2, fy * 2, 8, 8, &mut out, 8); + assert!(out.iter().all(|&v| v == 100)); + } + } + } + + #[test] + fn half_sample_horizontal_ramp() { + // linear ramp: 6-tap filter of a linear function is exact: b = (G + H) / 2 (rounded) + let d = plane(32, 32, |x, _| (x * 4) as u8); + let p = PlaneRef { data: &d, width: 32, height: 32, stride: 32 }; + let mut out = [0u8; 16]; + mc_luma(p, 8, 8, 2, 0, 4, 4, &mut out, 4); + assert_eq!(&out[..4], &[34, 38, 42, 46]); + mc_luma(p, 8, 8, 1, 0, 4, 4, &mut out, 4); + assert_eq!(&out[..4], &[33, 37, 41, 45]); + } + + #[test] + fn weighting() { + let a = [100u8; 4]; + let b = [50u8; 4]; + let mut d = [0u8; 4]; + weighted_store(&mut d, 4, Some(&a), Some(&b), 4, 4, 1, Weight::Default); + assert_eq!(d, [75; 4]); + weighted_store(&mut d, 4, Some(&a), None, 4, 4, 1, Weight::Weighted { log_wd: 5, w0: 16, w1: 0, o0: 3, o1: 0 }); + assert_eq!(d, [53; 4]); + } +} diff --git a/crates/h264/src/intra.rs b/crates/h264/src/intra.rs new file mode 100644 index 00000000..decf2965 --- /dev/null +++ b/crates/h264/src/intra.rs @@ -0,0 +1,462 @@ +//! Intra prediction (8.3): Intra_4x4, Intra_8x8 (with reference filtering), Intra_16x16 and chroma. + +/// Neighbouring sample availability for one block. +#[derive(Clone, Copy, Debug, Default)] +pub struct Avail { + pub left: bool, + pub top: bool, + pub top_left: bool, + pub top_right: bool, +} + +/// Neighbouring samples of an NxN block: p[x,-1] for x = 0..2N-1, p[-1,y] for y = 0..N-1, p[-1,-1]. +#[derive(Clone, Copy)] +struct Edge { + top: [i32; 16], + left: [i32; 8], + tl: i32, +} + +impl Edge { + /// p[x, -1] for x >= -1 + #[inline(always)] + fn t(&self, x: i32) -> i32 { + if x < 0 { self.tl } else { self.top[x as usize] } + } + /// p[-1, y] for y >= -1 + #[inline(always)] + fn l(&self, y: i32) -> i32 { + if y < 0 { self.tl } else { self.left[y as usize] } + } +} + +fn load_edge(plane: &[u8], stride: usize, x: usize, y: usize, n: usize, a: Avail) -> Edge { + let mut e = Edge { top: [128; 16], left: [128; 8], tl: 128 }; + if a.top { + let row = &plane[(y - 1) * stride + x..]; + for i in 0..n { + e.top[i] = row[i] as i32; + } + if a.top_right { + for i in n..2 * n { + e.top[i] = row[i] as i32; + } + } else { + let v = e.top[n - 1]; + for t in &mut e.top[n..2 * n] { + *t = v; + } + } + } + if a.left { + for i in 0..n { + e.left[i] = plane[(y + i) * stride + x - 1] as i32; + } + } + if a.top_left { + e.tl = plane[(y - 1) * stride + x - 1] as i32; + } + e +} + +/// Directional / DC prediction shared by Intra_4x4 and Intra_8x8 (8.3.1.2.x, 8.3.2.2.x). +fn pred_nxn(mode: u8, e: &Edge, n: i32, a: Avail, out: &mut [u8], stride: usize) { + let mut put = |x: i32, y: i32, v: i32| out[y as usize * stride + x as usize] = v as u8; + match mode { + 0 => { + for y in 0..n { + for x in 0..n { + put(x, y, e.t(x)); + } + } + } + 1 => { + for y in 0..n { + for x in 0..n { + put(x, y, e.l(y)); + } + } + } + 2 => { + let shift = if n == 4 { 2 } else { 3 }; + let st: i32 = (0..n).map(|x| e.t(x)).sum(); + let sl: i32 = (0..n).map(|y| e.l(y)).sum(); + let dc = match (a.top, a.left) { + (true, true) => (st + sl + n) >> (shift + 1), + (false, true) => (sl + (n >> 1)) >> shift, + (true, false) => (st + (n >> 1)) >> shift, + _ => 128, + }; + for y in 0..n { + for x in 0..n { + put(x, y, dc); + } + } + } + 3 => { + for y in 0..n { + for x in 0..n { + let v = if x == n - 1 && y == n - 1 { + (e.t(2 * n - 2) + 3 * e.t(2 * n - 1) + 2) >> 2 + } else { + (e.t(x + y) + 2 * e.t(x + y + 1) + e.t(x + y + 2) + 2) >> 2 + }; + put(x, y, v); + } + } + } + 4 => { + for y in 0..n { + for x in 0..n { + let v = if x > y { + (e.t(x - y - 2) + 2 * e.t(x - y - 1) + e.t(x - y) + 2) >> 2 + } else if x < y { + (e.l(y - x - 2) + 2 * e.l(y - x - 1) + e.l(y - x) + 2) >> 2 + } else { + (e.t(0) + 2 * e.tl + e.l(0) + 2) >> 2 + }; + put(x, y, v); + } + } + } + 5 => { + for y in 0..n { + for x in 0..n { + let z = 2 * x - y; + let v = if z >= 0 && z % 2 == 0 { + (e.t(x - (y >> 1) - 1) + e.t(x - (y >> 1)) + 1) >> 1 + } else if z > 0 { + (e.t(x - (y >> 1) - 2) + 2 * e.t(x - (y >> 1) - 1) + e.t(x - (y >> 1)) + 2) >> 2 + } else if z == -1 { + (e.l(0) + 2 * e.tl + e.t(0) + 2) >> 2 + } else { + (e.l(y - 2 * x - 1) + 2 * e.l(y - 2 * x - 2) + e.l(y - 2 * x - 3) + 2) >> 2 + }; + put(x, y, v); + } + } + } + 6 => { + for y in 0..n { + for x in 0..n { + let z = 2 * y - x; + let v = if z >= 0 && z % 2 == 0 { + (e.l(y - (x >> 1) - 1) + e.l(y - (x >> 1)) + 1) >> 1 + } else if z > 0 { + (e.l(y - (x >> 1) - 2) + 2 * e.l(y - (x >> 1) - 1) + e.l(y - (x >> 1)) + 2) >> 2 + } else if z == -1 { + (e.l(0) + 2 * e.tl + e.t(0) + 2) >> 2 + } else { + (e.t(x - 2 * y - 1) + 2 * e.t(x - 2 * y - 2) + e.t(x - 2 * y - 3) + 2) >> 2 + }; + put(x, y, v); + } + } + } + 7 => { + for y in 0..n { + for x in 0..n { + let i = x + (y >> 1); + let v = if y % 2 == 0 { (e.t(i) + e.t(i + 1) + 1) >> 1 } else { (e.t(i) + 2 * e.t(i + 1) + e.t(i + 2) + 2) >> 2 }; + put(x, y, v); + } + } + } + _ => { + let zmax = 2 * n - 3; + for y in 0..n { + for x in 0..n { + let z = x + 2 * y; + let i = y + (x >> 1); + let v = if z < zmax && z % 2 == 0 { + (e.l(i) + e.l(i + 1) + 1) >> 1 + } else if z < zmax { + (e.l(i) + 2 * e.l(i + 1) + e.l(i + 2) + 2) >> 2 + } else if z == zmax { + (e.l(n - 2) + 3 * e.l(n - 1) + 2) >> 2 + } else { + e.l(n - 1) + }; + put(x, y, v); + } + } + } + } +} + +/// Intra_4x4 prediction written into `plane` at (x, y). +pub fn pred4x4(plane: &mut [u8], stride: usize, x: usize, y: usize, mode: u8, a: Avail) { + let e = load_edge(plane, stride, x, y, 4, a); + pred_nxn(mode, &e, 4, a, &mut plane[y * stride + x..], stride); +} + +/// Intra_8x8 prediction (with reference sample filtering) written into `plane` at (x, y). +pub fn pred8x8(plane: &mut [u8], stride: usize, x: usize, y: usize, mode: u8, a: Avail) { + let p = load_edge(plane, stride, x, y, 8, a); + let mut f = p; + if a.top { + f.top[0] = if a.top_left { (p.tl + 2 * p.top[0] + p.top[1] + 2) >> 2 } else { (3 * p.top[0] + p.top[1] + 2) >> 2 }; + for i in 1..15 { + f.top[i] = (p.top[i - 1] + 2 * p.top[i] + p.top[i + 1] + 2) >> 2; + } + f.top[15] = (p.top[14] + 3 * p.top[15] + 2) >> 2; + } + if a.top_left { + f.tl = match (a.top, a.left) { + (true, true) => (p.top[0] + 2 * p.tl + p.left[0] + 2) >> 2, + (true, false) => (3 * p.tl + p.top[0] + 2) >> 2, + (false, true) => (3 * p.tl + p.left[0] + 2) >> 2, + _ => p.tl, + }; + } + if a.left { + f.left[0] = if a.top_left { (p.tl + 2 * p.left[0] + p.left[1] + 2) >> 2 } else { (3 * p.left[0] + p.left[1] + 2) >> 2 }; + for i in 1..7 { + f.left[i] = (p.left[i - 1] + 2 * p.left[i] + p.left[i + 1] + 2) >> 2; + } + f.left[7] = (p.left[6] + 3 * p.left[7] + 2) >> 2; + } + pred_nxn(mode, &f, 8, a, &mut plane[y * stride + x..], stride); +} + +#[inline(always)] +fn clip(v: i32) -> u8 { + v.clamp(0, 255) as u8 +} + +/// Intra_16x16 prediction (8.3.3) for the macroblock at luma (x, y). +pub fn pred16x16(plane: &mut [u8], stride: usize, x: usize, y: usize, mode: u8, a: Avail) { + let mut top = [0i32; 16]; + let mut left = [0i32; 16]; + if a.top { + for i in 0..16 { + top[i] = plane[(y - 1) * stride + x + i] as i32; + } + } + if a.left { + for i in 0..16 { + left[i] = plane[(y + i) * stride + x - 1] as i32; + } + } + let base = y * stride + x; + match mode { + 0 => { + for r in 0..16 { + for c in 0..16 { + plane[base + r * stride + c] = top[c] as u8; + } + } + } + 1 => { + for r in 0..16 { + plane[base + r * stride..base + r * stride + 16].fill(left[r] as u8); + } + } + 2 => { + let st: i32 = top.iter().sum(); + let sl: i32 = left.iter().sum(); + let dc = match (a.top, a.left) { + (true, true) => (st + sl + 16) >> 5, + (false, true) => (sl + 8) >> 4, + (true, false) => (st + 8) >> 4, + _ => 128, + } as u8; + for r in 0..16 { + plane[base + r * stride..base + r * stride + 16].fill(dc); + } + } + _ => { + let tl = if a.top_left { plane[(y - 1) * stride + x - 1] as i32 } else { 128 }; + let t = |i: i32| if i < 0 { tl } else { top[i as usize] }; + let l = |i: i32| if i < 0 { tl } else { left[i as usize] }; + let mut h = 0; + let mut v = 0; + for k in 0..8 { + h += (k + 1) * (t(8 + k) - t(6 - k)); + v += (k + 1) * (l(8 + k) - l(6 - k)); + } + let aa = 16 * (left[15] + top[15]); + let b = (5 * h + 32) >> 6; + let c = (5 * v + 32) >> 6; + for r in 0..16 { + for cc in 0..16 { + plane[base + r * stride + cc] = clip((aa + b * (cc as i32 - 7) + c * (r as i32 - 7) + 16) >> 5); + } + } + } + } +} + +/// Chroma intra prediction for a 4:2:0 macroblock (8x8 block at chroma (x, y)), 8.3.4. +/// Modes: 0 DC, 1 horizontal, 2 vertical, 3 plane. +pub fn pred_chroma(plane: &mut [u8], stride: usize, x: usize, y: usize, mode: u8, a: Avail) { + const W: usize = 8; + const H: usize = 8; + let mut top = [0i32; W]; + let mut left = [0i32; H]; + if a.top { + for i in 0..W { + top[i] = plane[(y - 1) * stride + x + i] as i32; + } + } + if a.left { + for i in 0..H { + left[i] = plane[(y + i) * stride + x - 1] as i32; + } + } + let base = y * stride + x; + match mode { + 0 => { + for by in 0..H / 4 { + for bx in 0..W / 4 { + let st: i32 = top[bx * 4..bx * 4 + 4].iter().sum(); + let sl: i32 = left[by * 4..by * 4 + 4].iter().sum(); + let dc = if (bx == 0 && by == 0) || (bx > 0 && by > 0) { + match (a.top, a.left) { + (true, true) => (st + sl + 4) >> 3, + (false, true) => (sl + 2) >> 2, + (true, false) => (st + 2) >> 2, + _ => 128, + } + } else if bx > 0 { + if a.top { + (st + 2) >> 2 + } else if a.left { + (sl + 2) >> 2 + } else { + 128 + } + } else if a.left { + (sl + 2) >> 2 + } else if a.top { + (st + 2) >> 2 + } else { + 128 + }; + for r in 0..4 { + let o = base + (by * 4 + r) * stride + bx * 4; + plane[o..o + 4].fill(dc as u8); + } + } + } + } + 1 => { + for r in 0..H { + plane[base + r * stride..base + r * stride + W].fill(left[r] as u8); + } + } + 2 => { + for r in 0..H { + for c in 0..W { + plane[base + r * stride + c] = top[c] as u8; + } + } + } + _ => { + let tl = if a.top_left { plane[(y - 1) * stride + x - 1] as i32 } else { 128 }; + let t = |i: i32| if i < 0 { tl } else { top[i as usize] }; + let l = |i: i32| if i < 0 { tl } else { left[i as usize] }; + let mut h = 0; + let mut v = 0; + for k in 0..4 { + h += (k + 1) * (t(4 + k) - t(2 - k)); + v += (k + 1) * (l(4 + k) - l(2 - k)); + } + let aa = 16 * (left[H - 1] + top[W - 1]); + let b = (34 * h + 32) >> 6; + let c = (34 * v + 32) >> 6; + for r in 0..H { + for cc in 0..W { + plane[base + r * stride + cc] = clip((aa + b * (cc as i32 - 3) + c * (r as i32 - 3) + 16) >> 5); + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn setup() -> (Vec, usize) { + // 12x12 plane; block at (4,4); top row y=3 = 10,20,..; left column x=3 = 1,2,3,... + let stride = 16; + let mut p = vec![0u8; stride * 16]; + for i in 0..12 { + p[3 * stride + 4 + i] = (10 * (i + 1)) as u8; + } + for i in 0..8 { + p[(4 + i) * stride + 3] = (i + 1) as u8; + } + p[3 * stride + 3] = 5; + (p, stride) + } + const ALL: Avail = Avail { left: true, top: true, top_left: true, top_right: true }; + + #[test] + fn i4_vertical_horizontal_dc() { + let (mut p, s) = setup(); + pred4x4(&mut p, s, 4, 4, 0, ALL); + assert_eq!(&p[4 * s + 4..4 * s + 8], &[10, 20, 30, 40]); + assert_eq!(&p[7 * s + 4..7 * s + 8], &[10, 20, 30, 40]); + pred4x4(&mut p, s, 4, 4, 1, ALL); + assert_eq!(&p[5 * s + 4..5 * s + 8], &[2, 2, 2, 2]); + pred4x4(&mut p, s, 4, 4, 2, ALL); + // (10+20+30+40 + 1+2+3+4 + 4) >> 3 = 114 >> 3 = 14 + assert_eq!(p[4 * s + 4], 14); + } + + #[test] + fn i4_diag_down_left_hand_computed() { + let (mut p, s) = setup(); + pred4x4(&mut p, s, 4, 4, 3, ALL); + // pred[0,0] = (t0 + 2 t1 + t2 + 2) >> 2 = (10 + 40 + 30 + 2) >> 2 = 20 + assert_eq!(p[4 * s + 4], 20); + // pred[3,3] = (t6 + 3 t7 + 2) >> 2 = (70 + 240 + 2) >> 2 = 78 + assert_eq!(p[7 * s + 7], 78); + } + + #[test] + fn i4_diag_down_right_hand_computed() { + let (mut p, s) = setup(); + pred4x4(&mut p, s, 4, 4, 4, ALL); + // x==y: (t0 + 2 tl + l0 + 2) >> 2 = (10 + 10 + 1 + 2) >> 2 = 5 + assert_eq!(p[4 * s + 4], 5); + // x=1,y=0: (tl + 2 t0 + t1 + 2)>>2 = (5 + 20 + 20 + 2) >> 2 = 11 + assert_eq!(p[4 * s + 5], 11); + // x=0,y=1: (tl + 2 l0 + l1 + 2) >> 2 = (5 + 2 + 2 + 2) >> 2 = 2 + assert_eq!(p[5 * s + 4], 2); + } + + #[test] + fn i4_top_right_substitution() { + let (mut p, s) = setup(); + let a = Avail { top_right: false, ..ALL }; + pred4x4(&mut p, s, 4, 4, 3, a); + // top-right replaced by t3 = 40: pred[3,3] = (40 + 120 + 2) >> 2 = 40 + assert_eq!(p[7 * s + 7], 40); + } + + #[test] + fn i16_plane_flat_is_flat() { + let s = 32; + let mut p = vec![77u8; s * 32]; + pred16x16(&mut p, s, 8, 8, 3, ALL); + assert!(p[8 * s + 8..8 * s + 24].iter().all(|&v| v == 77)); + } + + #[test] + fn chroma_dc_rules() { + let s = 16; + let mut p = vec![0u8; s * 16]; + for i in 0..8 { + p[3 * s + 4 + i] = if i < 4 { 40 } else { 80 }; + p[(4 + i) * s + 3] = if i < 4 { 8 } else { 16 }; + } + pred_chroma(&mut p, s, 4, 4, 0, ALL); + assert_eq!(p[4 * s + 4], (160 + 32 + 4) >> 3); // top-left block: both + assert_eq!(p[4 * s + 8], 80); // top-right block: top only + assert_eq!(p[8 * s + 4], 16); // bottom-left: left only + assert_eq!(p[8 * s + 8], ((320 + 64 + 4) >> 3) as u8); // bottom-right: both + } +} diff --git a/crates/h264/src/lib.rs b/crates/h264/src/lib.rs index a316c54e..82764cb6 100644 --- a/crates/h264/src/lib.rs +++ b/crates/h264/src/lib.rs @@ -1,13 +1,75 @@ //! Clean-room pure-Rust H.264 / AVC decoder, implemented from ITU-T Rec. H.264 (ISO/IEC 14496-10). //! -//! M2.1: NAL unit, parameter set, slice header and picture order count parsing. +//! Supported: progressive (frame) coding, 8-bit 4:2:0, Baseline / Main / High profiles — CAVLC and CABAC, +//! I/P/B slices, 8x8 transform, custom scaling matrices, weighted prediction, spatial/temporal direct, +//! multiple slices, long-term references and all MMCOs, deblocking. +//! +//! ```no_run +//! let mut dec = filmcraft_h264::Decoder::new(); +//! let stream = std::fs::read("video.h264").unwrap(); +//! for pic in dec.decode(&stream, 0).unwrap() { +//! println!("{}x{} poc {}", pic.width, pic.height, pic.poc); +//! } +//! for pic in dec.flush() { +//! let _ = pic; +//! } +//! ``` -// Constant tables are shared with the macroblock decoding stages that follow. +// Index loops over fixed-size blocks read more clearly than iterator chains in codec code. +#![allow(clippy::needless_range_loop, clippy::too_many_arguments)] +// Deblocking tables and per-slice filter parameters are consumed by the M2.5 loop filter. #![allow(dead_code)] +mod cavlc; +#[rustfmt::skip] +mod cavlc_tables; +mod decoder; +mod dpb; mod error; +mod inter; +mod intra; +mod mbtypes; pub mod params; +mod picture; pub mod slice; +mod slicedec; mod tables; +mod transform; +pub use decoder::Decoder; pub use error::{Error, Result}; + +/// Colour description from the VUI (ITU-T H.273 code points). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ColorInfo { + /// `video_full_range_flag`. + pub full_range: bool, + pub primaries: u8, + pub transfer: u8, + pub matrix: u8, +} + +/// A decoded, cropped picture in 8-bit planar 4:2:0. +#[derive(Clone, Debug)] +pub struct Picture { + /// Cropped luma width. + pub width: u32, + /// Cropped luma height. + pub height: u32, + pub chroma_width: u32, + pub chroma_height: u32, + pub y: Vec, + pub u: Vec, + pub v: Vec, + pub y_stride: usize, + pub uv_stride: usize, + /// Presentation timestamp passed to [`Decoder::decode`] with the access unit of this picture. + pub pts: i64, + /// Picture order count. + pub poc: i32, + /// IDR picture. + pub key: bool, + pub color: ColorInfo, + /// Sample aspect ratio (0, 0 when unspecified). + pub sar: (u16, u16), +} diff --git a/crates/h264/src/mbtypes.rs b/crates/h264/src/mbtypes.rs new file mode 100644 index 00000000..87aa206b --- /dev/null +++ b/crates/h264/src/mbtypes.rs @@ -0,0 +1,235 @@ +//! Macroblock / sub-macroblock type tables (7.4.5, Tables 7-11 .. 7-18). + +use crate::picture::MbKind; + +pub const PRED_L0: u8 = 1; +pub const PRED_L1: u8 = 2; +pub const PRED_BI: u8 = 3; + +/// Macroblock partitioning for inter macroblocks. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum Part { + #[default] + P16x16, + P16x8, + P8x16, + P8x8, +} + +impl Part { + pub fn num_parts(self) -> usize { + match self { + Part::P16x16 => 1, + Part::P16x8 | Part::P8x16 => 2, + Part::P8x8 => 4, + } + } + /// (x, y, w, h) of partition `i` in 4x4-block units. + pub fn rect(self, i: usize) -> (usize, usize, usize, usize) { + match self { + Part::P16x16 => (0, 0, 4, 4), + Part::P16x8 => (0, 2 * i, 4, 2), + Part::P8x16 => (2 * i, 0, 2, 4), + Part::P8x8 => ((i & 1) * 2, (i >> 1) * 2, 2, 2), + } + } +} + +/// Sub-macroblock partitioning. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum SubPart { + #[default] + S8x8, + S8x4, + S4x8, + S4x4, +} + +impl SubPart { + pub fn num_parts(self) -> usize { + match self { + SubPart::S8x8 => 1, + SubPart::S8x4 | SubPart::S4x8 => 2, + SubPart::S4x4 => 4, + } + } + /// (x, y, w, h) of sub-partition `j` relative to the 8x8 block, in 4x4 units. + pub fn rect(self, j: usize) -> (usize, usize, usize, usize) { + match self { + SubPart::S8x8 => (0, 0, 2, 2), + SubPart::S8x4 => (0, j, 2, 1), + SubPart::S4x8 => (j, 0, 1, 2), + SubPart::S4x4 => (j & 1, j >> 1, 1, 1), + } + } +} + +/// Decoded meaning of mb_type. +#[derive(Clone, Copy, Debug, Default)] +pub struct MbTypeInfo { + pub kind: MbKind, + /// Intra16x16PredMode. + pub i16_mode: u8, + /// coded_block_pattern implied by an Intra_16x16 mb_type. + pub i16_cbp: u8, + pub part: Part, + /// Prediction list flags per macroblock partition. + pub pred: [u8; 2], + /// P_8x8ref0. + pub ref0: bool, +} + +/// Intra mb_type (I-slice numbering 0..=25). +pub fn intra_mb_type(t: u32) -> Option { + let mut m = MbTypeInfo::default(); + match t { + 0 => m.kind = MbKind::I4x4, // or I8x8, decided by transform_size_8x8_flag + 1..=24 => { + m.kind = MbKind::I16x16; + m.i16_mode = ((t - 1) % 4) as u8; + let chroma = ((t - 1) / 4) % 3; + let luma = if t >= 13 { 15 } else { 0 }; + m.i16_cbp = (luma | (chroma << 4)) as u8; + } + 25 => m.kind = MbKind::IPcm, + _ => return None, + } + Some(m) +} + +/// P-slice mb_type 0..=4 (5.. are intra, handled by the caller). +pub fn p_mb_type(t: u32) -> Option { + let mut m = MbTypeInfo { kind: MbKind::Inter, pred: [PRED_L0; 2], ..Default::default() }; + m.part = match t { + 0 => Part::P16x16, + 1 => Part::P16x8, + 2 => Part::P8x16, + 3 => Part::P8x8, + 4 => { + m.ref0 = true; + Part::P8x8 + } + _ => return None, + }; + Some(m) +} + +/// B-slice mb_type 0..=22 (23.. are intra). +pub fn b_mb_type(t: u32) -> Option { + let mut m = MbTypeInfo { kind: MbKind::Inter, ..Default::default() }; + match t { + 0 => { + m.kind = MbKind::BDirect16x16; + m.pred = [PRED_BI; 2]; + } + 1..=3 => { + m.part = Part::P16x16; + m.pred = [t as u8; 2]; + } + 4..=21 => { + const PAIRS: [(u8, u8); 9] = [ + (PRED_L0, PRED_L0), + (PRED_L1, PRED_L1), + (PRED_L0, PRED_L1), + (PRED_L1, PRED_L0), + (PRED_L0, PRED_BI), + (PRED_L1, PRED_BI), + (PRED_BI, PRED_L0), + (PRED_BI, PRED_L1), + (PRED_BI, PRED_BI), + ]; + let (a, b) = PAIRS[((t - 4) / 2) as usize]; + m.pred = [a, b]; + m.part = if t.is_multiple_of(2) { Part::P16x8 } else { Part::P8x16 }; + } + 22 => m.part = Part::P8x8, + _ => return None, + } + Some(m) +} + +/// Sub-macroblock type: (shape, prediction flags, direct). +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct SubMbInfo { + pub shape: SubPart, + pub pred: u8, + pub direct: bool, +} + +pub fn p_sub_mb_type(t: u32) -> Option { + let shape = match t { + 0 => SubPart::S8x8, + 1 => SubPart::S8x4, + 2 => SubPart::S4x8, + 3 => SubPart::S4x4, + _ => return None, + }; + Some(SubMbInfo { shape, pred: PRED_L0, direct: false }) +} + +pub fn b_sub_mb_type(t: u32) -> Option { + let (shape, pred, direct) = match t { + 0 => (SubPart::S8x8, PRED_BI, true), + 1 => (SubPart::S8x8, PRED_L0, false), + 2 => (SubPart::S8x8, PRED_L1, false), + 3 => (SubPart::S8x8, PRED_BI, false), + 4 => (SubPart::S8x4, PRED_L0, false), + 5 => (SubPart::S4x8, PRED_L0, false), + 6 => (SubPart::S8x4, PRED_L1, false), + 7 => (SubPart::S4x8, PRED_L1, false), + 8 => (SubPart::S8x4, PRED_BI, false), + 9 => (SubPart::S4x8, PRED_BI, false), + 10 => (SubPart::S4x4, PRED_L0, false), + 11 => (SubPart::S4x4, PRED_L1, false), + 12 => (SubPart::S4x4, PRED_BI, false), + _ => return None, + }; + Some(SubMbInfo { shape, pred, direct }) +} + +/// Table 9-4 (a): codeNum -> coded_block_pattern for (Intra_4x4/Intra_8x8, Inter), ChromaArrayType 1 or 2. +#[rustfmt::skip] +pub const CBP_ME: [(u8, u8); 48] = [ + (47, 0), (31, 16), (15, 1), (0, 2), (23, 4), (27, 8), (29, 32), (30, 3), + (7, 5), (11, 10), (13, 12), (14, 15), (39, 47), (43, 7), (45, 11), (46, 13), + (16, 14), (3, 6), (5, 9), (10, 31), (12, 35), (19, 37), (21, 42), (26, 44), + (28, 33), (35, 34), (37, 36), (42, 40), (44, 39), (1, 43), (2, 45), (4, 46), + (8, 17), (17, 18), (18, 20), (20, 24), (24, 19), (6, 21), (9, 26), (22, 28), + (25, 23), (32, 27), (33, 29), (34, 30), (36, 22), (40, 25), (38, 38), (41, 41), +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cbp_table_is_permutation() { + let mut a = [false; 48]; + let mut b = [false; 48]; + for &(i, p) in &CBP_ME { + a[i as usize] = true; + b[p as usize] = true; + } + assert!(a.iter().all(|&x| x) && b.iter().all(|&x| x)); + } + + #[test] + fn i16_types() { + let m = intra_mb_type(1).unwrap(); + assert_eq!((m.kind, m.i16_mode, m.i16_cbp), (MbKind::I16x16, 0, 0)); + let m = intra_mb_type(24).unwrap(); + assert_eq!((m.i16_mode, m.i16_cbp), (3, 15 | (2 << 4))); + let m = intra_mb_type(13).unwrap(); + assert_eq!((m.i16_mode, m.i16_cbp), (0, 15)); + } + + #[test] + fn b_types() { + let m = b_mb_type(12).unwrap(); + assert_eq!((m.part, m.pred), (Part::P16x8, [PRED_L0, PRED_BI])); + let m = b_mb_type(21).unwrap(); + assert_eq!((m.part, m.pred), (Part::P8x16, [PRED_BI, PRED_BI])); + let m = b_mb_type(11).unwrap(); + assert_eq!((m.part, m.pred), (Part::P8x16, [PRED_L1, PRED_L0])); + } +} diff --git a/crates/h264/src/picture.rs b/crates/h264/src/picture.rs new file mode 100644 index 00000000..94d3335c --- /dev/null +++ b/crates/h264/src/picture.rs @@ -0,0 +1,144 @@ +//! Picture buffers and per-macroblock state. + +use std::sync::Arc; + +/// Planar 4:2:0 8-bit sample buffers (MB-aligned dimensions, stride == width). +#[derive(Clone)] +pub struct Planes { + pub y: Vec, + pub cb: Vec, + pub cr: Vec, + pub width: usize, + pub height: usize, + pub cwidth: usize, + pub cheight: usize, +} + +impl Planes { + pub fn new(width: usize, height: usize) -> Self { + let (cw, ch) = (width / 2, height / 2); + Planes { y: vec![0; width * height], cb: vec![128; cw * ch], cr: vec![128; cw * ch], width, height, cwidth: cw, cheight: ch } + } + pub fn gray(width: usize, height: usize) -> Self { + let mut p = Self::new(width, height); + p.y.fill(128); + p + } +} + +/// Macroblock coding category. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum MbKind { + #[default] + None, + I4x4, + I8x8, + I16x16, + IPcm, + PSkip, + BSkip, + BDirect16x16, + /// Any other P or B inter macroblock. + Inter, +} + +impl MbKind { + #[inline] + pub fn is_intra(self) -> bool { + matches!(self, MbKind::I4x4 | MbKind::I8x8 | MbKind::I16x16 | MbKind::IPcm) + } +} + +/// Per-macroblock state kept for the whole picture (neighbour derivations, deblocking, co-located data). +#[derive(Clone, Copy)] +pub struct MbState { + /// Slice number within the picture; `u32::MAX` = not (yet) decoded. + pub slice_num: u32, + pub kind: MbKind, + pub transform_8x8: bool, + /// coded_block_pattern: bits 0..3 luma, bits 4..5 chroma. + pub cbp: u8, + pub qp: u8, + /// QPc for Cb / Cr (used by deblocking). + pub qpc: [u8; 2], + pub intra_chroma_mode: u8, + /// Intra4x4PredMode / Intra8x8PredMode per 4x4 block (raster order). + pub intra_modes: [u8; 16], + /// CAVLC: TotalCoeff per luma 4x4 block (raster). CABAC: coded_block_flag. + pub nnz: [u8; 16], + /// Same for chroma AC blocks [Cb, Cr][raster 2x2]. + pub nnz_c: [[u8; 4]; 2], + /// CABAC coded_block_flag of DC blocks: bit0 luma (I16x16), bit1 Cb, bit2 Cr. + pub cbf_dc: u8, + /// Luma 4x4 blocks (raster bit) with non-zero coefficients, for deblocking bS = 2. + pub nz_mask: u16, + pub ref_idx: [[i8; 4]; 2], + /// Motion vectors per 4x4 block (raster order). + pub mv: [[[i16; 2]; 16]; 2], + /// |mvd| per 4x4 block (raster), saturated, for CABAC context selection. + #[allow(dead_code)] // used by CABAC context selection + pub mvd: [[[u8; 2]; 16]; 2], + /// Bit per 8x8 block predicted in direct mode. + pub direct8x8: u8, +} + +impl Default for MbState { + fn default() -> Self { + MbState { + slice_num: u32::MAX, + kind: MbKind::None, + transform_8x8: false, + cbp: 0, + qp: 0, + qpc: [0; 2], + intra_chroma_mode: 0, + intra_modes: [2; 16], + nnz: [0; 16], + nnz_c: [[0; 4]; 2], + cbf_dc: 0, + nz_mask: 0, + ref_idx: [[-1; 4]; 2], + mv: [[[0; 2]; 16]; 2], + mvd: [[[0; 2]; 16]; 2], + direct8x8: 0, + } + } +} + +/// Motion data of a decoded picture, used as co-located data for direct prediction. +#[derive(Clone, Default)] +pub struct MotionField { + /// Per 4x4 block: index mb * 16 + raster. + pub mv: [Vec<[i16; 2]>; 2], + /// Per 8x8 block: index mb * 4 + b8. + pub ref_idx: [Vec; 2], + /// Unique id of the referenced picture per 8x8 block (u32::MAX = none). + pub ref_id: [Vec; 2], + pub intra: Vec, +} + +/// A decoded frame as stored in the DPB and referenced by later pictures. +pub struct Frame { + pub id: u32, + pub poc: i32, + pub planes: Planes, + pub motion: MotionField, +} + +pub type FrameRef = Arc; + +/// An entry of RefPicList0/1 for the current slice. +#[derive(Clone)] +pub struct RefPic { + pub frame: FrameRef, + pub long_term: bool, +} + +impl RefPic { + pub fn poc(&self) -> i32 { + self.frame.poc + } + pub fn id(&self) -> u32 { + self.frame.id + } +} diff --git a/crates/h264/src/slicedec.rs b/crates/h264/src/slicedec.rs new file mode 100644 index 00000000..4b693f44 --- /dev/null +++ b/crates/h264/src/slicedec.rs @@ -0,0 +1,1340 @@ +//! Slice data decoding: macroblock parsing (CAVLC), motion vector derivation, and reconstruction. + +use crate::cavlc; +use crate::error::{Result, ensure, invalid}; +use crate::inter::{self, PlaneRef, Weight}; +use crate::intra::{self, Avail}; +use crate::mbtypes::*; +use crate::params::{Pps, Sps}; +use crate::picture::*; +use crate::slice::{SliceHeader, SliceType}; +use crate::tables::*; +use crate::transform::{self, LevelScale}; +use filmcraft_bitstream::BitReader; + +/// Per-slice data kept with the picture (deblocking, motion field export). +#[derive(Clone, Debug, Default)] +pub struct SliceInfo { + pub disable_deblocking_filter_idc: u32, + pub alpha_offset: i32, + pub beta_offset: i32, + /// Unique ids of the pictures in RefPicList0/1. + pub ref_ids: [Vec; 2], +} + +/// The picture under construction. +pub struct PicState { + pub planes: Planes, + pub mbs: Vec, + pub slices: Vec, + pub mb_w: usize, + pub mb_h: usize, + pub poc: i32, + pub decoded_mbs: usize, +} + +impl PicState { + pub fn new(mb_w: usize, mb_h: usize, poc: i32) -> Self { + PicState { planes: Planes::new(mb_w * 16, mb_h * 16), mbs: vec![MbState::default(); mb_w * mb_h], slices: Vec::new(), mb_w, mb_h, poc, decoded_mbs: 0 } + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum WpMode { + Default, + Explicit, + Implicit, +} + +/// Prediction weight tables of a slice. +struct Weights { + mode: WpMode, + luma_denom: i32, + chroma_denom: i32, + /// [list][ref] -> (weight, offset) for luma and two chroma components. + luma: [Vec<(i32, i32)>; 2], + chroma: [Vec<[(i32, i32); 2]>; 2], + /// Implicit weights (w0, w1) indexed by ref0 * n1 + ref1. + implicit: Vec<(i32, i32)>, + n1: usize, +} + +/// Current macroblock syntax (prediction part). +#[derive(Clone, Copy, Default)] +pub(crate) struct MbCur { + pub info: MbTypeInfo, + pub sub: [SubMbInfo; 4], + /// rem_intra_pred_mode per block, or -1 when prev_intra_pred_mode_flag was set. + pub rem_mode: [i8; 16], + pub cbp: u8, + pub ref_idx: [[i8; 4]; 2], + /// mvd per (partition * 4 + sub-partition). + pub mvd: [[[i16; 2]; 16]; 2], +} + +/// Coefficient scratch for one macroblock. +pub(crate) struct Scratch { + /// Luma: 16 4x4 blocks (raster block index * 16 + raster coefficient) or 4 8x8 blocks (b8 * 64 + raster). + pub coef: [i32; 256], + /// Chroma: [Cb, Cr][raster 2x2 block * 16 + raster coefficient]. + pub coef_c: [[i32; 64]; 2], + /// Luma blocks with coefficients: raster 4x4 bits, or b8 bits in 8x8 mode. + pub blk_nz: u16, + /// Chroma blocks with coefficients (bit per raster 2x2 block). + pub c_nz: [u8; 2], + pub pred: [[u8; 256]; 2], + pub pred_c: [[[u8; 64]; 2]; 2], +} + +impl Scratch { + pub fn new() -> Box { + Box::new(Scratch { coef: [0; 256], coef_c: [[0; 64]; 2], blk_nz: 0, c_nz: [0; 2], pred: [[0; 256]; 2], pred_c: [[[0; 64]; 2]; 2] }) + } +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct Nb { + pub avail: bool, + pub ref_idx: i8, + pub mv: [i16; 2], +} + +const NB_NA: Nb = Nb { avail: false, ref_idx: -1, mv: [0, 0] }; + +pub struct SliceDecoder<'a> { + pub sh: &'a SliceHeader, + pub pps: &'a Pps, + pub sps: &'a Sps, + pub pic: &'a mut PicState, + pub refs: &'a [Vec; 2], + pub ls: &'a LevelScale, + pub slice_num: u32, + pub mb_w: usize, + pub mb_h: usize, + pub qp: i32, + pub chroma_qp_offset: [i32; 2], + // current macroblock + pub mb_addr: usize, + pub mb_x: usize, + pub mb_y: usize, + /// Neighbouring macroblocks A (left), B (above), C (above right), D (above left) when available. + pub nb: [Option; 4], + pub mv_done: u16, + pub cur: MbCur, + pub s: Box, + weights: Weights, + /// Temporal direct: DistScaleFactor per refIdxL0, and whether pic0 is long-term. + tdirect: Vec<(i32, bool)>, + /// CABAC: previous MB in decoding order had a non-zero mb_qp_delta. + pub prev_qp_delta_nz: bool, +} + +#[inline(always)] +fn median(a: i16, b: i16, c: i16) -> i16 { + a.max(b).min(a.min(b).max(c)) +} + +impl<'a> SliceDecoder<'a> { + pub fn new(sh: &'a SliceHeader, pps: &'a Pps, sps: &'a Sps, pic: &'a mut PicState, refs: &'a [Vec; 2], ls: &'a LevelScale) -> Result { + let slice_num = pic.slices.len() as u32; + let chroma_qp_offset = [pps.chroma_qp_index_offset, pps.second_chroma_qp_index_offset]; + pic.slices.push(SliceInfo { + disable_deblocking_filter_idc: sh.disable_deblocking_filter_idc, + alpha_offset: sh.slice_alpha_c0_offset_div2 * 2, + beta_offset: sh.slice_beta_offset_div2 * 2, + ref_ids: [refs[0].iter().map(|r| r.id()).collect(), refs[1].iter().map(|r| r.id()).collect()], + }); + let weights = Self::build_weights(sh, pps, refs, pic.poc); + let mut tdirect = Vec::new(); + if sh.slice_type == SliceType::B && !sh.direct_spatial_mv_pred && !refs[1].is_empty() { + let poc1 = refs[1][0].poc(); + for r0 in refs[0].iter() { + let tb = (pic.poc - r0.poc()).clamp(-128, 127); + let td = (poc1 - r0.poc()).clamp(-128, 127); + if td == 0 || r0.long_term { + tdirect.push((256, true)); + } else { + let tx = (16384 + (td / 2).abs()) / td; + tdirect.push((((tb * tx + 32) >> 6).clamp(-1024, 1023), false)); + } + } + } + let (mb_w, mb_h) = (pic.mb_w, pic.mb_h); + Ok(SliceDecoder { + sh, + pps, + sps, + pic, + refs, + ls, + slice_num, + mb_w, + mb_h, + qp: sh.qp(pps), + chroma_qp_offset, + mb_addr: 0, + mb_x: 0, + mb_y: 0, + nb: [None; 4], + mv_done: 0, + cur: MbCur::default(), + s: Scratch::new(), + weights, + tdirect, + prev_qp_delta_nz: false, + }) + } + + fn build_weights(sh: &SliceHeader, pps: &Pps, refs: &[Vec; 2], cur_poc: i32) -> Weights { + let mut w = Weights { + mode: WpMode::Default, + luma_denom: 0, + chroma_denom: 0, + luma: [Vec::new(), Vec::new()], + chroma: [Vec::new(), Vec::new()], + implicit: Vec::new(), + n1: refs[1].len(), + }; + let explicit = (pps.weighted_pred && sh.slice_type.is_p()) || (pps.weighted_bipred_idc == 1 && sh.slice_type.is_b()); + if explicit { + if let Some(t) = &sh.pred_weight_table { + w.mode = WpMode::Explicit; + w.luma_denom = t.luma_log2_denom as i32; + w.chroma_denom = t.chroma_log2_denom as i32; + for (l, entries) in [&t.l0, &t.l1].into_iter().enumerate() { + w.luma[l] = entries.iter().map(|e| (e.luma_weight, e.luma_offset)).collect(); + w.chroma[l] = entries.iter().map(|e| [(e.chroma_weight[0], e.chroma_offset[0]), (e.chroma_weight[1], e.chroma_offset[1])]).collect(); + } + } + } else if pps.weighted_bipred_idc == 2 && sh.slice_type.is_b() { + w.mode = WpMode::Implicit; + for r0 in refs[0].iter() { + for r1 in refs[1].iter() { + let tb = (cur_poc - r0.poc()).clamp(-128, 127); + let td = (r1.poc() - r0.poc()).clamp(-128, 127); + let pair = if td == 0 || r0.long_term || r1.long_term { + (32, 32) + } else { + let tx = (16384 + (td / 2).abs()) / td; + let dsf = ((tb * tx + 32) >> 6).clamp(-1024, 1023); + if (dsf >> 2) < -64 || (dsf >> 2) > 128 { (32, 32) } else { (64 - (dsf >> 2), dsf >> 2) } + }; + w.implicit.push(pair); + } + } + } + w + } + + // ------------------------------------------------------------------------------------------ + // Neighbour helpers + // ------------------------------------------------------------------------------------------ + + /// Set up the current macroblock address and neighbour availability. + pub fn start_mb(&mut self, addr: usize) { + self.mb_addr = addr; + self.mb_x = addr % self.mb_w; + self.mb_y = addr / self.mb_w; + let sn = self.slice_num; + let mbs = &self.pic.mbs; + let w = self.mb_w; + let (x, y) = (self.mb_x, self.mb_y); + let ok = |a: usize| if mbs[a].slice_num == sn { Some(a) } else { None }; + self.nb = [ + if x > 0 { ok(addr - 1) } else { None }, + if y > 0 { ok(addr - w) } else { None }, + if y > 0 && x + 1 < w { ok(addr - w + 1) } else { None }, + if y > 0 && x > 0 { ok(addr - w - 1) } else { None }, + ]; + self.mv_done = 0; + self.pic.mbs[addr] = MbState { slice_num: sn, ..MbState::default() }; + } + + #[inline] + pub fn mb(&self) -> &MbState { + &self.pic.mbs[self.mb_addr] + } + #[inline] + pub fn mb_mut(&mut self) -> &mut MbState { + &mut self.pic.mbs[self.mb_addr] + } + + /// Neighbour MB usable for intra prediction (constrained_intra_pred aware). + fn intra_nb(&self, n: usize) -> bool { + match self.nb[n] { + None => false, + Some(a) => !self.pps.constrained_intra_pred || self.pic.mbs[a].kind.is_intra(), + } + } + + fn chroma_qp(&self, qp: i32, c: usize) -> i32 { + QPC_TABLE[(qp + self.chroma_qp_offset[c]).clamp(0, 51) as usize] as i32 + } + + /// nC for a luma 4x4 block (raster index) - 9.2.1. + fn nc_luma(&self, raster: usize) -> i32 { + let (bx, by) = (raster & 3, raster >> 2); + let cur = self.mb(); + let a = if bx > 0 { Some(cur.nnz[raster - 1]) } else { self.nb[0].map(|m| self.pic.mbs[m].nnz[by * 4 + 3]) }; + let b = if by > 0 { Some(cur.nnz[raster - 4]) } else { self.nb[1].map(|m| self.pic.mbs[m].nnz[12 + bx]) }; + match (a, b) { + (Some(a), Some(b)) => (a as i32 + b as i32 + 1) >> 1, + (Some(a), None) => a as i32, + (None, Some(b)) => b as i32, + _ => 0, + } + } + + fn nc_chroma(&self, c: usize, raster: usize) -> i32 { + let (bx, by) = (raster & 1, raster >> 1); + let cur = self.mb(); + let a = if bx > 0 { Some(cur.nnz_c[c][raster - 1]) } else { self.nb[0].map(|m| self.pic.mbs[m].nnz_c[c][by * 2 + 1]) }; + let b = if by > 0 { Some(cur.nnz_c[c][raster - 2]) } else { self.nb[1].map(|m| self.pic.mbs[m].nnz_c[c][2 + bx]) }; + match (a, b) { + (Some(a), Some(b)) => (a as i32 + b as i32 + 1) >> 1, + (Some(a), None) => a as i32, + (None, Some(b)) => b as i32, + _ => 0, + } + } + + /// Motion data of the neighbouring 4x4 block at (x, y) (4x4 units, relative to the current MB). + pub(crate) fn nb_motion(&self, list: usize, x: i32, y: i32) -> Nb { + let (mb, bx, by) = if x < 0 { + if y < 0 { + (self.nb[3], 3, 3) + } else if y > 3 { + (None, 0, 0) + } else { + (self.nb[0], 3, y as usize) + } + } else if x > 3 { + if y < 0 { (self.nb[2], 0, 3) } else { (None, 0, 0) } + } else if y < 0 { + (self.nb[1], x as usize, 3) + } else { + let r = (y * 4 + x) as usize; + if self.mv_done & (1 << r) == 0 { + return NB_NA; + } + (Some(self.mb_addr), x as usize, y as usize) + }; + let Some(m) = mb else { return NB_NA }; + let st = &self.pic.mbs[m]; + let b8 = (by >> 1) * 2 + (bx >> 1); + let r = st.ref_idx[list][b8]; + if r < 0 { + return Nb { avail: true, ref_idx: -1, mv: [0, 0] }; + } + Nb { avail: true, ref_idx: r, mv: st.mv[list][by * 4 + bx] } + } + + /// Returns (A, B, C) neighbour motion for a partition at (x, y) of width w (4x4 units), with C + /// replaced by D when unavailable. + fn nb_abc(&self, list: usize, x: usize, y: usize, w: usize) -> (Nb, Nb, Nb) { + let (x, y, w) = (x as i32, y as i32, w as i32); + let a = self.nb_motion(list, x - 1, y); + let b = self.nb_motion(list, x, y - 1); + let mut c = self.nb_motion(list, x + w, y - 1); + if !c.avail { + c = self.nb_motion(list, x - 1, y - 1); + } + (a, b, c) + } + + /// Luma motion vector prediction (8.4.1.3). `shape`: 1 = 16x8, 2 = 8x16, 0 = other. + #[allow(clippy::too_many_arguments)] + fn mvp(&self, list: usize, x: usize, y: usize, w: usize, ref_idx: i8, shape: u8, part: usize) -> [i16; 2] { + let (a, b, c) = self.nb_abc(list, x, y, w); + if shape == 1 { + if part == 0 && b.ref_idx == ref_idx { + return b.mv; + } + if part == 1 && a.ref_idx == ref_idx { + return a.mv; + } + } else if shape == 2 { + if part == 0 && a.ref_idx == ref_idx { + return a.mv; + } + if part == 1 && c.ref_idx == ref_idx { + return c.mv; + } + } + Self::median_pred(a, b, c, ref_idx) + } + + fn median_pred(a: Nb, mut b: Nb, mut c: Nb, ref_idx: i8) -> [i16; 2] { + if !b.avail && !c.avail && a.avail { + b = a; + c = a; + } + let ma = a.ref_idx == ref_idx; + let mb = b.ref_idx == ref_idx; + let mc = c.ref_idx == ref_idx; + match (ma, mb, mc) { + (true, false, false) => a.mv, + (false, true, false) => b.mv, + (false, false, true) => c.mv, + _ => [median(a.mv[0], b.mv[0], c.mv[0]), median(a.mv[1], b.mv[1], c.mv[1])], + } + } + + /// Store motion for a rectangle (4x4 units) of the current MB and mark it decoded. + fn set_motion(&mut self, x: usize, y: usize, w: usize, h: usize, refs: [i8; 2], mvs: [[i16; 2]; 2]) { + let st = &mut self.pic.mbs[self.mb_addr]; + for yy in y..y + h { + for xx in x..x + w { + let r = yy * 4 + xx; + for l in 0..2 { + st.mv[l][r] = if refs[l] >= 0 { mvs[l] } else { [0, 0] }; + } + self.mv_done |= 1 << r; + } + } + for l in 0..2 { + // ref_idx is per 8x8; rectangles never straddle 8x8 blocks with different refs + for yy in (y..y + h).step_by(1) { + for xx in x..x + w { + st.ref_idx[l][(yy >> 1) * 2 + (xx >> 1)] = refs[l]; + } + } + } + } + + // ------------------------------------------------------------------------------------------ + // Direct prediction (8.4.1.2) + // ------------------------------------------------------------------------------------------ + + /// Co-located (mvCol, refIdxCol, referenced picture id) for 4x4 block `raster` of the current MB. + fn colocated(&self, raster: usize) -> ([i16; 2], i8, u32) { + let Some(col) = self.refs[1].first() else { return ([0, 0], -1, u32::MAX) }; + let m = &col.frame.motion; + let raster = if self.sps.direct_8x8_inference { + let (bx, by) = (raster & 3, raster >> 2); + (if by >= 2 { 12 } else { 0 }) + if bx >= 2 { 3 } else { 0 } + } else { + raster + }; + let mb = self.mb_addr; + if m.intra.get(mb).copied().unwrap_or(true) { + return ([0, 0], -1, u32::MAX); + } + let b8 = mb * 4 + (raster >> 3) * 2 + ((raster & 3) >> 1); + let l = if m.ref_idx[0][b8] >= 0 { 0 } else { 1 }; + (m.mv[l][mb * 16 + raster], m.ref_idx[l][b8], m.ref_id[l][b8]) + } + + /// Spatial direct: MB-level reference indices, predicted mvs and directZeroPrediction (8.4.1.2.2). + fn spatial_direct_base(&self) -> ([i8; 2], [[i16; 2]; 2], bool) { + let mut refs = [-1i8; 2]; + let mut nbs = [(NB_NA, NB_NA, NB_NA); 2]; + let min_pos = |x: i8, y: i8| if x >= 0 && y >= 0 { x.min(y) } else { x.max(y) }; + for (l, nb) in nbs.iter_mut().enumerate() { + *nb = self.nb_abc(l, 0, 0, 4); + refs[l] = min_pos(nb.0.ref_idx, min_pos(nb.1.ref_idx, nb.2.ref_idx)); + } + let mut mvs = [[0i16; 2]; 2]; + if refs[0] < 0 && refs[1] < 0 { + return ([0, 0], mvs, true); + } + for l in 0..2 { + if refs[l] >= 0 { + let (a, b, c) = nbs[l]; + mvs[l] = Self::median_pred(a, b, c, refs[l]); + } + } + (refs, mvs, false) + } + + /// Derive direct-mode motion for the 8x8 blocks in `b8_mask` of the current MB. + fn derive_direct(&mut self, b8_mask: u8) -> Result<()> { + let spatial = self.sh.direct_spatial_mv_pred; + let (base_refs, base_mvs, direct_zero) = if spatial { self.spatial_direct_base() } else { ([0, 0], [[0, 0]; 2], false) }; + let col_short_term = self.refs[1].first().map(|r| !r.long_term).unwrap_or(false); + for b8 in 0..4 { + if b8_mask & (1 << b8) == 0 { + continue; + } + let (x8, y8) = ((b8 & 1) * 2, (b8 >> 1) * 2); + for sub in 0..4 { + let (x, y) = (x8 + (sub & 1), y8 + (sub >> 1)); + let raster = y * 4 + x; + let (mv_col, ref_col, id_col) = self.colocated(raster); + let (refs, mvs) = if spatial { + let col_zero = col_short_term && ref_col == 0 && (-1..=1).contains(&mv_col[0]) && (-1..=1).contains(&mv_col[1]); + let mut mvs = base_mvs; + for l in 0..2 { + if direct_zero || base_refs[l] < 0 || (base_refs[l] == 0 && col_zero) { + mvs[l] = [0, 0]; + } + } + (base_refs, mvs) + } else { + let ref0 = if ref_col < 0 { + 0 + } else { + match self.refs[0].iter().position(|r| r.id() == id_col) { + Some(i) => i as i8, + None => 0, + } + }; + let (dsf, ltr) = self.tdirect.get(ref0 as usize).copied().unwrap_or((256, true)); + let (mv0, mv1) = if ltr { + (mv_col, [0, 0]) + } else { + let s = |c: i16| ((dsf * c as i32 + 128) >> 8) as i16; + let m0 = [s(mv_col[0]), s(mv_col[1])]; + (m0, [m0[0] - mv_col[0], m0[1] - mv_col[1]]) + }; + ([ref0, 0], [mv0, mv1]) + }; + self.set_motion(x, y, 1, 1, refs, mvs); + } + } + let st = &mut self.pic.mbs[self.mb_addr]; + st.direct8x8 |= b8_mask; + Ok(()) + } + + // ------------------------------------------------------------------------------------------ + // Motion vector derivation for a parsed inter macroblock + // ------------------------------------------------------------------------------------------ + + /// P_Skip motion (8.4.1.1). + pub fn derive_p_skip(&mut self) { + let a = self.nb_motion(0, -1, 0); + let b = self.nb_motion(0, 0, -1); + let mv = if self.nb[0].is_none() || self.nb[1].is_none() || (a.ref_idx == 0 && a.mv == [0, 0]) || (b.ref_idx == 0 && b.mv == [0, 0]) { + [0, 0] + } else { + self.mvp(0, 0, 0, 4, 0, 0, 0) + }; + self.set_motion(0, 0, 4, 4, [0, -1], [mv, [0, 0]]); + } + + /// Derive mvs of the current (non-skip) inter MB from `self.cur` (ref_idx, mvd). + pub fn derive_inter_motion(&mut self) -> Result<()> { + let info = self.cur.info; + if info.kind == MbKind::BDirect16x16 || info.kind == MbKind::BSkip { + return self.derive_direct(0xf); + } + let n = [self.refs[0].len(), self.refs[1].len()]; + if info.part != Part::P8x8 { + let shape = match info.part { + Part::P16x8 => 1, + Part::P8x16 => 2, + _ => 0, + }; + for p in 0..info.part.num_parts() { + let (x, y, w, h) = info.part.rect(p); + let mut refs = [-1i8; 2]; + let mut mvs = [[0i16; 2]; 2]; + for l in 0..2 { + if info.pred[p] & (1 << l) != 0 { + let r = self.cur.ref_idx[l][p]; + ensure!((r as usize) < n[l], "ref_idx {r} out of range (list {l}, size {})", n[l]); + let mvp = self.mvp(l, x, y, w, r, shape, p); + let d = self.cur.mvd[l][p * 4]; + refs[l] = r; + mvs[l] = [mvp[0].wrapping_add(d[0]), mvp[1].wrapping_add(d[1])]; + } + } + self.set_motion(x, y, w, h, refs, mvs); + } + return Ok(()); + } + for p in 0..4 { + let sub = self.cur.sub[p]; + let (x8, y8) = ((p & 1) * 2, (p >> 1) * 2); + if sub.direct { + self.derive_direct(1 << p)?; + continue; + } + let mut refs = [-1i8; 2]; + for (l, r) in refs.iter_mut().enumerate() { + if sub.pred & (1 << l) != 0 { + let v = self.cur.ref_idx[l][p]; + ensure!((v as usize) < n[l], "ref_idx {v} out of range (list {l}, size {})", n[l]); + *r = v; + } + } + for j in 0..sub.shape.num_parts() { + let (sx, sy, w, h) = sub.shape.rect(j); + let (x, y) = (x8 + sx, y8 + sy); + let mut mvs = [[0i16; 2]; 2]; + for l in 0..2 { + if refs[l] >= 0 { + let mvp = self.mvp(l, x, y, w, refs[l], 0, 0); + let d = self.cur.mvd[l][p * 4 + j]; + mvs[l] = [mvp[0].wrapping_add(d[0]), mvp[1].wrapping_add(d[1])]; + } + } + self.set_motion(x, y, w, h, refs, mvs); + } + } + Ok(()) + } + + // ------------------------------------------------------------------------------------------ + // Inter prediction + // ------------------------------------------------------------------------------------------ + + fn weight_for(&self, refs: [i8; 2], comp: usize) -> Weight { + let w = &self.weights; + match w.mode { + WpMode::Default => Weight::Default, + WpMode::Implicit => { + if refs[0] >= 0 && refs[1] >= 0 { + let (w0, w1) = w.implicit[refs[0] as usize * w.n1 + refs[1] as usize]; + Weight::Weighted { log_wd: 5, w0, w1, o0: 0, o1: 0 } + } else { + Weight::Default + } + } + WpMode::Explicit => { + let get = |l: usize| -> (i32, i32) { + if refs[l] < 0 { + return (0, 0); + } + let r = refs[l] as usize; + if comp == 0 { + w.luma[l].get(r).copied().unwrap_or((1 << w.luma_denom, 0)) + } else { + w.chroma[l].get(r).map(|c| c[comp - 1]).unwrap_or((1 << w.chroma_denom, 0)) + } + }; + let (w0, o0) = get(0); + let (w1, o1) = get(1); + let log_wd = if comp == 0 { w.luma_denom } else { w.chroma_denom }; + Weight::Weighted { log_wd, w0, w1, o0, o1 } + } + } + } + + /// Motion-compensated prediction of a rectangle (4x4 units) written into the picture. + fn mc_rect(&mut self, x4: usize, y4: usize, w4: usize, h4: usize) { + let st = &self.pic.mbs[self.mb_addr]; + let r = y4 * 4 + x4; + let b8 = (y4 >> 1) * 2 + (x4 >> 1); + let refs = [st.ref_idx[0][b8], st.ref_idx[1][b8]]; + let mvs = [st.mv[0][r], st.mv[1][r]]; + let (bw, bh) = (w4 * 4, h4 * 4); + let wts = [self.weight_for(refs, 0), self.weight_for(refs, 1), self.weight_for(refs, 2)]; + let px = self.mb_x * 16 + x4 * 4; + let py = self.mb_y * 16 + y4 * 4; + let s = &mut *self.s; + for l in 0..2 { + if refs[l] < 0 { + continue; + } + let Some(rp) = self.refs[l].get(refs[l] as usize) else { continue }; + let f = &rp.frame.planes; + let mv = mvs[l]; + let yref = PlaneRef { data: &f.y, width: f.width, height: f.height, stride: f.width }; + inter::mc_luma( + yref, + px as i32 + (mv[0] as i32 >> 2), + py as i32 + (mv[1] as i32 >> 2), + (mv[0] & 3) as u32, + (mv[1] & 3) as u32, + bw, + bh, + &mut s.pred[l], + 16, + ); + let (cx, cy) = (px as i32 / 2, py as i32 / 2); + for (c, plane) in [&f.cb, &f.cr].into_iter().enumerate() { + let cref = PlaneRef { data: plane, width: f.cwidth, height: f.cheight, stride: f.cwidth }; + inter::mc_chroma( + cref, + cx + (mv[0] as i32 >> 3), + cy + (mv[1] as i32 >> 3), + (mv[0] & 7) as u32, + (mv[1] & 7) as u32, + bw / 2, + bh / 2, + &mut s.pred_c[l][c], + 8, + ); + } + } + let use0 = refs[0] >= 0 && self.refs[0].get(refs[0] as usize).is_some(); + let use1 = refs[1] >= 0 && self.refs[1].get(refs[1] as usize).is_some(); + let stride = self.pic.planes.width; + inter::weighted_store( + &mut self.pic.planes.y[py * stride + px..], + stride, + if use0 { Some(&s.pred[0][..]) } else { None }, + if use1 { Some(&s.pred[1][..]) } else { None }, + 16, + bw, + bh, + wts[0], + ); + let cstride = self.pic.planes.cwidth; + let (cx, cy) = (px / 2, py / 2); + for c in 0..2 { + let wc = wts[c + 1]; + let a = if use0 { Some(&s.pred_c[0][c][..]) } else { None }; + let b = if use1 { Some(&s.pred_c[1][c][..]) } else { None }; + let plane = if c == 0 { &mut self.pic.planes.cb } else { &mut self.pic.planes.cr }; + inter::weighted_store(&mut plane[cy * cstride + cx..], cstride, a, b, 8, bw / 2, bh / 2, wc); + } + } + + /// Inter prediction for the whole current MB, using the stored motion. + pub fn predict_inter(&mut self) { + let info = self.cur.info; + let direct_unit = if self.sps.direct_8x8_inference { 2 } else { 1 }; + match info.kind { + MbKind::PSkip => self.mc_rect(0, 0, 4, 4), + MbKind::BSkip | MbKind::BDirect16x16 => { + for y in (0..4).step_by(direct_unit) { + for x in (0..4).step_by(direct_unit) { + self.mc_rect(x, y, direct_unit, direct_unit); + } + } + } + _ => { + if info.part != Part::P8x8 { + for p in 0..info.part.num_parts() { + let (x, y, w, h) = info.part.rect(p); + self.mc_rect(x, y, w, h); + } + } else { + for p in 0..4 { + let sub = self.cur.sub[p]; + let (x8, y8) = ((p & 1) * 2, (p >> 1) * 2); + if sub.direct { + for y in (0..2).step_by(direct_unit) { + for x in (0..2).step_by(direct_unit) { + self.mc_rect(x8 + x, y8 + y, direct_unit, direct_unit); + } + } + } else { + for j in 0..sub.shape.num_parts() { + let (sx, sy, w, h) = sub.shape.rect(j); + self.mc_rect(x8 + sx, y8 + sy, w, h); + } + } + } + } + } + } + } + + // ------------------------------------------------------------------------------------------ + // Intra prediction mode derivation (8.3.1.1 / 8.3.2.1) + // ------------------------------------------------------------------------------------------ + + fn intra_mode_nb(&self, n: usize, bx: usize, by: usize) -> Option { + // returns Some(mode) or None if dcPredModePredictedFlag must be set + let m = self.nb[n]?; + let st = &self.pic.mbs[m]; + if self.pps.constrained_intra_pred && !st.kind.is_intra() { + return None; + } + Some(match st.kind { + MbKind::I4x4 | MbKind::I8x8 => st.intra_modes[by * 4 + bx], + _ => 2, + }) + } + + /// Resolve Intra4x4PredMode / Intra8x8PredMode for the current MB from `self.cur.rem_mode`. + pub fn derive_intra_modes(&mut self, is8x8: bool) { + let step = if is8x8 { 2 } else { 1 }; + let nblk = if is8x8 { 4 } else { 16 }; + for i in 0..nblk { + let (bx, by) = if is8x8 { + ((i & 1) * 2, (i >> 1) * 2) + } else { + let (x, y) = BLK4_XY[i]; + (x as usize, y as usize) + }; + let a = if bx > 0 { Some(self.mb().intra_modes[by * 4 + bx - 1]) } else { self.intra_mode_nb(0, 3, by) }; + let b = if by > 0 { Some(self.mb().intra_modes[(by - 1) * 4 + bx]) } else { self.intra_mode_nb(1, bx, 3) }; + let pred = match (a, b) { + (Some(a), Some(b)) => a.min(b), + _ => 2, + }; + let rem = self.cur.rem_mode[i]; + let mode = if rem < 0 { + pred + } else if (rem as u8) < pred { + rem as u8 + } else { + rem as u8 + 1 + }; + let st = self.mb_mut(); + for yy in by..by + step { + for xx in bx..bx + step { + st.intra_modes[yy * 4 + xx] = mode; + } + } + } + } + + // ------------------------------------------------------------------------------------------ + // Reconstruction + // ------------------------------------------------------------------------------------------ + + fn add_luma_block4(&mut self, raster: usize) { + if self.s.blk_nz & (1 << raster) == 0 { + return; + } + let stride = self.pic.planes.width; + let px = self.mb_x * 16 + (raster & 3) * 4; + let py = self.mb_y * 16 + (raster >> 2) * 4; + let blk: &mut [i32; 16] = (&mut self.s.coef[raster * 16..raster * 16 + 16]).try_into().unwrap(); + transform::idct4_add(blk, &mut self.pic.planes.y[py * stride + px..], stride); + blk.fill(0); + } + + fn add_luma_block8(&mut self, b8: usize) { + if self.s.blk_nz & (1 << b8) == 0 { + return; + } + let stride = self.pic.planes.width; + let px = self.mb_x * 16 + (b8 & 1) * 8; + let py = self.mb_y * 16 + (b8 >> 1) * 8; + let blk: &mut [i32; 64] = (&mut self.s.coef[b8 * 64..b8 * 64 + 64]).try_into().unwrap(); + transform::idct8_add(blk, &mut self.pic.planes.y[py * stride + px..], stride); + blk.fill(0); + } + + fn add_chroma(&mut self) { + let stride = self.pic.planes.cwidth; + for c in 0..2 { + if self.s.c_nz[c] == 0 { + continue; + } + for b in 0..4 { + if self.s.c_nz[c] & (1 << b) == 0 { + continue; + } + let px = self.mb_x * 8 + (b & 1) * 4; + let py = self.mb_y * 8 + (b >> 1) * 4; + let blk: &mut [i32; 16] = (&mut self.s.coef_c[c][b * 16..b * 16 + 16]).try_into().unwrap(); + let plane = if c == 0 { &mut self.pic.planes.cb } else { &mut self.pic.planes.cr }; + transform::idct4_add(blk, &mut plane[py * stride + px..], stride); + blk.fill(0); + } + self.s.c_nz[c] = 0; + } + } + + fn luma_avail4(&self, bx: usize, by: usize) -> Avail { + let a = self.intra_nb(0); + let b = self.intra_nb(1); + let c = self.intra_nb(2); + let d = self.intra_nb(3); + let left = if bx > 0 { true } else { a }; + let top = if by > 0 { true } else { b }; + let top_left = match (bx > 0, by > 0) { + (true, true) => true, + (false, true) => a, + (true, false) => b, + (false, false) => d, + }; + let top_right = if by == 0 { + if bx < 3 { b } else { c } + } else if bx == 3 { + false + } else { + RASTER_TO_BLK4[(by - 1) * 4 + bx + 1] < RASTER_TO_BLK4[by * 4 + bx] + }; + Avail { left, top, top_left, top_right } + } + + /// Reconstruct an intra MB (prediction + residual) after its syntax has been parsed. + pub fn reconstruct_intra(&mut self) { + let kind = self.mb().kind; + let stride = self.pic.planes.width; + let (x0, y0) = (self.mb_x * 16, self.mb_y * 16); + match kind { + MbKind::I4x4 => { + for (blk, &(bx, by)) in BLK4_XY.iter().enumerate() { + let (bx, by) = (bx as usize, by as usize); + let _ = blk; + let av = self.luma_avail4(bx, by); + let mode = self.mb().intra_modes[by * 4 + bx]; + intra::pred4x4(&mut self.pic.planes.y, stride, x0 + bx * 4, y0 + by * 4, mode, av); + self.add_luma_block4(by * 4 + bx); + } + } + MbKind::I8x8 => { + for b8 in 0..4 { + let (bx, by) = ((b8 & 1) * 2, (b8 >> 1) * 2); + let mut av = self.luma_avail4(bx, by); + av.top_right = match b8 { + 0 => self.intra_nb(1), + 1 => self.intra_nb(2), + 2 => true, + _ => false, + }; + let mode = self.mb().intra_modes[by * 4 + bx]; + intra::pred8x8(&mut self.pic.planes.y, stride, x0 + bx * 4, y0 + by * 4, mode, av); + self.add_luma_block8(b8); + } + } + MbKind::I16x16 => { + let av = Avail { left: self.intra_nb(0), top: self.intra_nb(1), top_left: self.intra_nb(3), top_right: false }; + intra::pred16x16(&mut self.pic.planes.y, stride, x0, y0, self.cur.info.i16_mode, av); + for r in 0..16 { + self.add_luma_block4(r); + } + } + _ => {} + } + let av = Avail { left: self.intra_nb(0), top: self.intra_nb(1), top_left: self.intra_nb(3), top_right: false }; + let cs = self.pic.planes.cwidth; + let mode = self.mb().intra_chroma_mode; + let (cx, cy) = (self.mb_x * 8, self.mb_y * 8); + intra::pred_chroma(&mut self.pic.planes.cb, cs, cx, cy, mode, av); + intra::pred_chroma(&mut self.pic.planes.cr, cs, cx, cy, mode, av); + self.add_chroma(); + self.s.blk_nz = 0; + } + + /// Add the residual of an inter MB (prediction already written). + pub fn reconstruct_inter_residual(&mut self) { + if self.mb().transform_8x8 { + for b8 in 0..4 { + self.add_luma_block8(b8); + } + } else if self.s.blk_nz != 0 { + for r in 0..16 { + self.add_luma_block4(r); + } + } + self.add_chroma(); + self.s.blk_nz = 0; + } + + /// Write I_PCM samples. + pub fn write_pcm(&mut self, samples: &[u8]) { + let stride = self.pic.planes.width; + let (x0, y0) = (self.mb_x * 16, self.mb_y * 16); + for r in 0..16 { + self.pic.planes.y[(y0 + r) * stride + x0..(y0 + r) * stride + x0 + 16].copy_from_slice(&samples[r * 16..r * 16 + 16]); + } + let cs = self.pic.planes.cwidth; + let (cx, cy) = (self.mb_x * 8, self.mb_y * 8); + for c in 0..2 { + let src = &samples[256 + c * 64..256 + c * 64 + 64]; + let plane = if c == 0 { &mut self.pic.planes.cb } else { &mut self.pic.planes.cr }; + for r in 0..8 { + plane[(cy + r) * cs + cx..(cy + r) * cs + cx + 8].copy_from_slice(&src[r * 8..r * 8 + 8]); + } + } + } + + /// Common state for an I_PCM MB. + pub fn finish_pcm(&mut self) { + let qpc = [self.chroma_qp(0, 0) as u8, self.chroma_qp(0, 1) as u8]; + let st = self.mb_mut(); + st.kind = MbKind::IPcm; + st.qp = 0; + st.qpc = qpc; + st.nnz = [16; 16]; + st.nnz_c = [[16; 4]; 2]; + st.cbf_dc = 7; + st.cbp = 0x2f; + st.nz_mask = 0xffff; + st.ref_idx = [[-1; 4]; 2]; + } + + /// Update QP after mb_qp_delta. + pub fn apply_qp_delta(&mut self, delta: i32) -> Result<()> { + ensure!((-26..=25).contains(&delta), "mb_qp_delta {delta} out of range"); + self.qp = (self.qp + delta + 52) % 52; + Ok(()) + } + + pub fn store_qp(&mut self) { + let qp = self.qp; + let qpc = [self.chroma_qp(qp, 0) as u8, self.chroma_qp(qp, 1) as u8]; + let st = self.mb_mut(); + st.qp = qp as u8; + st.qpc = qpc; + } + + // ------------------------------------------------------------------------------------------ + // Residual storage helpers (shared by CAVLC and CABAC) + // ------------------------------------------------------------------------------------------ + + fn luma_list4(&self) -> usize { + if self.mb().kind.is_intra() { 0 } else { 3 } + } + + /// Store a 4x4 luma block given in scan order (levels[0..16]); `ac_only` skips position 0. + pub fn put_luma4(&mut self, raster: usize, levels: &[i32; 16], ac_only: bool) { + let qp = self.qp; + let ls = &self.ls.ls4[self.luma_list4()][(qp % 6) as usize]; + let blk = &mut self.s.coef[raster * 16..raster * 16 + 16]; + let mut any = false; + for (k, &lv) in levels.iter().enumerate().skip(ac_only as usize) { + if lv != 0 { + let p = ZIGZAG4[k] as usize; + blk[p] = transform::scale4(lv, ls[p], qp); + any = true; + } + } + if any { + self.s.blk_nz |= 1 << raster; + self.pic.mbs[self.mb_addr].nz_mask |= 1 << raster; + } + } + + /// Store an 8x8 luma block given in scan order. + pub fn put_luma8(&mut self, b8: usize, levels: &[i32; 64]) { + let qp = self.qp; + let list = if self.mb().kind.is_intra() { 0 } else { 1 }; + let ls = &self.ls.ls8[list][(qp % 6) as usize]; + let blk = &mut self.s.coef[b8 * 64..b8 * 64 + 64]; + let mut any = false; + for (k, &lv) in levels.iter().enumerate() { + if lv != 0 { + let p = ZIGZAG8[k] as usize; + blk[p] = transform::scale8(lv, ls[p], qp); + any = true; + } + } + if any { + self.s.blk_nz |= 1 << b8; + let (bx, by) = ((b8 & 1) * 2, (b8 >> 1) * 2); + let m = (0b11u16 << (by * 4 + bx)) | (0b11u16 << ((by + 1) * 4 + bx)); + self.pic.mbs[self.mb_addr].nz_mask |= m; + } + } + + /// Intra16x16 DC levels in scan order -> scaled DC of each 4x4 block. + pub fn put_luma_dc(&mut self, levels: &[i32; 16]) { + let mut c = [0i32; 16]; + for (k, &lv) in levels.iter().enumerate() { + c[ZIGZAG4[k] as usize] = lv; + } + let qp = self.qp; + let ls00 = self.ls.ls4[0][(qp % 6) as usize][0]; + transform::luma_dc_dequant(&mut c, qp, ls00); + for (r, &dc) in c.iter().enumerate() { + if dc != 0 { + self.s.coef[r * 16] = dc; + self.s.blk_nz |= 1 << r; + } + } + } + + /// Chroma DC levels (4:2:0: 4 values in raster order) for component `c`. + pub fn put_chroma_dc(&mut self, c: usize, levels: &[i32; 4]) { + let mut v = *levels; + let qpc = self.chroma_qp(self.qp, c); + let list = if self.mb().kind.is_intra() { 1 + c } else { 4 + c }; + let ls00 = self.ls.ls4[list][(qpc % 6) as usize][0]; + transform::chroma_dc_dequant_420(&mut v, qpc, ls00); + for (b, &dc) in v.iter().enumerate() { + if dc != 0 { + self.s.coef_c[c][b * 16] = dc; + self.s.c_nz[c] |= 1 << b; + } + } + } + + /// Chroma AC levels (scan positions 1..16 in levels[1..]) for block `b` of component `c`. + pub fn put_chroma_ac(&mut self, c: usize, b: usize, levels: &[i32; 16]) { + let qpc = self.chroma_qp(self.qp, c); + let list = if self.mb().kind.is_intra() { 1 + c } else { 4 + c }; + let ls = &self.ls.ls4[list][(qpc % 6) as usize]; + let blk = &mut self.s.coef_c[c][b * 16..b * 16 + 16]; + let mut any = false; + for (k, &lv) in levels.iter().enumerate().skip(1) { + if lv != 0 { + let p = ZIGZAG4[k] as usize; + blk[p] = transform::scale4(lv, ls[p], qpc); + any = true; + } + } + if any { + self.s.c_nz[c] |= 1 << b; + } + } + + // ------------------------------------------------------------------------------------------ + // CAVLC macroblock layer + // ------------------------------------------------------------------------------------------ + + /// Decode a P_Skip or B_Skip macroblock at the current address. + pub fn decode_skip(&mut self) -> Result<()> { + let b = self.sh.slice_type.is_b(); + self.cur = MbCur::default(); + self.cur.info.kind = if b { MbKind::BSkip } else { MbKind::PSkip }; + self.mb_mut().kind = self.cur.info.kind; + self.store_qp(); + if b { + self.derive_direct(0xf)?; + } else { + ensure!(!self.refs[0].is_empty(), "P_Skip with empty reference list"); + self.derive_p_skip(); + } + self.predict_inter(); + self.prev_qp_delta_nz = false; + Ok(()) + } + + fn read_te(r: &mut BitReader, range: u32) -> Result { + Ok(if range > 1 { r.read_ue()? } else { !r.read_bit()? as u32 }) + } + + /// Parse and reconstruct one macroblock_layer() with CAVLC. + pub fn decode_mb_cavlc(&mut self, r: &mut BitReader) -> Result<()> { + let mb_type = r.read_ue()?; + let st = self.sh.slice_type; + let info = match st { + SliceType::I | SliceType::Si => intra_mb_type(mb_type), + SliceType::P | SliceType::Sp => { + if mb_type < 5 { + p_mb_type(mb_type) + } else { + intra_mb_type(mb_type - 5) + } + } + SliceType::B => { + if mb_type < 23 { + b_mb_type(mb_type) + } else { + intra_mb_type(mb_type - 23) + } + } + }; + let Some(mut info) = info else { return invalid(format!("mb_type {mb_type} out of range")) }; + self.cur = MbCur { info, ..Default::default() }; + if info.kind == MbKind::IPcm { + r.byte_align(); + let pos = r.byte_pos(); + let data = r.data(); + ensure!(data.len() >= pos + 384, "truncated I_PCM macroblock"); + let samples: [u8; 384] = data[pos..pos + 384].try_into().unwrap(); + r.skip(384 * 8)?; + self.finish_pcm(); + self.write_pcm(&samples); + return Ok(()); + } + let mut transform_8x8 = false; + if info.kind == MbKind::I4x4 && self.pps.transform_8x8_mode && r.read_flag()? { + transform_8x8 = true; + info.kind = MbKind::I8x8; + self.cur.info.kind = MbKind::I8x8; + } + self.mb_mut().kind = info.kind; + self.mb_mut().transform_8x8 = transform_8x8; + let n_ref = [self.sh.num_ref_idx_active[0], self.sh.num_ref_idx_active[1]]; + match info.kind { + MbKind::I4x4 | MbKind::I8x8 => { + let n = if info.kind == MbKind::I8x8 { 4 } else { 16 }; + for i in 0..n { + self.cur.rem_mode[i] = if r.read_flag()? { -1 } else { r.read_bits(3)? as i8 }; + } + self.derive_intra_modes(info.kind == MbKind::I8x8); + let m = r.read_ue()?; + ensure!(m <= 3, "intra_chroma_pred_mode out of range"); + self.mb_mut().intra_chroma_mode = m as u8; + } + MbKind::I16x16 => { + let m = r.read_ue()?; + ensure!(m <= 3, "intra_chroma_pred_mode out of range"); + self.mb_mut().intra_chroma_mode = m as u8; + } + MbKind::BDirect16x16 => {} + _ => { + if info.part == Part::P8x8 { + for p in 0..4 { + let t = r.read_ue()?; + let sub = if st.is_b() { b_sub_mb_type(t) } else { p_sub_mb_type(t) }; + let Some(sub) = sub else { return invalid(format!("sub_mb_type {t} out of range")) }; + self.cur.sub[p] = sub; + } + for l in 0..2 { + for p in 0..4 { + let sub = self.cur.sub[p]; + if sub.direct || sub.pred & (1 << l) == 0 { + continue; + } + self.cur.ref_idx[l][p] = if n_ref[l] > 1 && !info.ref0 { Self::read_te(r, n_ref[l] - 1)? as i8 } else { 0 }; + } + } + for l in 0..2 { + for p in 0..4 { + let sub = self.cur.sub[p]; + if sub.direct || sub.pred & (1 << l) == 0 { + continue; + } + for j in 0..sub.shape.num_parts() { + let x = r.read_se()?; + let y = r.read_se()?; + self.cur.mvd[l][p * 4 + j] = [x as i16, y as i16]; + } + } + } + } else { + let np = info.part.num_parts(); + for l in 0..2 { + for p in 0..np { + if info.pred[p] & (1 << l) != 0 { + self.cur.ref_idx[l][p] = if n_ref[l] > 1 { Self::read_te(r, n_ref[l] - 1)? as i8 } else { 0 }; + } + } + } + for l in 0..2 { + for p in 0..np { + if info.pred[p] & (1 << l) != 0 { + let x = r.read_se()?; + let y = r.read_se()?; + self.cur.mvd[l][p * 4] = [x as i16, y as i16]; + } + } + } + } + } + } + // coded_block_pattern + let cbp = if info.kind == MbKind::I16x16 { + info.i16_cbp + } else { + let code = r.read_ue()?; + ensure!(code < 48, "coded_block_pattern out of range"); + let (i, p) = CBP_ME[code as usize]; + if info.kind.is_intra() { i } else { p } + }; + self.cur.cbp = cbp; + self.mb_mut().cbp = cbp; + if !info.kind.is_intra() && cbp & 15 != 0 && self.pps.transform_8x8_mode && self.no_sub_8x8_lt() { + let t = r.read_flag()?; + self.mb_mut().transform_8x8 = t; + } + if !info.kind.is_intra() { + self.derive_inter_motion()?; + } + if cbp != 0 || info.kind == MbKind::I16x16 { + let d = r.read_se()?; + self.apply_qp_delta(d)?; + } + self.store_qp(); + self.residual_cavlc(r, cbp)?; + if info.kind.is_intra() { + self.reconstruct_intra(); + } else { + self.predict_inter(); + self.reconstruct_inter_residual(); + } + Ok(()) + } + + /// Whether transform_size_8x8_flag may be present for an inter MB (NoSubMbPartSizeLessThan8x8 etc.). + pub fn no_sub_8x8_lt(&self) -> bool { + let info = self.cur.info; + match info.kind { + MbKind::BDirect16x16 => self.sps.direct_8x8_inference, + MbKind::Inter if info.part == Part::P8x8 => { + self.cur.sub.iter().all(|s| if s.direct { self.sps.direct_8x8_inference } else { s.shape == SubPart::S8x8 }) + } + _ => true, + } + } + + fn residual_cavlc(&mut self, r: &mut BitReader, cbp: u8) -> Result<()> { + let kind = self.mb().kind; + let t8 = self.mb().transform_8x8; + if kind == MbKind::I16x16 { + let nc = self.nc_luma(0); + let mut lv = [0i32; 16]; + cavlc::residual_block(r, &mut lv, 0, 15, 16, nc)?; + self.put_luma_dc(&lv); + } + for b8 in 0..4 { + let (bx8, by8) = ((b8 & 1) * 2, (b8 >> 1) * 2); + if cbp & (1 << b8) == 0 { + continue; + } + if t8 { + let mut lv8 = [0i32; 64]; + for i4 in 0..4 { + let raster = (by8 + (i4 >> 1)) * 4 + bx8 + (i4 & 1); + let nc = self.nc_luma(raster); + let mut lv = [0i32; 16]; + let tc = cavlc::residual_block(r, &mut lv, 0, 15, 16, nc)?; + self.mb_mut().nnz[raster] = tc; + for (i, &v) in lv.iter().enumerate() { + lv8[4 * i + i4] = v; + } + } + self.put_luma8(b8, &lv8); + } else { + for i4 in 0..4 { + let raster = (by8 + (i4 >> 1)) * 4 + bx8 + (i4 & 1); + let nc = self.nc_luma(raster); + let mut lv = [0i32; 16]; + let tc = if kind == MbKind::I16x16 { + cavlc::residual_block(r, &mut lv[1..], 0, 14, 15, nc)? + } else { + cavlc::residual_block(r, &mut lv, 0, 15, 16, nc)? + }; + self.mb_mut().nnz[raster] = tc; + self.put_luma4(raster, &lv, kind == MbKind::I16x16); + } + } + } + let cbp_c = cbp >> 4; + if cbp_c & 3 != 0 { + for c in 0..2 { + let mut lv = [0i32; 4]; + cavlc::residual_block(r, &mut lv, 0, 3, 4, -1)?; + self.put_chroma_dc(c, &lv); + } + } + if cbp_c & 2 != 0 { + for c in 0..2 { + for b in 0..4 { + let nc = self.nc_chroma(c, b); + let mut lv = [0i32; 16]; + let tc = cavlc::residual_block(r, &mut lv[1..], 0, 14, 15, nc)?; + self.mb_mut().nnz_c[c][b] = tc; + self.put_chroma_ac(c, b, &lv); + } + } + } + Ok(()) + } + + /// Decode slice_data() with CAVLC. `r` is positioned after the slice header. + pub fn decode_cavlc(&mut self, r: &mut BitReader) -> Result<()> { + let total = self.mb_w * self.mb_h; + let mut addr = self.sh.first_mb_in_slice as usize; + let intra_slice = self.sh.slice_type.is_intra(); + loop { + ensure!(addr < total, "macroblock address {addr} out of range"); + if !intra_slice { + let run = r.read_ue()? as usize; + ensure!(addr + run <= total, "mb_skip_run out of range"); + for _ in 0..run { + self.start_mb(addr); + self.decode_skip()?; + self.pic.decoded_mbs += 1; + addr += 1; + } + if run > 0 && !r.more_rbsp_data() { + break; + } + ensure!(addr < total, "macroblock address {addr} out of range"); + } + self.start_mb(addr); + self.decode_mb_cavlc(r)?; + self.pic.decoded_mbs += 1; + addr += 1; + if !r.more_rbsp_data() { + break; + } + } + Ok(()) + } +} diff --git a/crates/h264/src/transform.rs b/crates/h264/src/transform.rs new file mode 100644 index 00000000..0247f186 --- /dev/null +++ b/crates/h264/src/transform.rs @@ -0,0 +1,228 @@ +//! Scaling (dequantisation) and inverse transforms (8.5.10 - 8.5.13). + +use crate::params::ScalingMatrices; +use crate::tables::{norm_adjust4, norm_adjust8}; + +/// LevelScale4x4 / LevelScale8x8 for all six scaling lists and qP % 6, raster order. +#[derive(Clone)] +pub struct LevelScale { + /// [list 0..6 (IntraY, IntraCb, IntraCr, InterY, InterCb, InterCr)][qp % 6][pos] + pub ls4: [[[i32; 16]; 6]; 6], + /// [list 0..6 (IntraY, InterY, IntraCb, InterCb, IntraCr, InterCr)][qp % 6][pos] + pub ls8: [[[i32; 64]; 6]; 6], +} + +impl LevelScale { + pub fn new(m: &ScalingMatrices) -> Box { + let mut s = Box::new(LevelScale { ls4: [[[0; 16]; 6]; 6], ls8: [[[0; 64]; 6]; 6] }); + for l in 0..6 { + for q in 0..6 { + for p in 0..16 { + s.ls4[l][q][p] = m.m4[l][p] as i32 * norm_adjust4(q, p) as i32; + } + for p in 0..64 { + s.ls8[l][q][p] = m.m8[l][p] as i32 * norm_adjust8(q, p) as i32; + } + } + } + s + } +} + +/// Scale one 4x4 AC coefficient (8.5.12.1, not the DC of Intra16x16 / chroma blocks). +#[inline(always)] +pub fn scale4(c: i32, ls: i32, qp: i32) -> i32 { + if qp >= 24 { (c * ls) << (qp / 6 - 4) } else { (c * ls + (1 << (3 - qp / 6))) >> (4 - qp / 6) } +} + +/// Scale one 8x8 coefficient (8.5.13.1). +#[inline(always)] +pub fn scale8(c: i32, ls: i32, qp: i32) -> i32 { + if qp >= 36 { (c * ls) << (qp / 6 - 6) } else { (c * ls + (1 << (5 - qp / 6))) >> (6 - qp / 6) } +} + +/// Intra16x16 luma DC: inverse Hadamard + scaling (8.5.10). `c` is raster order (row*4+col); +/// output replaces `c` with dcY (raster: row i -> 4x4 block row, col j -> block column). +pub fn luma_dc_dequant(c: &mut [i32; 16], qp: i32, ls00: i32) { + let mut t = [0i32; 16]; + // rows + for i in 0..4 { + let r = &c[i * 4..i * 4 + 4]; + let s01 = r[0] + r[1]; + let d01 = r[0] - r[1]; + let s23 = r[2] + r[3]; + let d23 = r[2] - r[3]; + t[i * 4] = s01 + s23; + t[i * 4 + 1] = s01 - s23; + t[i * 4 + 2] = d01 - d23; + t[i * 4 + 3] = d01 + d23; + } + // columns + for j in 0..4 { + let s01 = t[j] + t[4 + j]; + let d01 = t[j] - t[4 + j]; + let s23 = t[8 + j] + t[12 + j]; + let d23 = t[8 + j] - t[12 + j]; + let f = [s01 + s23, s01 - s23, d01 - d23, d01 + d23]; + for i in 0..4 { + c[i * 4 + j] = if qp >= 36 { (f[i] * ls00) << (qp / 6 - 6) } else { (f[i] * ls00 + (1 << (5 - qp / 6))) >> (6 - qp / 6) }; + } + } +} + +/// 4:2:0 chroma DC: 2x2 transform + scaling (8.5.11). `c` = [c00, c01, c10, c11]. +pub fn chroma_dc_dequant_420(c: &mut [i32; 4], qp: i32, ls00: i32) { + let s0 = c[0] + c[1]; + let d0 = c[0] - c[1]; + let s1 = c[2] + c[3]; + let d1 = c[2] - c[3]; + let f = [s0 + s1, d0 + d1, s0 - s1, d0 - d1]; + for k in 0..4 { + c[k] = ((f[k] * ls00) << (qp / 6)) >> 5; + } +} + +#[inline(always)] +fn clip_u8(v: i32) -> u8 { + v.clamp(0, 255) as u8 +} + +/// Inverse 4x4 transform of raster-order `d` and add to `dst` (8.5.12.2, 8.5.14). +pub fn idct4_add(d: &[i32; 16], dst: &mut [u8], stride: usize) { + let mut t = [0i32; 16]; + for i in 0..4 { + let r = &d[i * 4..i * 4 + 4]; + let e = r[0] + r[2]; + let f = r[0] - r[2]; + let g = (r[1] >> 1) - r[3]; + let h = r[1] + (r[3] >> 1); + t[i * 4] = e + h; + t[i * 4 + 1] = f + g; + t[i * 4 + 2] = f - g; + t[i * 4 + 3] = e - h; + } + for j in 0..4 { + let e = t[j] + t[8 + j]; + let f = t[j] - t[8 + j]; + let g = (t[4 + j] >> 1) - t[12 + j]; + let h = t[4 + j] + (t[12 + j] >> 1); + let col = [e + h, f + g, f - g, e - h]; + for i in 0..4 { + let p = &mut dst[i * stride + j]; + *p = clip_u8(*p as i32 + ((col[i] + 32) >> 6)); + } + } +} + +/// DC-only shortcut: all AC coefficients zero. +#[allow(dead_code)] +pub fn idct4_dc_add(dc: i32, dst: &mut [u8], stride: usize) { + let v = (dc + 32) >> 6; + for i in 0..4 { + for p in &mut dst[i * stride..i * stride + 4] { + *p = clip_u8(*p as i32 + v); + } + } +} + +#[inline(always)] +fn idct8_1d(d: [i32; 8]) -> [i32; 8] { + let a0 = d[0] + d[4]; + let a4 = d[0] - d[4]; + let a2 = (d[2] >> 1) - d[6]; + let a6 = d[2] + (d[6] >> 1); + let b0 = a0 + a6; + let b2 = a4 + a2; + let b4 = a4 - a2; + let b6 = a0 - a6; + let a1 = -d[3] + d[5] - d[7] - (d[7] >> 1); + let a3 = d[1] + d[7] - d[3] - (d[3] >> 1); + let a5 = -d[1] + d[7] + d[5] + (d[5] >> 1); + let a7 = d[3] + d[5] + d[1] + (d[1] >> 1); + let b1 = a1 + (a7 >> 2); + let b7 = a7 - (a1 >> 2); + let b3 = a3 + (a5 >> 2); + let b5 = (a3 >> 2) - a5; + [b0 + b7, b2 + b5, b4 + b3, b6 + b1, b6 - b1, b4 - b3, b2 - b5, b0 - b7] +} + +/// Inverse 8x8 transform of raster-order `d` and add to `dst` (8.5.13.2). +pub fn idct8_add(d: &[i32; 64], dst: &mut [u8], stride: usize) { + let mut t = [0i32; 64]; + for i in 0..8 { + let row: [i32; 8] = d[i * 8..i * 8 + 8].try_into().unwrap(); + t[i * 8..i * 8 + 8].copy_from_slice(&idct8_1d(row)); + } + for j in 0..8 { + let col = idct8_1d(std::array::from_fn(|i| t[i * 8 + j])); + for i in 0..8 { + let p = &mut dst[i * stride + j]; + *p = clip_u8(*p as i32 + ((col[i] + 32) >> 6)); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn idct4_dc_only_matches_shortcut() { + let mut d = [0i32; 16]; + d[0] = 640; // -> +10 everywhere + let mut a = [100u8; 16]; + let mut b = [100u8; 16]; + idct4_add(&d, &mut a, 4); + idct4_dc_add(640, &mut b, 4); + assert_eq!(a, b); + assert!(a.iter().all(|&v| v == 110)); + } + + #[test] + fn idct4_hand_example() { + // d = [[64, 64, 0, 0], 0...]: row0: e=64,f=64,g=32,h=64 -> [128, 96, 32, 0]; columns: each col only row0 set + // col j: e=f=t0j, g=0, h=0 -> all rows = t0j. Result (t+32)>>6 = [2, 2, 1, 0] for every row. + let mut d = [0i32; 16]; + d[0] = 64; + d[1] = 64; + let mut out = [0u8; 16]; + idct4_add(&d, &mut out, 4); + for i in 0..4 { + assert_eq!(&out[i * 4..i * 4 + 4], &[2, 2, 1, 0]); + } + } + + #[test] + fn idct8_dc() { + let mut d = [0i32; 64]; + d[0] = 64 * 5; + let mut out = [10u8; 64]; + idct8_add(&d, &mut out, 8); + assert!(out.iter().all(|&v| v == 15)); + } + + #[test] + fn luma_dc_flat() { + // single DC coefficient spreads equally to all 16 blocks + let mut c = [0i32; 16]; + c[0] = 4; + luma_dc_dequant(&mut c, 36, 16 * 16); + assert!(c.iter().all(|&v| v == 4 * 256)); + } + + #[test] + fn chroma_dc() { + let mut c = [1, 0, 0, 0]; + chroma_dc_dequant_420(&mut c, 6, 16 * 16); + // f = [1,1,1,1] -> ((1*256) << 1) >> 5 = 16 + assert_eq!(c, [16, 16, 16, 16]); + } + + #[test] + fn scale_functions() { + assert_eq!(scale4(2, 160, 24), 2 * 160); + assert_eq!(scale4(2, 160, 12), (2 * 160 + 2) >> 2); + assert_eq!(scale8(1, 320, 36), 320); + assert_eq!(scale8(1, 320, 30), (320 + 1) >> 1); + } +} diff --git a/crates/h264/tests/common/mod.rs b/crates/h264/tests/common/mod.rs index 63206541..7b51b4cb 100644 --- a/crates/h264/tests/common/mod.rs +++ b/crates/h264/tests/common/mod.rs @@ -1,6 +1,7 @@ //! Test fixture generation (ffmpeg/libx264 as an external oracle) and comparison helpers. #![allow(dead_code)] +use filmcraft_h264::{Decoder, Picture}; use std::path::{Path, PathBuf}; use std::process::Command; @@ -364,3 +365,89 @@ pub fn split_access_units(data: &[u8]) -> Vec<&[u8]> { } aus } + +/// Decode a whole Annex-B file access unit by access unit (pts = AU index). +/// Failed decode: (access unit index, error, pictures output before the error). +pub type DecodeFailure = (usize, filmcraft_h264::Error, Vec); + +pub fn decode_file(path: &Path) -> Result, DecodeFailure> { + let data = std::fs::read(path).unwrap(); + let mut dec = Decoder::new(); + let mut out = Vec::new(); + for (i, au) in split_access_units(&data).into_iter().enumerate() { + match dec.decode(au, i as i64) { + Ok(p) => out.extend(p), + Err(e) => return Err((i, e, out)), + } + } + out.extend(dec.flush()); + Ok(out) +} + +/// Compare decoded pictures with a raw yuv420p reference; returns a diagnostic on the first mismatch. +pub fn compare(pics: &[Picture], reference: &[u8], w: usize, h: usize) -> Result<(), String> { + let (cw, ch) = (w.div_ceil(2), h.div_ceil(2)); + let fsize = w * h + 2 * cw * ch; + let nref = reference.len() / fsize; + for (i, p) in pics.iter().enumerate() { + if i >= nref { + return Err(format!("decoder produced {} frames, reference has {}", pics.len(), nref)); + } + if p.width as usize != w || p.height as usize != h { + return Err(format!("frame {i}: size {}x{} != {}x{}", p.width, p.height, w, h)); + } + let f = &reference[i * fsize..(i + 1) * fsize]; + let planes = + [("Y", &p.y[..], &f[..w * h], w, h, 16), ("U", &p.u[..], &f[w * h..w * h + cw * ch], cw, ch, 8), ("V", &p.v[..], &f[w * h + cw * ch..], cw, ch, 8)]; + for (name, got, exp, pw, ph, mbs) in planes { + if got != exp { + let mut first = None; + let mut count = 0usize; + let mut maxd = 0i32; + for y in 0..ph { + for x in 0..pw { + let (a, b) = (got[y * pw + x], exp[y * pw + x]); + if a != b { + count += 1; + maxd = maxd.max((a as i32 - b as i32).abs()); + if first.is_none() { + first = Some((x, y, a, b)); + } + } + } + } + let (x, y, a, b) = first.unwrap(); + let mb_w = pw.div_ceil(mbs); + return Err(format!( + "frame {i} (poc {}, pts {}) plane {name}: first mismatch at ({x},{y}) MB {} (mb_x {}, mb_y {}): got {a} want {b}; {count} samples differ, max diff {maxd}", + p.poc, + p.pts, + (y / mbs) * mb_w + x / mbs, + x / mbs, + y / mbs + )); + } + } + } + if pics.len() != nref { + return Err(format!("decoder produced {} frames, reference has {}", pics.len(), nref)); + } + Ok(()) +} + +/// Full check of one fixture: generate, decode, compare. Returns Ok(false) when skipped. +pub fn check_fixture(name: &str) -> Result { + let f = fixture(name); + let Some((h264, yuv)) = ensure(f) else { return Ok(false) }; + let reference = std::fs::read(&yuv).unwrap(); + let pics = match decode_file(&h264) { + Ok(p) => p, + Err((au, e, partial)) => { + // still report comparison of what was decoded + let cmp = compare(&partial, &reference, f.width as usize, f.height as usize).err().unwrap_or_default(); + return Err(format!("{name}: decode error at access unit {au}: {e} (after {} pictures) {cmp}", partial.len())); + } + }; + compare(&pics, &reference, f.width as usize, f.height as usize).map_err(|e| format!("{name}: {e}"))?; + Ok(true) +} diff --git a/crates/h264/tests/conformance.rs b/crates/h264/tests/conformance.rs new file mode 100644 index 00000000..bf06c310 --- /dev/null +++ b/crates/h264/tests/conformance.rs @@ -0,0 +1,20 @@ +//! Bit-exactness against ffmpeg's decoder on libx264-generated fixtures (skipped when ffmpeg is absent). + +mod common; + +macro_rules! fixture_tests { + ($($name:ident),* $(,)?) => { + $( + #[test] + fn $name() { + match common::check_fixture(stringify!($name)) { + Ok(true) => {} + Ok(false) => eprintln!("skipped {}", stringify!($name)), + Err(e) => panic!("{e}"), + } + } + )* + }; +} + +fixture_tests!(intra_cavlc, intra_cavlc_noise, intra_cavlc_8x8, intra_cavlc_cqm, p_cavlc_nodeblock,);