diff --git a/crates/h264/src/cabac_mb.rs b/crates/h264/src/cabac_mb.rs --- a/crates/h264/src/cabac_mb.rs +++ b/crates/h264/src/cabac_mb.rs @@ -7,6 +7,7 @@ use crate::mbtypes::*; use crate::picture::MbKind; use crate::slice::SliceType; use crate::slicedec::{MbCur, SliceDecoder}; +use crate::tables::{ZIGZAG4, ZIGZAG8}; /// Table 9-43: significant_coeff_flag ctxIdxInc for 8x8 frame blocks. #[rustfmt::skip] @@ -375,16 +376,17 @@ impl SliceDecoder<'_> { (st.nnz_c[comp][y * 2 + x] != 0) as usize } - /// residual_block_cabac for ctxBlockCat `cat` (0..=5). Coefficients are written to `coeff` - /// in list order (indices 0..max_num). `cbf_inc` is the coded_block_flag ctxIdxInc (ignored for cat 5). - /// Returns the number of non-zero coefficients. - fn residual_block_cabac(&mut self, c: &mut Cabac, coeff: &mut [i32], cat: usize, max_num: usize, cbf_inc: usize) -> u8 { - if cat != 5 && c.decode_decision(85 + CAT_OFFSET[cat].0 + cbf_inc) == 0 { + /// residual_block_cabac for ctxBlockCat `CAT` (0..=5). Each non-zero coefficient is passed to + /// `put` as (index in the coefficient list, level); the others are zero. `cbf_inc` is the + /// coded_block_flag ctxIdxInc (ignored for cat 5). Returns the number of non-zero coefficients. + fn residual_block_cabac(c: &mut Cabac, cbf_inc: usize, mut put: impl FnMut(usize, i32)) -> u8 { + if CAT != 5 && c.decode_decision(85 + CAT_OFFSET[CAT].0 + cbf_inc) == 0 { return 0; } + let max_num = [16, 15, 16, 4, 15, 64][CAT]; let (sig_base, last_base, abs_base) = - if cat == 5 { (402, 417, 426) } else { (105 + CAT_OFFSET[cat].1, 166 + CAT_OFFSET[cat].1, 227 + CAT_OFFSET[cat].2) }; - let (sig_tab, last_tab): (&[u8], &[u8]) = match cat { + if CAT == 5 { (402, 417, 426) } else { (105 + CAT_OFFSET[CAT].1, 166 + CAT_OFFSET[CAT].1, 227 + CAT_OFFSET[CAT].2) }; + let (sig_tab, last_tab): (&[u8], &[u8]) = match CAT { 3 => (&CHROMA_DC_CTX, &CHROMA_DC_CTX), 5 => (&SIG8_FRAME, &LAST8), _ => (&IDENT_CTX, &IDENT_CTX), @@ -410,7 +412,7 @@ impl SliceDecoder<'_> { } let mut eq1 = 0usize; let mut gt1 = 0usize; - let gt1_cap = if cat == 3 { 3 } else { 4 }; + let gt1_cap = if CAT == 3 { 3 } else { 4 }; for &p in pos[..n].iter().rev() { let inc0 = if gt1 != 0 { 0 } else { (1 + eq1).min(4) }; let mut abs = 1i32; @@ -441,7 +443,7 @@ impl SliceDecoder<'_> { } else { eq1 += 1; } - coeff[p as usize] = if c.decode_bypass() == 1 { -abs } else { abs }; + put(p as usize, if c.decode_bypass() == 1 { -abs } else { abs }); } n as u8 } @@ -453,7 +455,7 @@ impl SliceDecoder<'_> { if kind == MbKind::I16x16 { let inc = self.cbf_dc_cond(self.nb[0], 0, intra) + 2 * self.cbf_dc_cond(self.nb[1], 0, intra); let mut lv = [0i32; 16]; - if self.residual_block_cabac(c, &mut lv, 0, 16, inc) > 0 { + if Self::residual_block_cabac::<0>(c, inc, |k, v| lv[k] = v) > 0 { self.mb_mut().cbf_dc |= 1; self.put_luma_dc(&lv); } @@ -464,27 +466,36 @@ impl SliceDecoder<'_> { continue; } if t8 { - let mut lv = [0i32; 64]; - self.residual_block_cabac(c, &mut lv, 5, 64, 0); + let (blk, ls, dq) = self.luma8_dst(b8); + Self::residual_block_cabac::<5>(c, 0, |k, v| { + let p = ZIGZAG8[k] as usize; + blk[p] = dq.apply(v, ls[p]); + }); + self.mark_luma8(b8); let st = self.mb_mut(); for r in [by8 * 4 + bx8, by8 * 4 + bx8 + 1, (by8 + 1) * 4 + bx8, (by8 + 1) * 4 + bx8 + 1] { st.nnz[r] = 1; } - self.put_luma8(b8, &lv); } else { for i4 in 0..4 { let (bx, by) = (bx8 + (i4 & 1), by8 + (i4 >> 1)); let raster = by * 4 + bx; let inc = self.cbf_luma_cond(bx as i32 - 1, by as i32, intra) + 2 * self.cbf_luma_cond(bx as i32, by as i32 - 1, intra); - let mut lv = [0i32; 16]; + let (blk, ls, dq) = self.luma4_dst(raster); let n = if kind == MbKind::I16x16 { - self.residual_block_cabac(c, &mut lv[1..], 1, 15, inc) + Self::residual_block_cabac::<1>(c, inc, |k, v| { + let p = ZIGZAG4[k + 1] as usize; + blk[p] = dq.apply(v, ls[p]); + }) } else { - self.residual_block_cabac(c, &mut lv, 2, 16, inc) + Self::residual_block_cabac::<2>(c, inc, |k, v| { + let p = ZIGZAG4[k] as usize; + blk[p] = dq.apply(v, ls[p]); + }) }; self.mb_mut().nnz[raster] = n; if n > 0 { - self.put_luma4(raster, &lv, kind == MbKind::I16x16); + self.mark_luma4(raster); } } } @@ -495,7 +506,7 @@ impl SliceDecoder<'_> { let bit = 1 + comp as u8; let inc = self.cbf_dc_cond(self.nb[0], bit, intra) + 2 * self.cbf_dc_cond(self.nb[1], bit, intra); let mut lv = [0i32; 4]; - if self.residual_block_cabac(c, &mut lv, 3, 4, inc) > 0 { + if Self::residual_block_cabac::<3>(c, inc, |k, v| lv[k] = v) > 0 { self.mb_mut().cbf_dc |= 1 << bit; self.put_chroma_dc(comp, &lv); } @@ -506,11 +517,14 @@ impl SliceDecoder<'_> { for b in 0..4 { let (bx, by) = ((b & 1) as i32, (b >> 1) as i32); let inc = self.cbf_chroma_ac_cond(comp, bx - 1, by, intra) + 2 * self.cbf_chroma_ac_cond(comp, bx, by - 1, intra); - let mut lv = [0i32; 16]; - let n = self.residual_block_cabac(c, &mut lv[1..], 4, 15, inc); + let (blk, ls, dq) = self.chroma_ac_dst(comp, b); + let n = Self::residual_block_cabac::<4>(c, inc, |k, v| { + let p = ZIGZAG4[k + 1] as usize; + blk[p] = dq.apply(v, ls[p]); + }); self.mb_mut().nnz_c[comp][b] = n; if n > 0 { - self.put_chroma_ac(comp, b, &lv); + self.mark_chroma_ac(comp, b); } } } diff --git a/crates/h264/src/cavlc.rs b/crates/h264/src/cavlc.rs --- a/crates/h264/src/cavlc.rs +++ b/crates/h264/src/cavlc.rs @@ -89,10 +89,17 @@ 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). +/// Parse residual_block_cavlc(). Each non-zero coefficient is passed to `put` as (scan index in +/// `start_idx..=end_idx`, level); the others are zero. /// `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 { +pub(crate) fn residual_block( + r: &mut BitReader, + start_idx: usize, + end_idx: usize, + max_num_coeff: usize, + nc: i32, + mut put: impl FnMut(usize, i32), +) -> Result { let t = tables(); let col = match nc { -1 => 4, @@ -162,7 +169,7 @@ pub(crate) fn residual_block(r: &mut BitReader, coeff: &mut [i32], start_idx: us // 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; + put(pos, lv); if i + 1 == total_coeff { break; } @@ -274,7 +281,7 @@ mod tests { 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(); + let tc = residual_block(&mut r, 0, 15, 16, 0, |k, v| c[k] = v).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/slicedec.rs b/crates/h264/src/slicedec.rs --- a/crates/h264/src/slicedec.rs +++ b/crates/h264/src/slicedec.rs @@ -10,7 +10,7 @@ use crate::params::{Pps, Sps}; use crate::picture::*; use crate::slice::{SliceHeader, SliceType}; use crate::tables::*; -use crate::transform::{self, LevelScale}; +use crate::transform::{self, Dequant, LevelScale}; use filmcraft_bitstream::BitReader; /// Per-slice data kept with the picture (deblocking, motion field export). @@ -1076,45 +1076,36 @@ impl<'a> SliceDecoder<'a> { 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) { + /// Coefficients, level scales and dequantiser of luma 4x4 block `raster`: the residual + /// decoders write each coefficient straight to its raster position. + pub fn luma4_dst(&mut self, raster: usize) -> (&mut [i32; 16], &'a [i32; 16], Dequant) { 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; - } + let blk = &mut self.s.coef.as_chunks_mut::<16>().0[raster]; + (blk, ls, Dequant::new4(qp)) + } + + /// Luma 4x4 block `raster` has coefficients. + pub fn mark_luma4(&mut self, raster: usize) { + 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]) { + /// `luma4_dst` for luma 8x8 block `b8`. + pub fn luma8_dst(&mut self, b8: usize) -> (&mut [i32; 64], &'a [i32; 64], Dequant) { 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; - } + let blk = &mut self.s.coef.as_chunks_mut::<64>().0[b8]; + (blk, ls, Dequant::new8(qp)) + } + + /// Luma 8x8 block `b8` has coefficients. + pub fn mark_luma8(&mut self, b8: usize) { + 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. @@ -1149,23 +1140,18 @@ impl<'a> SliceDecoder<'a> { } } - /// 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]) { + /// `luma4_dst` for chroma AC block `b` of component `c`. + pub fn chroma_ac_dst(&mut self, c: usize, b: usize) -> (&mut [i32; 16], &'a [i32; 16], Dequant) { 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; - } + let blk = &mut self.s.coef_c[c].as_chunks_mut::<16>().0[b]; + (blk, ls, Dequant::new4(qpc)) + } + + /// Chroma AC block `b` of component `c` has coefficients. + pub fn mark_chroma_ac(&mut self, c: usize, b: usize) { + self.s.c_nz[c] |= 1 << b; } // ------------------------------------------------------------------------------------------ @@ -1358,7 +1344,7 @@ impl<'a> SliceDecoder<'a> { 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)?; + cavlc::residual_block(r, 0, 15, 16, nc, |k, v| lv[k] = v)?; self.put_luma_dc(&lv); } for b8 in 0..4 { @@ -1367,30 +1353,41 @@ impl<'a> SliceDecoder<'a> { continue; } if t8 { - let mut lv8 = [0i32; 64]; + let mut any = false; 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)?; + let (blk, ls, dq) = self.luma8_dst(b8); + let tc = cavlc::residual_block(r, 0, 15, 16, nc, |k, v| { + let p = ZIGZAG8[4 * k + i4] as usize; + blk[p] = dq.apply(v, ls[p]); + })?; self.mb_mut().nnz[raster] = tc; - for (i, &v) in lv.iter().enumerate() { - lv8[4 * i + i4] = v; - } + any |= tc > 0; + } + if any { + self.mark_luma8(b8); } - 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 (blk, ls, dq) = self.luma4_dst(raster); let tc = if kind == MbKind::I16x16 { - cavlc::residual_block(r, &mut lv[1..], 0, 14, 15, nc)? + cavlc::residual_block(r, 0, 14, 15, nc, |k, v| { + let p = ZIGZAG4[k + 1] as usize; + blk[p] = dq.apply(v, ls[p]); + })? } else { - cavlc::residual_block(r, &mut lv, 0, 15, 16, nc)? + cavlc::residual_block(r, 0, 15, 16, nc, |k, v| { + let p = ZIGZAG4[k] as usize; + blk[p] = dq.apply(v, ls[p]); + })? }; self.mb_mut().nnz[raster] = tc; - self.put_luma4(raster, &lv, kind == MbKind::I16x16); + if tc > 0 { + self.mark_luma4(raster); + } } } } @@ -1398,7 +1395,7 @@ impl<'a> SliceDecoder<'a> { 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)?; + cavlc::residual_block(r, 0, 3, 4, -1, |k, v| lv[k] = v)?; self.put_chroma_dc(c, &lv); } } @@ -1406,10 +1403,15 @@ impl<'a> SliceDecoder<'a> { 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)?; + let (blk, ls, dq) = self.chroma_ac_dst(c, b); + let tc = cavlc::residual_block(r, 0, 14, 15, nc, |k, v| { + let p = ZIGZAG4[k + 1] as usize; + blk[p] = dq.apply(v, ls[p]); + })?; self.mb_mut().nnz_c[c][b] = tc; - self.put_chroma_ac(c, b, &lv); + if tc > 0 { + self.mark_chroma_ac(c, b); + } } } } diff --git a/crates/h264/src/transform.rs b/crates/h264/src/transform.rs --- a/crates/h264/src/transform.rs +++ b/crates/h264/src/transform.rs @@ -29,16 +29,34 @@ impl LevelScale { } } -/// 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.wrapping_mul(ls) << (qp / 6 - 4) } else { c.wrapping_mul(ls).wrapping_add(1 << (3 - qp / 6)) >> (4 - qp / 6) } +/// Coefficient scaling of one block (8.5.12.1 for 4x4 AC, 8.5.13.1 for 8x8) with the qP-dependent +/// shift and rounding resolved once per block. +#[derive(Clone, Copy)] +pub struct Dequant { + shl: u32, + shr: u32, + round: i32, } -/// Scale one 8x8 coefficient (8.5.13.1). -#[inline(always)] -pub fn scale8(c: i32, ls: i32, qp: i32) -> i32 { - if qp >= 36 { c.wrapping_mul(ls) << (qp / 6 - 6) } else { c.wrapping_mul(ls).wrapping_add(1 << (5 - qp / 6)) >> (6 - qp / 6) } +impl Dequant { + /// 4x4 blocks (not the DC of Intra16x16 / chroma blocks). + pub fn new4(qp: i32) -> Self { + Self::new(qp / 6 - 4) + } + + /// 8x8 blocks. + pub fn new8(qp: i32) -> Self { + Self::new(qp / 6 - 6) + } + + fn new(shift: i32) -> Self { + if shift >= 0 { Dequant { shl: shift as u32, shr: 0, round: 0 } } else { Dequant { shl: 0, shr: -shift as u32, round: 1 << (-shift - 1) } } + } + + #[inline(always)] + pub fn apply(self, c: i32, ls: i32) -> i32 { + (c.wrapping_mul(ls) << self.shl).wrapping_add(self.round) >> self.shr + } } /// Intra16x16 luma DC: inverse Hadamard + scaling (8.5.10). `c` is raster order (row*4+col); @@ -283,10 +301,26 @@ mod tests { #[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); + assert_eq!(Dequant::new4(24).apply(2, 160), 2 * 160); + assert_eq!(Dequant::new4(12).apply(2, 160), (2 * 160 + 2) >> 2); + assert_eq!(Dequant::new8(36).apply(1, 320), 320); + assert_eq!(Dequant::new8(30).apply(1, 320), (320 + 1) >> 1); + } + + /// The 8.5.12.1 / 8.5.13.1 formulas, per coefficient. + #[test] + fn dequant_matches_spec_formulas() { + for qp in 0..52 { + let (d4, d8) = (Dequant::new4(qp), Dequant::new8(qp)); + for c in [-2047, -300, -17, -1, 1, 2, 5, 64, 1000, 2047] { + for ls in [10, 16, 25, 160, 288, 2880, 18 * 255] { + let want4 = if qp >= 24 { (c * ls) << (qp / 6 - 4) } else { (c * ls + (1 << (3 - qp / 6))) >> (4 - qp / 6) }; + let want8 = if qp >= 36 { (c * ls) << (qp / 6 - 6) } else { (c * ls + (1 << (5 - qp / 6))) >> (6 - qp / 6) }; + assert_eq!(d4.apply(c, ls), want4, "4x4 qp {qp} c {c} ls {ls}"); + assert_eq!(d8.apply(c, ls), want8, "8x8 qp {qp} c {c} ls {ls}"); + } + } + } } /// Corrupt streams can carry huge coefficient levels: dequantisation wraps (as release @@ -296,9 +330,9 @@ mod tests { fn huge_levels_dequantise_without_overflow_panics() { let mut c = [i32::MAX / 2, i32::MAX / 3, -i32::MAX / 2, 7]; chroma_dc_dequant_420(&mut c, 51, 16 * 25); - let _ = scale4(i32::MAX, 16 * 25, 51); - let _ = scale4(i32::MAX, 16 * 25, 0); - let _ = scale8(i32::MIN, 16 * 25, 51); - let _ = scale8(i32::MIN, 16 * 25, 0); + let _ = Dequant::new4(51).apply(i32::MAX, 16 * 25); + let _ = Dequant::new4(0).apply(i32::MAX, 16 * 25); + let _ = Dequant::new8(51).apply(i32::MIN, 16 * 25); + let _ = Dequant::new8(0).apply(i32::MIN, 16 * 25); } }