diff --git a/crates/hevcenc/src/bitstream.rs b/crates/hevcenc/src/bitstream.rs new file mode 100644 index 000000000..27ed3063e --- /dev/null +++ b/crates/hevcenc/src/bitstream.rs @@ -0,0 +1,135 @@ +//! RBSP bit writer (7.2, 9.2 Exp-Golomb) and NAL unit encapsulation with emulation prevention (7.3.1). + +/// MSB-first bit writer. +#[derive(Default)] +pub struct BitWriter { + buf: Vec, + acc: u64, + nbits: u32, +} + +impl BitWriter { + pub fn new() -> Self { + Self::default() + } + + /// Write the `n` (<= 32) low bits of `v`. + #[inline] + pub fn put(&mut self, v: u32, n: u32) { + debug_assert!(n <= 32); + if n == 0 { + return; + } + let v = (v as u64) & ((1u64 << n) - 1); + self.acc = (self.acc << n) | v; + self.nbits += n; + while self.nbits >= 8 { + self.nbits -= 8; + self.buf.push((self.acc >> self.nbits) as u8); + } + self.acc &= (1u64 << self.nbits) - 1; + } + + #[inline] + pub fn bit(&mut self, b: u32) { + self.put(b & 1, 1); + } + + pub fn flag(&mut self, b: bool) { + self.put(b as u32, 1); + } + + /// ue(v). + pub fn ue(&mut self, v: u32) { + let x = v as u64 + 1; + let len = 64 - x.leading_zeros(); // bits of x + // len - 1 leading zeros, then x in len bits + self.put(0, len - 1); + if len > 32 { + self.put((x >> 32) as u32, len - 32); + self.put(x as u32, 32); + } else { + self.put(x as u32, len); + } + } + + /// se(v). + pub fn se(&mut self, v: i32) { + let k = if v > 0 { 2 * v as u32 - 1 } else { (-(v as i64) as u32) * 2 }; + self.ue(k); + } + + pub fn is_aligned(&self) -> bool { + self.nbits == 0 + } + + /// rbsp_trailing_bits(): stop bit and zero alignment. + pub fn trailing(&mut self) { + self.bit(1); + self.align_zero(); + } + + pub fn align_zero(&mut self) { + if self.nbits > 0 { + self.put(0, 8 - self.nbits); + } + } + + pub fn into_bytes(mut self) -> Vec { + self.align_zero(); + self.buf + } +} + +/// Build a NAL unit (2-byte header, `nuh_layer_id` 0, `TemporalId` 0) from an RBSP, inserting +/// emulation prevention bytes. +pub fn nal_unit(nal_type: u8, rbsp: &[u8]) -> Vec { + let mut out = Vec::with_capacity(rbsp.len() + rbsp.len() / 64 + 4); + out.push(nal_type << 1); + out.push(1); + let mut zeros = 0; + for &b in rbsp { + if zeros >= 2 && b <= 3 { + out.push(3); + zeros = 0; + } + out.push(b); + if b == 0 { + zeros += 1; + } else { + zeros = 0; + } + } + // cabac_zero_words / trailing zero bytes never occur: every RBSP ends with the stop bit. + out +} + +pub const NAL_TRAIL_R: u8 = 1; +pub const NAL_IDR_N_LP: u8 = 20; +pub const NAL_VPS: u8 = 32; +pub const NAL_SPS: u8 = 33; +pub const NAL_PPS: u8 = 34; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn exp_golomb() { + let mut w = BitWriter::new(); + w.ue(0); // 1 + w.ue(1); // 010 + w.ue(2); // 011 + w.ue(3); // 00100 + w.se(-1); // ue(2) = 011 + w.trailing(); + // 1 010 011 00100 011 + stop bit: 1010 0110 0100 0111 + assert_eq!(w.into_bytes(), vec![0xa6, 0x47]); + } + + #[test] + fn emulation_prevention() { + let n = nal_unit(1, &[0, 0, 1, 0, 0, 0, 0, 0, 4]); + assert_eq!(&n[2..], &[0, 0, 3, 1, 0, 0, 3, 0, 0, 3, 0, 4]); + } +} diff --git a/crates/hevcenc/src/cabac.rs b/crates/hevcenc/src/cabac.rs new file mode 100644 index 000000000..220b999cd --- /dev/null +++ b/crates/hevcenc/src/cabac.rs @@ -0,0 +1,220 @@ +//! CABAC: context initialisation (9.3.2.2), the arithmetic encoder (9.3.4.x, informative encoder +//! description: EncodeDecision, EncodeBypass, EncodeTerminate, EncodeFlush) and a fractional-bit +//! estimator with the same context-state evolution, used for rate-distortion decisions. + +use crate::bitstream::BitWriter; +use crate::tables::{CABAC_INIT, NEXT_STATE, NUM_CTX, RANGE_TAB_LPS}; +use std::sync::OnceLock; + +/// Context states, each `pStateIdx << 1 | valMps`. +pub type Contexts = [u8; NUM_CTX]; + +/// Initialise all context variables for `slice_qp` and `init_type` (9.3.2.2). +pub fn init_contexts(slice_qp: i32, init_type: usize) -> Contexts { + let mut ctx = [0u8; NUM_CTX]; + let qp = slice_qp.clamp(0, 51); + for (c, &v) in ctx.iter_mut().zip(CABAC_INIT[init_type].iter()) { + let slope = (v >> 4) as i32; + let offset = (v & 15) as i32; + let m = slope * 5 - 45; + let n = (offset << 3) - 16; + let pre = (((m * qp) >> 4) + n).clamp(1, 126); + *c = if pre <= 63 { ((63 - pre) << 1) as u8 } else { (((pre - 64) << 1) | 1) as u8 }; + } + ctx +} + +/// Destination of binarised syntax elements: the real arithmetic coder or the bit estimator. +pub trait Sink { + fn decision(&mut self, ctx: usize, bin: u32); + fn bypass(&mut self, bin: u32); + /// `n` bypass bins carrying the low `n` bits of `v`, MSB first. + fn bypass_bits(&mut self, v: u32, n: u32) { + for i in (0..n).rev() { + self.bypass((v >> i) & 1); + } + } + fn terminate(&mut self, bin: u32); +} + +/// Arithmetic encoder writing into a bit writer. +pub struct CabacEncoder { + pub ctx: Contexts, + low: u32, + range: u32, + outstanding: u32, + first: bool, + pub w: BitWriter, +} + +impl CabacEncoder { + /// Start encoding slice data into `w` (which must be byte aligned). + pub fn new(w: BitWriter, ctx: Contexts) -> Self { + debug_assert!(w.is_aligned()); + CabacEncoder { ctx, low: 0, range: 510, outstanding: 0, first: true, w } + } + + #[inline] + fn put_bit(&mut self, b: u32) { + if self.first { + self.first = false; + } else { + self.w.bit(b); + } + while self.outstanding > 0 { + self.w.bit(1 - b); + self.outstanding -= 1; + } + } + + #[inline] + fn renorm(&mut self) { + while self.range < 256 { + if self.low < 256 { + self.put_bit(0); + } else if self.low >= 512 { + self.low -= 512; + self.put_bit(1); + } else { + self.low -= 256; + self.outstanding += 1; + } + self.range <<= 1; + self.low <<= 1; + } + } + + /// EncodeFlush after a terminating bin equal to 1; the final bit written is the + /// rbsp_stop_one_bit. Returns the writer (byte aligned with zero bits). + pub fn finish(mut self) -> BitWriter { + self.range = 2; + self.renorm(); + let b = (self.low >> 9) & 1; + self.put_bit(b); + self.w.put(((self.low >> 7) & 3) | 1, 2); + self.w.align_zero(); + self.w + } +} + +impl Sink for CabacEncoder { + #[inline] + fn decision(&mut self, ctx_idx: usize, bin: u32) { + let s = self.ctx[ctx_idx] as usize; + let mps = (s & 1) as u32; + let q = ((self.range >> 6) & 3) as usize; + let lps = RANGE_TAB_LPS[s >> 1][q] as u32; + self.range -= lps; + if bin != mps { + self.low += self.range; + self.range = lps; + self.ctx[ctx_idx] = NEXT_STATE[s][1]; + } else { + self.ctx[ctx_idx] = NEXT_STATE[s][0]; + } + self.renorm(); + } + + #[inline] + fn bypass(&mut self, bin: u32) { + self.low <<= 1; + if bin != 0 { + self.low += self.range; + } + if self.low >= 1024 { + self.put_bit(1); + self.low -= 1024; + } else if self.low < 512 { + self.put_bit(0); + } else { + self.low -= 512; + self.outstanding += 1; + } + } + + fn terminate(&mut self, bin: u32) { + self.range -= 2; + if bin != 0 { + self.low += self.range; + // the caller flushes with finish() + } else { + self.renorm(); + } + } +} + +/// Fractional bit cost (1/32768 bit units) of coding the MPS / LPS in each state. +fn cost_table() -> &'static [[u32; 2]; 64] { + static T: OnceLock<[[u32; 2]; 64]> = OnceLock::new(); + T.get_or_init(|| { + let mut t = [[0u32; 2]; 64]; + // pLPS(s) = 0.5 * alpha^s with alpha = (0.01875 / 0.5)^(1/63) (9.3.4.3 design) + let alpha = (0.01875f64 / 0.5).powf(1.0 / 63.0); + for (s, e) in t.iter_mut().enumerate() { + let p = 0.5 * alpha.powi(s as i32); + e[0] = (-(1.0 - p).log2() * 32768.0).round() as u32; + e[1] = (-p.log2() * 32768.0).round() as u32; + } + t + }) +} + +pub const BIT: u64 = 32768; + +/// Bit estimator: tracks context states like the encoder and accumulates the ideal code length. +#[derive(Clone)] +pub struct Estimator { + pub ctx: Contexts, + /// Accumulated cost in 1/32768 bits. + pub bits: u64, + table: &'static [[u32; 2]; 64], +} + +impl Estimator { + pub fn new(ctx: Contexts) -> Self { + Estimator { ctx, bits: 0, table: cost_table() } + } +} + +impl Sink for Estimator { + #[inline] + fn decision(&mut self, ctx_idx: usize, bin: u32) { + let s = self.ctx[ctx_idx] as usize; + let lps = (bin != (s & 1) as u32) as usize; + self.bits += self.table[s >> 1][lps] as u64; + self.ctx[ctx_idx] = NEXT_STATE[s][lps]; + } + #[inline] + fn bypass(&mut self, _bin: u32) { + self.bits += BIT; + } + #[inline] + fn bypass_bits(&mut self, _v: u32, n: u32) { + self.bits += BIT * n as u64; + } + fn terminate(&mut self, _bin: u32) { + self.bits += 7; // ~0 bits for a 0 bin + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn context_init_formula() { + let ctx = init_contexts(30, 0); + // initValue 154: m = 0, n = 64 -> pStateIdx 0, valMps 1 + assert_eq!(ctx[crate::tables::CU_TRANSQUANT_BYPASS], 1); + // split_cu_flag ctx 0 (initValue 139) at QP 30: preCtxState 62 + assert_eq!(ctx[crate::tables::SPLIT_CU], (63 - 62) << 1); + } + + #[test] + fn costs_are_sane() { + let t = cost_table(); + assert_eq!(t[0][0], 32768); + assert_eq!(t[0][1], 32768); + assert!(t[62][0] < 1000 && t[62][1] > 5 * 32768); + } +} diff --git a/crates/hevcenc/src/frame.rs b/crates/hevcenc/src/frame.rs new file mode 100644 index 000000000..68fe3750d --- /dev/null +++ b/crates/hevcenc/src/frame.rs @@ -0,0 +1,1390 @@ +//! Picture coding: CTU loop, rate-distortion mode decisions (with a CABAC bit estimator) and the +//! slice_segment_data() syntax (7.3.8): coding quadtree, coding unit, prediction unit, transform tree +//! and residual coding. Reconstruction mirrors the decoding process exactly (no in-loop filters). + +use crate::cabac::{BIT, CabacEncoder, Estimator, Sink}; +use crate::inter::{self, PlaneRef}; +use crate::intra::{self, Refs}; +use crate::me::MeRef; +use crate::params::{LOG2_CTB, LOG2_MAX_TB, LOG2_MIN_CB, MAX_MERGE_CAND}; +use crate::tables::*; +use crate::transform; + +pub const F_INTRA: u8 = 1; +pub const F_SKIP: u8 = 2; + +/// Per 4x4 luma block state of the current picture. +#[derive(Clone, Copy, Default)] +pub struct Blk { + pub depth: u8, + pub flags: u8, + /// IntraPredModeY (DC for inter blocks). + pub ipm: u8, + /// L0 motion vector (quarter samples) of inter blocks. + pub mv: [i16; 2], +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CuMode { + Intra { nxn: bool, modes: [u8; 4] }, + Skip { idx: u8 }, + Merge { idx: u8 }, + Amvp { mv: [i16; 2], mvp: u8 }, +} + +#[derive(Clone, Copy, Debug)] +struct CuDec { + x: u32, + y: u32, + log2: u32, + mode: CuMode, +} + +/// One coded transform block of the current CU. +#[derive(Clone, Copy)] +struct Tb { + c: u8, + /// Luma-coordinate region the block belongs to (for cbf derivation). + lx: u32, + ly: u32, + log2: u32, + nz: bool, + off: usize, + scan: u8, +} + +/// Saved samples and block state of a square region. +struct Snap { + x: usize, + y: usize, + n: usize, + planes: [Vec; 3], + blk: Vec, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum SliceKind { + I, + P, +} + +pub struct FrameCoder<'a> { + pub bd: u32, + pub w: usize, + pub h: usize, + pub cw: usize, + pub ch: usize, + w4: usize, + wctb: usize, + hctb: usize, + pub kind: SliceKind, + /// QP' (with QpBdOffset) for luma and chroma. + qp_y: i32, + qp_c: i32, + lambda: f64, + lambda_sad: f64, + src: &'a [Vec; 3], + pub rec: [Vec; 3], + refp: Option<&'a [Vec; 3]>, + me: Option<&'a MeRef>, + blk: Vec, + tbs: Vec, + coefs: Vec, + pred: Vec, +} + +#[inline] +fn morton(x: u32, y: u32) -> u32 { + let mut r = 0; + for i in 0..5 { + r |= ((x >> i) & 1) << (2 * i); + r |= ((y >> i) & 1) << (2 * i + 1); + } + r +} + +/// Intra modes tried in the coarse pass of the SATD search. +const COARSE_MODES: [u32; 11] = [0, 1, 2, 6, 10, 14, 18, 22, 26, 30, 34]; + +impl<'a> FrameCoder<'a> { + #[allow(clippy::too_many_arguments)] + pub fn new( + bd: u32, + w: usize, + h: usize, + kind: SliceKind, + slice_qp: i32, + src: &'a [Vec; 3], + refp: Option<&'a [Vec; 3]>, + me: Option<&'a MeRef>, + ) -> Self { + let off = 6 * (bd as i32 - 8); + let qpi = slice_qp.clamp(-off, 57); + let qp_c = qpc_420(qpi) + off; + let base = if kind == SliceKind::I { 0.57 } else { 0.62 }; + let lambda = base * 2f64.powf((slice_qp as f64 - 12.0) / 3.0) * (1u64 << (2 * (bd - 8))) as f64; + let ctb = 1usize << LOG2_CTB; + FrameCoder { + bd, + w, + h, + cw: w / 2, + ch: h / 2, + w4: w / 4, + wctb: w.div_ceil(ctb), + hctb: h.div_ceil(ctb), + kind, + qp_y: slice_qp + off, + qp_c, + lambda, + lambda_sad: lambda.sqrt(), + src, + rec: [vec![0; w * h], vec![0; w * h / 4], vec![0; w * h / 4]], + refp, + me, + blk: vec![Blk::default(); (w / 4) * (h / 4)], + tbs: Vec::with_capacity(16), + coefs: Vec::with_capacity(64 * 64 * 2), + pred: vec![0; 64 * 64], + } + } + + // --------------------------------------------------------------------------------------- + // neighbourhood + + /// z-scan order availability (6.4.1) of (xn, yn) for the block at (xc, yc): single slice, no tiles. + #[inline] + fn zavail(&self, xc: i32, yc: i32, xn: i32, yn: i32) -> bool { + if xn < 0 || yn < 0 || xn >= self.w as i32 || yn >= self.h as i32 { + return false; + } + let s = LOG2_CTB; + let ctb_n = (yn >> s) as usize * self.wctb + (xn >> s) as usize; + let ctb_c = (yc >> s) as usize * self.wctb + (xc >> s) as usize; + if ctb_n != ctb_c { + ctb_n < ctb_c + } else { + let m = (1 << s) - 1; + morton(((xn & m) >> 2) as u32, ((yn & m) >> 2) as u32) <= morton(((xc & m) >> 2) as u32, ((yc & m) >> 2) as u32) + } + } + + #[inline] + fn blk_at(&self, x: i32, y: i32) -> &Blk { + &self.blk[(y as usize >> 2) * self.w4 + (x as usize >> 2)] + } + + fn fill(&mut self, x: usize, y: usize, n: usize, f: impl Fn(&mut Blk)) { + for by in (y >> 2)..((y + n) >> 2) { + for b in &mut self.blk[by * self.w4 + (x >> 2)..by * self.w4 + ((x + n) >> 2)] { + f(b); + } + } + } + + fn snap(&self, x: usize, y: usize, n: usize) -> Snap { + let mut planes: [Vec; 3] = Default::default(); + for (c, p) in planes.iter_mut().enumerate() { + let (px, py, pn, stride) = if c == 0 { (x, y, n, self.w) } else { (x / 2, y / 2, n / 2, self.cw) }; + p.reserve(pn * pn); + for r in 0..pn { + p.extend_from_slice(&self.rec[c][(py + r) * stride + px..][..pn]); + } + } + let mut blk = Vec::with_capacity((n / 4) * (n / 4)); + for by in (y >> 2)..((y + n) >> 2) { + blk.extend_from_slice(&self.blk[by * self.w4 + (x >> 2)..by * self.w4 + ((x + n) >> 2)]); + } + Snap { x, y, n, planes, blk } + } + + fn restore(&mut self, s: &Snap) { + let (x, y, n) = (s.x, s.y, s.n); + for c in 0..3 { + let (px, py, pn, stride) = if c == 0 { (x, y, n, self.w) } else { (x / 2, y / 2, n / 2, self.cw) }; + for r in 0..pn { + self.rec[c][(py + r) * stride + px..][..pn].copy_from_slice(&s.planes[c][r * pn..r * pn + pn]); + } + } + let bw = n / 4; + for (i, by) in ((y >> 2)..((y + n) >> 2)).enumerate() { + self.blk[by * self.w4 + (x >> 2)..by * self.w4 + (x >> 2) + bw].copy_from_slice(&s.blk[i * bw..i * bw + bw]); + } + } + + // --------------------------------------------------------------------------------------- + // picture loop + + /// Code all CTUs of the picture into `enc` (one slice); the caller flushes the encoder. + pub fn encode(&mut self, enc: &mut CabacEncoder) { + let ctb = 1usize << LOG2_CTB; + let total = self.wctb * self.hctb; + for i in 0..total { + let (x, y) = ((i % self.wctb) * ctb, (i / self.wctb) * ctb); + let mut est = Estimator::new(enc.ctx); + let mut decs = Vec::with_capacity(64); + self.decide(x, y, LOG2_CTB, 0, &mut est, &mut decs); + let mut idx = 0; + self.write_cq(enc, x, y, LOG2_CTB, 0, &decs, &mut idx); + debug_assert_eq!(idx, decs.len()); + debug_assert_eq!(est.ctx, enc.ctx); + enc.terminate((i + 1 == total) as u32); // end_of_slice_segment_flag + } + } + + fn write_split(&self, s: &mut S, x: usize, y: usize, depth: u32, split: bool) { + let (x, y) = (x as i32, y as i32); + let mut inc = 0; + if self.zavail(x, y, x - 1, y) && self.blk_at(x - 1, y).depth as u32 > depth { + inc += 1; + } + if self.zavail(x, y, x, y - 1) && self.blk_at(x, y - 1).depth as u32 > depth { + inc += 1; + } + s.decision(SPLIT_CU + inc, split as u32); + } + + /// Final pass: re-code the decided quadtree with the real arithmetic coder. + fn write_cq(&mut self, enc: &mut CabacEncoder, x: usize, y: usize, log2: u32, depth: u32, decs: &[CuDec], idx: &mut usize) { + if x >= self.w || y >= self.h { + return; + } + let n = 1usize << log2; + let inside = x + n <= self.w && y + n <= self.h; + let d = decs[*idx]; + let leaf = d.x as usize == x && d.y as usize == y && d.log2 == log2; + if inside && log2 > LOG2_MIN_CB { + self.write_split(enc, x, y, depth, !leaf); + } else { + debug_assert_eq!(leaf, log2 == LOG2_MIN_CB); + } + if leaf { + *idx += 1; + let (eff, _) = self.encode_cu(enc, x, y, log2, depth, d.mode); + debug_assert_eq!(eff, d.mode); + } else { + let h = n / 2; + for (dx, dy) in [(0, 0), (h, 0), (0, h), (h, h)] { + self.write_cq(enc, x + dx, y + dy, log2 - 1, depth + 1, decs, idx); + } + } + } + + #[inline] + fn cost(&self, dist: u64, bits: u64) -> f64 { + dist as f64 + self.lambda * bits as f64 / BIT as f64 + } + + /// Rate-distortion decision for the quadtree node at (x, y); leaves the chosen reconstruction + /// and block state in place, advances `est` and appends the decided CUs. Returns the RD cost. + fn decide(&mut self, x: usize, y: usize, log2: u32, depth: u32, est: &mut Estimator, decs: &mut Vec) -> f64 { + if x >= self.w || y >= self.h { + return 0.0; + } + let n = 1usize << log2; + let inside = x + n <= self.w && y + n <= self.h; + let can_split = log2 > LOG2_MIN_CB; + let base = est.bits; + let leaf_ok = inside && (self.kind == SliceKind::P || log2 <= LOG2_MAX_TB); + let mut best: Option<(f64, Estimator, Snap, CuMode)> = None; + if leaf_ok { + let mut e0 = est.clone(); + if can_split { + self.write_split(&mut e0, x, y, depth, false); + } + let try_mode = |me: &mut Self, mode: CuMode, best: &mut Option<(f64, Estimator, Snap, CuMode)>| { + let mut e = e0.clone(); + let (eff, dist) = me.encode_cu(&mut e, x, y, log2, depth, mode); + let cost = me.cost(dist, e.bits - base); + if best.as_ref().is_none_or(|b| cost < b.0) { + *best = Some((cost, e, me.snap(x, y, n), eff)); + } + }; + let mut inter_satd = u64::MAX; + if self.kind == SliceKind::P { + let merge = self.merge_list(x, y, n); + let (midx, msatd) = self.best_merge(x, y, n, &merge); + try_mode(self, CuMode::Skip { idx: midx as u8 }, &mut best); + try_mode(self, CuMode::Merge { idx: midx as u8 }, &mut best); + let amvp = self.amvp(x, y, n); + let (mv, mvp, asatd) = self.motion_search(x, y, n, &amvp, &merge); + if mv != merge[midx] { + try_mode(self, CuMode::Amvp { mv, mvp }, &mut best); + } + inter_satd = msatd.min(asatd); + } + if log2 <= LOG2_MAX_TB && !matches!(best, Some((_, _, _, CuMode::Skip { .. }))) { + let (mode, satd) = self.search_intra_mode(x, y, log2); + if satd < inter_satd { + try_mode(self, CuMode::Intra { nxn: false, modes: [mode as u8; 4] }, &mut best); + if log2 == LOG2_MIN_CB { + let modes = self.search_nxn_modes(x, y); + try_mode(self, CuMode::Intra { nxn: true, modes }, &mut best); + } + } + } + if let Some(b) = &best { + self.restore(&b.2); + } + } + let early = matches!(best, Some((_, _, _, CuMode::Skip { .. }))); + if can_split && (!leaf_ok || !early) { + let mut e = est.clone(); + if inside { + self.write_split(&mut e, x, y, depth, true); + } + let limit = best.as_ref().map_or(f64::INFINITY, |b| b.0); + let mut cost = self.cost(0, e.bits - base); + let mut d2 = Vec::new(); + let h = n / 2; + for (dx, dy) in [(0, 0), (h, 0), (0, h), (h, h)] { + cost += self.decide(x + dx, y + dy, log2 - 1, depth + 1, &mut e, &mut d2); + if cost >= limit { + break; + } + } + if cost < limit { + *est = e; + decs.extend(d2); + return cost; + } + } + let (cost, e, snap, mode) = best.expect("a coding unit decision"); + self.restore(&snap); + *est = e; + decs.push(CuDec { x: x as u32, y: y as u32, log2, mode }); + cost + } + + // --------------------------------------------------------------------------------------- + // coding unit + + fn skip_ctx(&self, x: usize, y: usize) -> usize { + let (x, y) = (x as i32, y as i32); + let mut inc = 0; + if self.zavail(x, y, x - 1, y) && self.blk_at(x - 1, y).flags & F_SKIP != 0 { + inc += 1; + } + if self.zavail(x, y, x, y - 1) && self.blk_at(x, y - 1).flags & F_SKIP != 0 { + inc += 1; + } + inc + } + + /// Predict, transform, quantise and reconstruct one CU and write its syntax into `s`. + /// Returns the effective mode (a residual-free merge becomes a skip) and the SSE. + fn encode_cu(&mut self, s: &mut S, x: usize, y: usize, log2: u32, depth: u32, mode: CuMode) -> (CuMode, u64) { + let n = 1usize << log2; + self.tbs.clear(); + self.coefs.clear(); + let p_slice = self.kind == SliceKind::P; + let skip_ctx = if p_slice { self.skip_ctx(x, y) } else { 0 }; + let d8 = depth as u8; + match mode { + CuMode::Intra { nxn, modes } => { + self.fill(x, y, n, |b| { + b.depth = d8; + b.flags = F_INTRA; + b.ipm = modes[0]; + }); + if nxn { + let h = n / 2; + for k in 1..4 { + let m = modes[k]; + self.fill(x + (k & 1) * h, y + (k >> 1) * h, h, |b| b.ipm = m); + } + for k in 0..4 { + let (px, py) = (x + (k & 1) * h, y + (k >> 1) * h); + self.code_tb(0, px, py, 2, Some(modes[k] as u32), (x, y)); + } + for c in 1..3 { + self.code_tb(c, x / 2, y / 2, 2, Some(modes[0] as u32), (x, y)); + } + } else { + self.code_tus(x, y, log2, Some(modes[0] as u32)); + } + if p_slice { + s.decision(CU_SKIP + skip_ctx, 0); + s.decision(PRED_MODE, 1); + } + if log2 == LOG2_MIN_CB { + s.decision(PART_MODE, (!nxn) as u32); + } + self.write_intra_modes(s, x, y, n, nxn, &modes); + s.decision(INTRA_CHROMA, 0); // intra_chroma_pred_mode 4 (DM) + let mut ti = 0; + self.write_tt(s, x, y, x, y, log2, 0, 0, [true, true], true, nxn, &mut ti); + (mode, self.cu_dist(x, y, n)) + } + _ => { + let mv = match mode { + CuMode::Skip { idx } | CuMode::Merge { idx } => self.merge_list(x, y, n)[idx as usize], + CuMode::Amvp { mv, .. } => mv, + CuMode::Intra { .. } => unreachable!(), + }; + let amvp = if let CuMode::Amvp { .. } = mode { self.amvp(x, y, n) } else { [[0; 2]; 2] }; + let skip = matches!(mode, CuMode::Skip { .. }); + self.fill(x, y, n, |b| { + b.depth = d8; + b.flags = if skip { F_SKIP } else { 0 }; + b.ipm = 1; + b.mv = mv; + }); + self.predict_inter(x, y, n, mv); + let mut eff = mode; + let mut root = false; + if !skip { + self.code_tus(x, y, log2, None); + root = self.tbs.iter().any(|t| t.nz); + if !root && let CuMode::Merge { idx } = mode { + // rqt_root_cbf would be inferred: code a skipped CU instead + eff = CuMode::Skip { idx }; + self.fill(x, y, n, |b| b.flags = F_SKIP); + } + } + match eff { + CuMode::Skip { idx } => { + s.decision(CU_SKIP + skip_ctx, 1); + write_merge_idx(s, idx as usize); + } + CuMode::Merge { idx } => { + s.decision(CU_SKIP + skip_ctx, 0); + s.decision(PRED_MODE, 0); + s.decision(PART_MODE, 1); // PART_2Nx2N + s.decision(MERGE_FLAG, 1); + write_merge_idx(s, idx as usize); + let mut ti = 0; + self.write_tt(s, x, y, x, y, log2, 0, 0, [true, true], false, false, &mut ti); + } + CuMode::Amvp { mv, mvp } => { + s.decision(CU_SKIP + skip_ctx, 0); + s.decision(PRED_MODE, 0); + s.decision(PART_MODE, 1); + s.decision(MERGE_FLAG, 0); + let p = amvp[mvp as usize]; + write_mvd(s, [mv[0] as i32 - p[0] as i32, mv[1] as i32 - p[1] as i32]); + s.decision(MVP_FLAG, mvp as u32); + s.decision(RQT_ROOT_CBF, root as u32); + if root { + let mut ti = 0; + self.write_tt(s, x, y, x, y, log2, 0, 0, [true, true], false, false, &mut ti); + } + } + CuMode::Intra { .. } => unreachable!(), + } + (eff, self.cu_dist(x, y, n)) + } + } + } + + /// Transform units of a 2Nx2N CU (split to the maximum TB size), luma then Cb, Cr per TU. + fn code_tus(&mut self, x: usize, y: usize, log2: u32, intra_mode: Option) { + let tl = log2.min(LOG2_MAX_TB); + let t = 1usize << tl; + let k = 1usize << (log2 - tl); + for i in 0..k * k { + // z-order of at most 2x2 TUs + let (tx, ty) = (x + (i & 1) * t, y + (i >> 1) * t); + self.code_tb(0, tx, ty, tl, intra_mode, (tx, ty)); + self.code_tb(1, tx / 2, ty / 2, tl - 1, intra_mode, (tx, ty)); + self.code_tb(2, tx / 2, ty / 2, tl - 1, intra_mode, (tx, ty)); + } + } + + fn cu_dist(&self, x: usize, y: usize, n: usize) -> u64 { + let mut d = 0u64; + for c in 0..3 { + let (px, py, pn, stride) = if c == 0 { (x, y, n, self.w) } else { (x / 2, y / 2, n / 2, self.cw) }; + for r in 0..pn { + let o = (py + r) * stride + px; + let a = &self.src[c][o..o + pn]; + let b = &self.rec[c][o..o + pn]; + d += a.iter().zip(b).map(|(&p, &q)| (p as i64 - q as i64).pow(2) as u64).sum::(); + } + } + d + } + + /// Code one transform block of component `c` at component position (x, y): optional intra + /// prediction, residual, transform, quantisation and reconstruction. + fn code_tb(&mut self, c: usize, x: usize, y: usize, log2: u32, intra_mode: Option, lpos: (usize, usize)) { + let n = 1usize << log2; + let stride = if c == 0 { self.w } else { self.cw }; + if let Some(m) = intra_mode { + let refs = self.build_refs(c, x, y, log2); + let bd = self.bd; + predict_from_refs(&refs, c, m, n, bd, &mut self.rec[c][y * stride + x..], stride); + } + let mut r = [0i32; 32 * 32]; + for yy in 0..n { + let o = (y + yy) * stride + x; + for xx in 0..n { + r[yy * n + xx] = self.src[c][o + xx] as i32 - self.rec[c][o + xx] as i32; + } + } + let intra = intra_mode.is_some(); + let dst = intra && c == 0 && n == 4; + let mut coef = [0i32; 32 * 32]; + transform::forward(&r[..n * n], &mut coef[..n * n], n, dst, self.bd); + let qp = if c == 0 { self.qp_y } else { self.qp_c }; + let mut lv = [0i32; 32 * 32]; + let nz = transform::quantize(&coef[..n * n], &mut lv[..n * n], n, qp, intra, self.bd) > 0; + if nz { + let mut d = [0i32; 32 * 32]; + transform::dequantize(&lv[..n * n], &mut d[..n * n], n, qp, self.bd); + transform::inverse(&mut d[..n * n], n, dst, self.bd); + let max = (1i32 << self.bd) - 1; + for yy in 0..n { + let o = (y + yy) * stride + x; + for xx in 0..n { + let v = &mut self.rec[c][o + xx]; + *v = (*v as i32 + d[yy * n + xx]).clamp(0, max) as u16; + } + } + } + let scan = match intra_mode { + Some(m) if log2 == 2 || (log2 == 3 && c == 0) => scan_idx_for_mode(m), + _ => 0, + }; + let off = self.coefs.len(); + self.coefs.extend_from_slice(&lv[..n * n]); + self.tbs.push(Tb { c: c as u8, lx: lpos.0 as u32, ly: lpos.1 as u32, log2, nz, off, scan }); + } + + /// Reference samples (substituted, unfiltered) of an intra block, as in 8.4.4.2.1 / 8.4.4.2.2. + fn build_refs(&self, c: usize, x0: usize, y0: usize, log2: u32) -> Refs { + let n = 1usize << log2; + let sc = if c == 0 { 1 } else { 2 }; + let (x0, y0) = (x0 as i32, y0 as i32); + let (xl, yl) = (x0 * sc, y0 * sc); + let unit = 4 / sc as usize; + let stride = if c == 0 { self.w } else { self.cw }; + let mut al = [false; 129]; + let mut at = [false; 129]; + let plane = &self.rec[c]; + let mut r = Refs::default(); + if self.zavail(xl, yl, xl - 1, yl - 1) { + al[0] = true; + at[0] = true; + let v = plane[(y0 - 1) as usize * stride + (x0 - 1) as usize]; + r.left[0] = v; + r.top[0] = v; + } + let mut k = 0; + while k < 2 * n { + let yn = y0 + k as i32; + if self.zavail(xl, yl, xl - 1, yn * sc) { + for i in 0..unit { + al[1 + k + i] = true; + r.left[1 + k + i] = plane[(yn as usize + i) * stride + (x0 - 1) as usize]; + } + } + let xn = x0 + k as i32; + if self.zavail(xl, yl, xn * sc, yl - 1) { + let o = (y0 - 1) as usize * stride + xn as usize; + for i in 0..unit { + at[1 + k + i] = true; + r.top[1 + k + i] = plane[o + i]; + } + } + k += unit; + } + intra::substitute(&mut r, &al, &at, n, self.bd); + r + } + + // --------------------------------------------------------------------------------------- + // syntax writers + + /// Candidate modes candModeList (8.4.2) of the prediction block at (xp, yp). + fn mpm(&self, xp: usize, yp: usize) -> [u32; 3] { + let (xp, yp) = (xp as i32, yp as i32); + let cand_a = if self.zavail(xp, yp, xp - 1, yp) { + let b = self.blk_at(xp - 1, yp); + if b.flags & F_INTRA != 0 { b.ipm as u32 } else { 1 } + } else { + 1 + }; + let cand_b = if self.zavail(xp, yp, xp, yp - 1) && yp > ((yp >> LOG2_CTB) << LOG2_CTB) { + let b = self.blk_at(xp, yp - 1); + if b.flags & F_INTRA != 0 { b.ipm as u32 } else { 1 } + } else { + 1 + }; + if cand_a == cand_b { + if cand_a < 2 { [0, 1, 26] } else { [cand_a, 2 + ((cand_a + 29) % 32), 2 + ((cand_a - 2 + 1) % 32)] } + } else { + let c = if cand_a != 0 && cand_b != 0 { + 0 + } else if cand_a != 1 && cand_b != 1 { + 1 + } else { + 26 + }; + [cand_a, cand_b, c] + } + } + + fn write_intra_modes(&self, s: &mut S, x: usize, y: usize, n: usize, nxn: bool, modes: &[u8; 4]) { + let nb = if nxn { 2 } else { 1 }; + let pb = n / nb; + let mut codes = [(false, 0u32); 4]; + for j in 0..nb { + for i in 0..nb { + let k = j * nb + i; + let list = self.mpm(x + i * pb, y + j * pb); + let m = modes[k] as u32; + codes[k] = match list.iter().position(|&c| c == m) { + Some(idx) => (true, idx as u32), + None => { + let mut sorted = list; + sorted.sort_unstable(); + let mut r = m; + for &c in sorted.iter().rev() { + if r > c { + r -= 1; + } + } + (false, r) + } + }; + } + } + for code in codes.iter().take(nb * nb) { + s.decision(PREV_INTRA_LUMA, code.0 as u32); + } + for &(prev, v) in codes.iter().take(nb * nb) { + if prev { + s.bypass((v > 0) as u32); + if v > 0 { + s.bypass((v > 1) as u32); + } + } else { + s.bypass_bits(v, 5); + } + } + } + + fn chroma_cbf(&self, c: u8, x0: usize, y0: usize, n: usize) -> bool { + self.tbs.iter().any(|t| t.c == c && t.nz && (t.lx as usize) >= x0 && (t.lx as usize) < x0 + n && (t.ly as usize) >= y0 && (t.ly as usize) < y0 + n) + } + + /// transform_tree() (7.3.8.8) over the coded transform blocks of the CU; no split flag is ever + /// coded (max_transform_hierarchy_depth_* = 0), splits are inferred. + #[allow(clippy::too_many_arguments)] + fn write_tt( + &self, + s: &mut S, + x0: usize, + y0: usize, + xb: usize, + yb: usize, + log2: u32, + depth: u32, + blk_idx: u32, + parent: [bool; 2], + intra: bool, + nxn: bool, + ti: &mut usize, + ) { + let n = 1usize << log2; + let split = log2 > LOG2_MAX_TB || (nxn && depth == 0); + let mut cbf_c = [false; 2]; + if log2 > 2 { + for c in 0..2 { + if depth == 0 || parent[c] { + cbf_c[c] = self.chroma_cbf(c as u8 + 1, x0, y0, n); + s.decision(CBF_CHROMA + depth as usize, cbf_c[c] as u32); + } + } + } else if depth > 0 { + cbf_c = parent; + } + if split { + let h = n / 2; + for (i, (dx, dy)) in [(0, 0), (h, 0), (0, h), (h, h)].into_iter().enumerate() { + self.write_tt(s, x0 + dx, y0 + dy, x0, y0, log2 - 1, depth + 1, i as u32, cbf_c, intra, nxn, ti); + } + return; + } + let _ = (xb, yb); + let tb = self.tbs[*ti]; + *ti += 1; + debug_assert_eq!(tb.c, 0); + if intra || depth != 0 || cbf_c[0] || cbf_c[1] { + s.decision(CBF_LUMA + (depth == 0) as usize, tb.nz as u32); + } else { + debug_assert!(tb.nz, "inferred cbf_luma needs coefficients"); + } + if tb.nz { + let n2 = 1usize << (2 * tb.log2); + residual_coding(s, &self.coefs[tb.off..tb.off + n2], tb.log2, 0, tb.scan as usize); + } + if log2 > 2 || blk_idx == 3 { + for c in 0..2 { + let tb = self.tbs[*ti]; + *ti += 1; + debug_assert_eq!(tb.c as usize, c + 1); + debug_assert_eq!(tb.nz, cbf_c[c]); + if cbf_c[c] { + let n2 = 1usize << (2 * tb.log2); + residual_coding(s, &self.coefs[tb.off..tb.off + n2], tb.log2, c + 1, tb.scan as usize); + } + } + } + } + + // --------------------------------------------------------------------------------------- + // inter prediction + + /// Prediction block availability (6.4.2) for a 2Nx2N CU: available and not intra. + #[inline] + fn pb_avail(&self, xp: i32, yp: i32, xn: i32, yn: i32) -> bool { + self.zavail(xp, yp, xn, yn) && self.blk_at(xn, yn).flags & F_INTRA == 0 + } + + /// Merge candidate list (8.5.3.2.2 - 8.5.3.2.5) of a 2Nx2N P-slice CU: spatial candidates and + /// zero candidates (no temporal candidate, one reference picture). + pub fn merge_list(&self, x: usize, y: usize, n: usize) -> [[i16; 2]; MAX_MERGE_CAND] { + let (xp, yp, w, h) = (x as i32, y as i32, n as i32, n as i32); + let mv = |p: (i32, i32)| self.blk_at(p.0, p.1).mv; + let mut list = [[0i16; 2]; MAX_MERGE_CAND]; + let mut k = 0; + let pa1 = (xp - 1, yp + h - 1); + let av_a1 = self.pb_avail(xp, yp, pa1.0, pa1.1); + if av_a1 { + list[k] = mv(pa1); + k += 1; + } + let pb1 = (xp + w - 1, yp - 1); + let av_b1 = self.pb_avail(xp, yp, pb1.0, pb1.1); + let fl_b1 = av_b1 && !(av_a1 && mv(pa1) == mv(pb1)); + if fl_b1 { + list[k] = mv(pb1); + k += 1; + } + let pb0 = (xp + w, yp - 1); + let fl_b0 = self.pb_avail(xp, yp, pb0.0, pb0.1) && !(av_b1 && mv(pb1) == mv(pb0)); + if fl_b0 { + list[k] = mv(pb0); + k += 1; + } + let pa0 = (xp - 1, yp + h); + let fl_a0 = self.pb_avail(xp, yp, pa0.0, pa0.1) && !(av_a1 && mv(pa1) == mv(pa0)); + if fl_a0 { + list[k] = mv(pa0); + k += 1; + } + let pb2 = (xp - 1, yp - 1); + let fl_b2 = self.pb_avail(xp, yp, pb2.0, pb2.1) + && !((av_a1 && mv(pa1) == mv(pb2)) || (av_b1 && mv(pb1) == mv(pb2))) + && (fl_a0 as u32 + av_a1 as u32 + fl_b0 as u32 + fl_b1 as u32) != 4; + if fl_b2 { + list[k] = mv(pb2); + } + // remaining entries are zero candidates with refIdx 0 + list + } + + /// AMVP candidates (8.5.3.2.6, 8.5.3.2.7) of a 2Nx2N CU with one reference picture. + pub fn amvp(&self, x: usize, y: usize, n: usize) -> [[i16; 2]; 2] { + let (xp, yp, w, h) = (x as i32, y as i32, n as i32, n as i32); + let pa = [(xp - 1, yp + h), (xp - 1, yp + h - 1)]; + let av_a = [self.pb_avail(xp, yp, pa[0].0, pa[0].1), self.pb_avail(xp, yp, pa[1].0, pa[1].1)]; + let is_scaled = av_a[0] || av_a[1]; + let mut a = None; + for k in 0..2 { + if av_a[k] && a.is_none() { + a = Some(self.blk_at(pa[k].0, pa[k].1).mv); + } + } + let pb = [(xp + w, yp - 1), (xp + w - 1, yp - 1), (xp - 1, yp - 1)]; + let mut b = None; + for p in pb { + if b.is_none() && self.pb_avail(xp, yp, p.0, p.1) { + b = Some(self.blk_at(p.0, p.1).mv); + } + } + if !is_scaled && b.is_some() { + a = b; + } + let mut list = [[0i16; 2]; 2]; + match (a, b) { + (Some(a), Some(b)) => { + list[0] = a; + if b != a { + list[1] = b; + } + } + (Some(a), None) => list[0] = a, + (None, Some(b)) => list[0] = b, + (None, None) => {} + } + list + } + + fn predict_inter(&mut self, x: usize, y: usize, n: usize, mv: [i16; 2]) { + let refp = self.refp.expect("reference picture"); + for c in 0..3 { + let (px, py, pn, pw, ph) = if c == 0 { (x, y, n, self.w, self.h) } else { (x / 2, y / 2, n / 2, self.cw, self.ch) }; + let pr = PlaneRef { data: &refp[c], w: pw, h: ph }; + if c == 0 { + inter::mc_luma(pr, self.bd, px as i32, py as i32, mv, pn, pn, &mut self.pred); + } else { + inter::mc_chroma(pr, self.bd, px as i32, py as i32, mv, pn, pn, &mut self.pred); + } + inter::put_uni(&self.pred, pn, pn, self.bd, &mut self.rec[c][py * pw + px..], pw); + } + } + + // --------------------------------------------------------------------------------------- + // estimation helpers + + /// Best merge index by SATD of the (approximate) luma prediction. + fn best_merge(&self, x: usize, y: usize, n: usize, list: &[[i16; 2]; MAX_MERGE_CAND]) -> (usize, u64) { + let me = self.me.expect("motion search reference"); + let mut buf = vec![0u16; n * n]; + let mut best = (0, u64::MAX); + for (i, &mv) in list.iter().enumerate() { + if list[..i].contains(&mv) { + continue; + } + me.approx_pred(x as i32, y as i32, n, mv, &mut buf); + let satd = satd(&self.src[0][y * self.w + x..], self.w, &buf, n, n) as u64; + let cost = satd + (self.lambda_sad * (i + 1) as f64) as u64; + if cost < best.1 { + best = (i, cost); + } + } + best + } + + /// Integer then fractional motion search. Returns (mv, mvp index, SATD cost). + fn motion_search(&self, x: usize, y: usize, n: usize, amvp: &[[i16; 2]; 2], merge: &[[i16; 2]; MAX_MERGE_CAND]) -> ([i16; 2], u8, u64) { + let me = self.me.expect("motion search reference"); + let (xi, yi) = (x as i32, y as i32); + let src = &self.src[0][y * self.w + x..]; + let ls = self.lambda_sad; + let mv_bits = |mv: [i32; 2]| -> f64 { amvp.iter().map(|p| mvd_bits(mv[0] - p[0] as i32) + mvd_bits(mv[1] - p[1] as i32)).fold(f64::MAX, f64::min) }; + let (lo_x, hi_x, lo_y, hi_y) = me.int_range(xi, yi, n); + let clamp_i = |v: [i32; 2]| [v[0].clamp(lo_x, hi_x), v[1].clamp(lo_y, hi_y)]; + let sad_cost = |v: [i32; 2]| -> f64 { me.sad(src, self.w, xi + v[0], yi + v[1], n) as f64 + ls * mv_bits([v[0] * 4, v[1] * 4]) }; + // starting points + let mut best = [0i32; 2]; + let mut best_cost = sad_cost(best); + for mv in amvp.iter().chain(merge.iter()) { + let v = clamp_i([(mv[0] as i32 + 2) >> 2, (mv[1] as i32 + 2) >> 2]); + let c = sad_cost(v); + if c < best_cost { + best_cost = c; + best = v; + } + } + // square pattern search with decreasing steps + let mut step = 16; + while step >= 1 { + let mut moved = true; + let mut iters = 0; + while moved && iters < 4 { + moved = false; + iters += 1; + let center = best; + for (dx, dy) in [(-1, -1), (0, -1), (1, -1), (-1, 0), (1, 0), (-1, 1), (0, 1), (1, 1)] { + let v = clamp_i([center[0] + dx * step, center[1] + dy * step]); + if v == center { + continue; + } + let c = sad_cost(v); + if c < best_cost { + best_cost = c; + best = v; + moved = true; + } + } + } + step /= 2; + } + // fractional refinement with SATD on the approximate interpolated planes + let mut buf = vec![0u16; n * n]; + let mut satd_cost = |mv: [i32; 2]| -> f64 { + me.approx_pred(xi, yi, n, [mv[0] as i16, mv[1] as i16], &mut buf); + satd(src, self.w, &buf, n, n) as f64 + ls * mv_bits(mv) + }; + let mut bq = [best[0] * 4, best[1] * 4]; + let mut bc = satd_cost(bq); + for step in [2, 1] { + let center = bq; + for (dx, dy) in [(-1, -1), (0, -1), (1, -1), (-1, 0), (1, 0), (-1, 1), (0, 1), (1, 1)] { + let v = [center[0] + dx * step, center[1] + dy * step]; + let c = satd_cost(v); + if c < bc { + bc = c; + bq = v; + } + } + } + let mv = [bq[0] as i16, bq[1] as i16]; + let bits = |p: &[i16; 2]| mvd_bits(bq[0] - p[0] as i32) + mvd_bits(bq[1] - p[1] as i32); + let mvp = if bits(&amvp[1]) < bits(&amvp[0]) { 1 } else { 0 }; + (mv, mvp, bc as u64) + } + + /// Best 2Nx2N intra luma mode by SATD (coarse grid, then refinement around the best angle). + fn search_intra_mode(&self, x: usize, y: usize, log2: u32) -> (u32, u64) { + let n = 1usize << log2; + let refs = self.build_refs(0, x, y, log2); + let mpm = self.mpm(x, y); + let mut buf = vec![0u16; n * n]; + let src = &self.src[0][y * self.w + x..]; + let mut eval = |m: u32| -> u64 { + predict_from_refs(&refs, 0, m, n, self.bd, &mut buf, n); + let bits = match mpm.iter().position(|&c| c == m) { + Some(0) => 2.0, + Some(_) => 3.0, + None => 6.0, + }; + satd(src, self.w, &buf, n, n) as u64 + (self.lambda_sad * bits) as u64 + }; + let mut best = (1u32, u64::MAX); + let mut tried = [false; 35]; + for &m in COARSE_MODES.iter().chain(mpm.iter()) { + if !tried[m as usize] { + tried[m as usize] = true; + let c = eval(m); + if c < best.1 { + best = (m, c); + } + } + } + if best.0 >= 2 { + for d in [2i32, 1] { + let center = best.0 as i32; + for m in [center - d, center + d] { + if (2..=34).contains(&m) && !tried[m as usize] { + tried[m as usize] = true; + let c = eval(m as u32); + if c < best.1 { + best = (m as u32, c); + } + } + } + } + } + best + } + + /// Modes of the four 4x4 PUs of an NxN intra CU, chosen sequentially (each PU reconstructed + /// before the next one is searched, as in the real coding order). + fn search_nxn_modes(&mut self, x: usize, y: usize) -> [u8; 4] { + let mut modes = [1u8; 4]; + self.fill(x, y, 8, |b| { + b.flags = F_INTRA; + b.ipm = 1; + }); + for k in 0..4 { + let (px, py) = (x + (k & 1) * 4, y + (k >> 1) * 4); + let (m, _) = self.search_intra_mode(px, py, 2); + modes[k] = m as u8; + self.fill(px, py, 4, |b| b.ipm = m as u8); + self.tbs.clear(); + self.coefs.clear(); + self.code_tb(0, px, py, 2, Some(m), (px, py)); + } + modes + } +} + +/// Filter (luma) and predict an intra block from substituted references. +fn predict_from_refs(refs: &Refs, c: usize, mode: u32, n: usize, bd: u32, out: &mut [u16], os: usize) { + if c == 0 { + let mut r = refs.clone(); + intra::filter(&mut r, mode, n, true, true, bd); + intra::predict(&r, mode, n, true, bd, out, os); + } else { + intra::predict(refs, mode, n, false, bd, out, os); + } +} + +/// scanIdx (7.4.9.11) for an intra prediction mode. +fn scan_idx_for_mode(m: u32) -> u8 { + if (6..=14).contains(&m) { + 2 + } else if (22..=30).contains(&m) { + 1 + } else { + 0 + } +} + +/// Approximate bits of one motion vector difference component. +fn mvd_bits(v: i32) -> f64 { + if v == 0 { 1.0 } else { (2 * (32 - (v.unsigned_abs()).leading_zeros()) + 1) as f64 } +} + +fn write_merge_idx(s: &mut S, idx: usize) { + if MAX_MERGE_CAND > 1 { + s.decision(MERGE_IDX, (idx > 0) as u32); + if idx > 0 { + let mut i = 1; + while i < MAX_MERGE_CAND - 1 { + let b = idx > i; + s.bypass(b as u32); + if !b { + break; + } + i += 1; + } + } + } +} + +/// mvd_coding() (7.3.8.9). +fn write_mvd(s: &mut S, d: [i32; 2]) { + let g0 = [d[0] != 0, d[1] != 0]; + s.decision(ABS_MVD_GT0, g0[0] as u32); + s.decision(ABS_MVD_GT0, g0[1] as u32); + for k in 0..2 { + if g0[k] { + s.decision(ABS_MVD_GT1, (d[k].abs() > 1) as u32); + } + } + for k in 0..2 { + if g0[k] { + let a = d[k].unsigned_abs(); + if a > 1 { + // abs_mvd_minus2: EG1 + let mut v = a - 2; + let mut kk = 1; + while v >= (1 << kk) { + s.bypass(1); + v -= 1 << kk; + kk += 1; + } + s.bypass(0); + s.bypass_bits(v, kk); + } + s.bypass((d[k] < 0) as u32); + } + } +} + +/// coeff_abs_level_remaining binarisation (9.3.3.11): TR prefix up to 3 with Rice suffix, then +/// an escape with an Exp-Golomb-like suffix. +fn write_remaining(s: &mut S, v: u32, rice: u32) { + if v < (4 << rice) { + let p = v >> rice; + for _ in 0..p { + s.bypass(1); + } + s.bypass(0); + s.bypass_bits(v & ((1 << rice) - 1), rice); + } else { + let u = v - (2 << rice); + let kr = 31 - u.leading_zeros(); // k + rice + let k = kr - rice; + for _ in 0..k + 3 { + s.bypass(1); + } + s.bypass(0); + s.bypass_bits(u - (1 << kr), kr); + } +} + +/// last_sig_coeff_{x,y}_prefix value and suffix for a position. +fn last_prefix(v: u32) -> (u32, u32, u32) { + if v < 4 { + return (v, 0, 0); + } + let mut p = 4; + loop { + let nb = (p >> 1) - 1; + let min = (1 << nb) * (2 + (p & 1)); + let next_nb = ((p + 1) >> 1) - 1; + let next = (1 << next_nb) * (2 + ((p + 1) & 1)); + if v < next { + return (p, v - min, nb); + } + p += 1; + } +} + +/// residual_coding() (7.3.8.11) of one transform block of levels (row-major, n x n). No transform +/// skip, no sign data hiding. +fn residual_coding(s: &mut S, lv: &[i32], log2: u32, c: usize, scan_idx: usize) { + let n = 1usize << log2; + let log2sb = log2 - 2; + let sb_w = 1usize << log2sb; + let sb_scan = scan_order(log2sb, scan_idx); + let pos_scan = &SCAN_4[scan_idx]; + let at = |xs: usize, ys: usize, p: usize| lv[(ys * 4 + pos_scan[p].1 as usize) * n + xs * 4 + pos_scan[p].0 as usize]; + // last significant coefficient in scan order + let mut last = None; + 'outer: for i in (0..sb_scan.len()).rev() { + let (xs, ys) = (sb_scan[i].0 as usize, sb_scan[i].1 as usize); + for p in (0..16).rev() { + if at(xs, ys, p) != 0 { + last = Some((i, p)); + break 'outer; + } + } + } + let (last_sb, last_pos) = last.expect("residual_coding of an all-zero block"); + let lx = sb_scan[last_sb].0 as u32 * 4 + pos_scan[last_pos].0 as u32; + let ly = sb_scan[last_sb].1 as u32 * 4 + pos_scan[last_pos].1 as u32; + let (cx, cy) = if scan_idx == 2 { (ly, lx) } else { (lx, ly) }; + let cmax = (log2 << 1) - 1; + let (off, shift) = if c == 0 { (3 * (log2 - 2) + ((log2 - 1) >> 2), (log2 + 1) >> 2) } else { (15, log2 - 2) }; + let (px, sx, nbx) = last_prefix(cx); + let (py, sy, nby) = last_prefix(cy); + for (p, base) in [(px, LAST_X_PREFIX), (py, LAST_Y_PREFIX)] { + for i in 0..p { + s.decision(base + (off + (i >> shift)) as usize, 1); + } + if p < cmax { + s.decision(base + (off + (p >> shift)) as usize, 0); + } + } + if px > 3 { + s.bypass_bits(sx, nbx); + } + if py > 3 { + s.bypass_bits(sy, nby); + } + let mut csbf = [0u8; 64]; + let mut prev_g1: Option = None; + let c_off_sig = if c == 0 { 0 } else { 27 }; + for i in (0..=last_sb).rev() { + let (xs, ys) = (sb_scan[i].0 as usize, sb_scan[i].1 as usize); + let right = if xs + 1 < sb_w { csbf[ys * 8 + xs + 1] } else { 0 }; + let below = if ys + 1 < sb_w { csbf[(ys + 1) * 8 + xs] } else { 0 }; + let mut infer_dc = false; + if i < last_sb && i > 0 { + let has = (0..16).any(|p| at(xs, ys, p) != 0); + let inc = (right | below).min(1) as usize + if c > 0 { 2 } else { 0 }; + s.decision(CODED_SUB_BLOCK + inc, has as u32); + csbf[ys * 8 + xs] = has as u8; + infer_dc = true; + } else { + csbf[ys * 8 + xs] = 1; + } + if csbf[ys * 8 + xs] == 0 { + continue; + } + let mut sig = [0u8; 16]; + let mut nsig = 0; + let start: i32 = if i == last_sb { + sig[0] = last_pos as u8; + nsig = 1; + last_pos as i32 - 1 + } else { + 15 + }; + let prev_csbf = right + (below << 1); + let mut np = start; + while np >= 0 { + let p = np as usize; + let (xp, yp) = (pos_scan[p].0 as usize, pos_scan[p].1 as usize); + let (xc, yc) = (xs * 4 + xp, ys * 4 + yp); + let v = lv[yc * n + xc] != 0; + if np > 0 || !infer_dc { + let sig_ctx = if log2 == 2 { + CTX_IDX_MAP[(yc << 2) + xc] as usize + } else if xc + yc == 0 { + 0 + } else { + let mut sc = match prev_csbf { + 0 => { + if xp + yp == 0 { + 2 + } else if xp + yp < 3 { + 1 + } else { + 0 + } + } + 1 => { + if yp == 0 { + 2 + } else if yp == 1 { + 1 + } else { + 0 + } + } + 2 => { + if xp == 0 { + 2 + } else if xp == 1 { + 1 + } else { + 0 + } + } + _ => 2, + }; + if c == 0 { + if xs + ys > 0 { + sc += 3; + } + sc += if log2 == 3 { if scan_idx == 0 { 9 } else { 15 } } else { 21 }; + } else { + sc += if log2 == 3 { 9 } else { 12 }; + } + sc + }; + s.decision(SIG_COEFF + c_off_sig + sig_ctx, v as u32); + if v { + sig[nsig] = p as u8; + nsig += 1; + infer_dc = false; + } + } else { + debug_assert!(v, "inferred DC coefficient must be non-zero"); + sig[nsig] = 0; + nsig += 1; + } + np -= 1; + } + if nsig == 0 { + continue; + } + let abs_at = |k: usize| at(xs, ys, sig[k] as usize).unsigned_abs(); + let mut ctx_set = if i == 0 || c > 0 { 0 } else { 2 }; + if prev_g1 == Some(0) { + ctx_set += 1; + } + let mut g1ctx = 1u32; + let mut gt1 = [0u8; 16]; + let mut first_g1: Option = None; + let c_off_g1 = if c > 0 { 16 } else { 0 }; + for k in 0..nsig.min(8) { + let f = abs_at(k) > 1; + s.decision(GT1 + ctx_set * 4 + g1ctx.min(3) as usize + c_off_g1, f as u32); + gt1[k] = f as u8; + if f { + g1ctx = 0; + if first_g1.is_none() { + first_g1 = Some(k); + } + } else if g1ctx > 0 { + g1ctx += 1; + } + } + prev_g1 = Some(g1ctx); + let mut gt2 = 0u32; + if let Some(k) = first_g1 { + gt2 = (abs_at(k) > 2) as u32; + s.decision(GT2 + ctx_set + if c > 0 { 4 } else { 0 }, gt2); + } + for k in 0..nsig { + s.bypass((at(xs, ys, sig[k] as usize) < 0) as u32); + } + let mut rice = 0u32; + for k in 0..nsig { + let mut base = 1 + gt1[k] as u32; + if Some(k) == first_g1 { + base += gt2; + } + let thresh = if k < 8 { if Some(k) == first_g1 { 3 } else { 2 } } else { 1 }; + if base == thresh { + let level = abs_at(k); + write_remaining(s, level - base, rice); + if level > 3 * (1 << rice) { + rice = (rice + 1).min(4); + } + } + } + } +} + +/// Sum of absolute Hadamard-transformed differences (8x8 blocks, 4x4 for 4-sample blocks). +pub fn satd(src: &[u16], ss: usize, pred: &[u16], ps: usize, n: usize) -> u32 { + if n == 4 { + return satd4(src, ss, pred, ps); + } + let mut total = 0; + for by in (0..n).step_by(8) { + for bx in (0..n).step_by(8) { + total += satd8(&src[by * ss + bx..], ss, &pred[by * ps + bx..], ps); + } + } + total +} + +fn satd4(src: &[u16], ss: usize, pred: &[u16], ps: usize) -> u32 { + let mut d = [0i32; 16]; + for y in 0..4 { + for x in 0..4 { + d[y * 4 + x] = src[y * ss + x] as i32 - pred[y * ps + x] as i32; + } + } + let mut m = [0i32; 16]; + for y in 0..4 { + let r = &d[y * 4..y * 4 + 4]; + let (a, b, c, e) = (r[0] + r[3], r[1] + r[2], r[1] - r[2], r[0] - r[3]); + m[y * 4] = a + b; + m[y * 4 + 1] = e + c; + m[y * 4 + 2] = a - b; + m[y * 4 + 3] = e - c; + } + let mut sum = 0; + for x in 0..4 { + let (r0, r1, r2, r3) = (m[x], m[4 + x], m[8 + x], m[12 + x]); + let (a, b, c, e) = (r0 + r3, r1 + r2, r1 - r2, r0 - r3); + sum += (a + b).abs() + (e + c).abs() + (a - b).abs() + (e - c).abs(); + } + ((sum + 1) >> 1) as u32 +} + +fn hadamard8(v: &mut [i32; 8]) { + let mut t = *v; + let mut h = 1; + while h < 8 { + for i in (0..8).step_by(h * 2) { + for j in i..i + h { + let (a, b) = (t[j], t[j + h]); + t[j] = a + b; + t[j + h] = a - b; + } + } + h *= 2; + } + *v = t; +} + +fn satd8(src: &[u16], ss: usize, pred: &[u16], ps: usize) -> u32 { + let mut m = [[0i32; 8]; 8]; + for y in 0..8 { + for x in 0..8 { + m[y][x] = src[y * ss + x] as i32 - pred[y * ps + x] as i32; + } + hadamard8(&mut m[y]); + } + let mut sum = 0; + for x in 0..8 { + let mut col = [0i32; 8]; + for y in 0..8 { + col[y] = m[y][x]; + } + hadamard8(&mut col); + sum += col.iter().map(|v| v.abs()).sum::(); + } + ((sum + 2) >> 2) as u32 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn last_prefix_groups() { + assert_eq!(last_prefix(3), (3, 0, 0)); + assert_eq!(last_prefix(4), (4, 0, 1)); + assert_eq!(last_prefix(5), (4, 1, 1)); + assert_eq!(last_prefix(7), (5, 1, 1)); + assert_eq!(last_prefix(8), (6, 0, 2)); + assert_eq!(last_prefix(15), (7, 3, 2)); + assert_eq!(last_prefix(16), (8, 0, 3)); + assert_eq!(last_prefix(31), (9, 7, 3)); + } +} diff --git a/crates/hevcenc/src/inter.rs b/crates/hevcenc/src/inter.rs new file mode 100644 index 000000000..23c7a8d40 --- /dev/null +++ b/crates/hevcenc/src/inter.rs @@ -0,0 +1,156 @@ +//! Fractional sample interpolation (8.5.3.3.3) and default weighted prediction (8-262), mirrored +//! from the filmcraft-hevc decoder. Reference sample fetches clamp to the picture (8-228, 8-229). + +use crate::tables::{CHROMA_FILTER, LUMA_FILTER}; + +/// Max prediction block size + filter margin. +const WIN: usize = 64 + 7; + +/// One plane of a reference picture. +#[derive(Clone, Copy)] +pub struct PlaneRef<'a> { + pub data: &'a [u16], + pub w: usize, + pub h: usize, +} + +impl PlaneRef<'_> { + /// Copy a w x h window at (x0, y0) with edge replication. + fn window(&self, x0: i32, y0: i32, w: usize, h: usize, out: &mut [i16], os: usize) { + let (pw, ph) = (self.w as i32, self.h as i32); + let inside_x = x0 >= 0 && x0 + w as i32 <= pw; + for r in 0..h { + let yy = (y0 + r as i32).clamp(0, ph - 1) as usize; + let row = &self.data[yy * self.w..yy * self.w + self.w]; + let dst = &mut out[r * os..r * os + w]; + if inside_x { + for (d, &s) in dst.iter_mut().zip(&row[x0 as usize..x0 as usize + w]) { + *d = s as i16; + } + } else { + for (c, d) in dst.iter_mut().enumerate() { + *d = row[(x0 + c as i32).clamp(0, pw - 1) as usize] as i16; + } + } + } + } +} + +/// Luma prediction samples (14-bit intermediate) of a w x h block at (x, y) displaced by the +/// quarter-sample vector `mv`, into `out` (stride w). +pub fn mc_luma(f: PlaneRef, bd: u32, x: i32, y: i32, mv: [i16; 2], w: usize, h: usize, out: &mut [i16]) { + let (fx, fy) = ((mv[0] & 3) as usize, (mv[1] & 3) as usize); + let xi = x + (mv[0] as i32 >> 2); + let yi = y + (mv[1] as i32 >> 2); + let shift1 = bd.min(12) - 8; + let shift3 = 14 - bd; + let mut win = [0i16; WIN * WIN]; + if fx == 0 && fy == 0 { + f.window(xi, yi, w, h, &mut win, WIN); + for r in 0..h { + for c in 0..w { + out[r * w + c] = win[r * WIN + c] << shift3; + } + } + return; + } + f.window(xi - 3, yi - 3, w + 7, h + 7, &mut win, WIN); + let hf = &LUMA_FILTER[fx]; + let vf = &LUMA_FILTER[fy]; + if fy == 0 { + for r in 0..h { + let row = &win[(r + 3) * WIN..]; + for c in 0..w { + let s: i32 = (0..8).map(|k| hf[k] as i32 * row[c + k] as i32).sum(); + out[r * w + c] = (s >> shift1) as i16; + } + } + } else if fx == 0 { + for r in 0..h { + for c in 0..w { + let s: i32 = (0..8).map(|k| vf[k] as i32 * win[(r + k) * WIN + c + 3] as i32).sum(); + out[r * w + c] = (s >> shift1) as i16; + } + } + } else { + let mut tmp = [0i16; WIN * WIN]; + for r in 0..h + 7 { + let row = &win[r * WIN..]; + for c in 0..w { + let s: i32 = (0..8).map(|k| hf[k] as i32 * row[c + k] as i32).sum(); + tmp[r * WIN + c] = (s >> shift1) as i16; + } + } + for r in 0..h { + for c in 0..w { + let s: i32 = (0..8).map(|k| vf[k] as i32 * tmp[(r + k) * WIN + c] as i32).sum(); + out[r * w + c] = (s >> 6) as i16; + } + } + } +} + +/// Chroma prediction samples of a w x h block at chroma position (x, y) with a 1/8-sample vector. +pub fn mc_chroma(f: PlaneRef, bd: u32, x: i32, y: i32, mv: [i16; 2], w: usize, h: usize, out: &mut [i16]) { + let (fx, fy) = ((mv[0] & 7) as usize, (mv[1] & 7) as usize); + let xi = x + (mv[0] as i32 >> 3); + let yi = y + (mv[1] as i32 >> 3); + let shift1 = bd.min(12) - 8; + let shift3 = 14 - bd; + let mut win = [0i16; WIN * WIN]; + if fx == 0 && fy == 0 { + f.window(xi, yi, w, h, &mut win, WIN); + for r in 0..h { + for c in 0..w { + out[r * w + c] = win[r * WIN + c] << shift3; + } + } + return; + } + f.window(xi - 1, yi - 1, w + 3, h + 3, &mut win, WIN); + let hf = &CHROMA_FILTER[fx]; + let vf = &CHROMA_FILTER[fy]; + if fy == 0 { + for r in 0..h { + let row = &win[(r + 1) * WIN..]; + for c in 0..w { + let s: i32 = (0..4).map(|k| hf[k] as i32 * row[c + k] as i32).sum(); + out[r * w + c] = (s >> shift1) as i16; + } + } + } else if fx == 0 { + for r in 0..h { + for c in 0..w { + let s: i32 = (0..4).map(|k| vf[k] as i32 * win[(r + k) * WIN + c + 1] as i32).sum(); + out[r * w + c] = (s >> shift1) as i16; + } + } + } else { + let mut tmp = [0i16; WIN * WIN]; + for r in 0..h + 3 { + let row = &win[r * WIN..]; + for c in 0..w { + let s: i32 = (0..4).map(|k| hf[k] as i32 * row[c + k] as i32).sum(); + tmp[r * WIN + c] = (s >> shift1) as i16; + } + } + for r in 0..h { + for c in 0..w { + let s: i32 = (0..4).map(|k| vf[k] as i32 * tmp[(r + k) * WIN + c] as i32).sum(); + out[r * w + c] = (s >> 6) as i16; + } + } + } +} + +/// Default weighted prediction, one list (8-262): 14-bit samples to `bd`-bit samples. +pub fn put_uni(p: &[i16], w: usize, h: usize, bd: u32, dst: &mut [u16], ds: usize) { + let shift = 14 - bd; + let off = if shift > 0 { 1 << (shift - 1) } else { 0 }; + let max = (1i32 << bd) - 1; + for r in 0..h { + for c in 0..w { + dst[r * ds + c] = ((p[r * w + c] as i32 + off) >> shift).clamp(0, max) as u16; + } + } +} diff --git a/crates/hevcenc/src/intra.rs b/crates/hevcenc/src/intra.rs new file mode 100644 index 000000000..f1fe9e4cb --- /dev/null +++ b/crates/hevcenc/src/intra.rs @@ -0,0 +1,254 @@ +//! Intra sample prediction (8.4.4.2), mirrored from the filmcraft-hevc decoder so the encoder +//! reconstructs exactly like a decoder: reference substitution, filtering (incl. strong intra +//! smoothing), planar, DC and the 33 angular modes. + +use crate::tables::{INTRA_PRED_ANGLE, INV_ANGLE}; + +/// Reference samples of an nTbS block: `left[0]` = `top[0]` = p[-1][-1], `left[1 + y]` = p[-1][y], +/// `top[1 + x]` = p[x][-1] for x, y = 0..2*nTbS-1. +#[derive(Clone)] +pub struct Refs { + pub left: [u16; 129], + pub top: [u16; 129], +} + +impl Default for Refs { + fn default() -> Self { + Refs { left: [0; 129], top: [0; 129] } + } +} + +/// Substitution process (8.4.4.2.2). `avail_left[1 + y]` / `avail_top[1 + x]` (index 0 = corner) tell +/// which samples are available; unavailable ones are replaced. +pub fn substitute(r: &mut Refs, avail_left: &[bool; 129], avail_top: &[bool; 129], n: usize, bit_depth: u32) { + let n2 = 2 * n; + // Sequence order: p[-1][2n-1] .. p[-1][-1], p[0][-1] .. p[2n-1][-1] + let any = avail_left[..=n2].iter().any(|&a| a) || avail_top[1..=n2].iter().any(|&a| a); + if !any { + let v = 1u16 << (bit_depth - 1); + r.left[..=n2].fill(v); + r.top[..=n2].fill(v); + return; + } + if !avail_left[n2] { + // search upwards then along the top + let mut found = None; + for y in (0..n2).rev() { + if avail_left[y] { + found = Some(r.left[y]); + break; + } + } + if found.is_none() { + for x in 1..=n2 { + if avail_top[x] { + found = Some(r.top[x]); + break; + } + } + } + r.left[n2] = found.unwrap_or(0); + } + for y in (0..n2).rev() { + if !avail_left[y] { + r.left[y] = r.left[y + 1]; + } + } + r.top[0] = r.left[0]; + for x in 1..=n2 { + if !avail_top[x] { + r.top[x] = r.top[x - 1]; + } + } +} + +/// Filtering process of neighbouring samples (8.4.4.2.3) for luma (and 4:4:4 chroma). +pub fn filter(r: &mut Refs, mode: u32, n: usize, strong_enabled: bool, is_luma: bool, bit_depth: u32) { + if mode == 1 || n == 4 { + return; + } + let min_dist = (mode as i32 - 26).abs().min((mode as i32 - 10).abs()); + let thres = match n { + 8 => 7, + 16 => 1, + _ => 0, + }; + if min_dist <= thres { + return; + } + let n2 = 2 * n; + let c = r.left[0] as i32; + if strong_enabled && is_luma && n == 32 { + let (bl, tr) = (r.left[n2] as i32, r.top[n2] as i32); + let thr = 1 << (bit_depth - 5); + if (c + tr - 2 * r.top[n] as i32).abs() < thr && (c + bl - 2 * r.left[n] as i32).abs() < thr { + for i in 0..63 { + r.left[1 + i] = (((63 - i as i32) * c + (i as i32 + 1) * bl + 32) >> 6) as u16; + r.top[1 + i] = (((63 - i as i32) * c + (i as i32 + 1) * tr + 32) >> 6) as u16; + } + return; + } + } + let (l, t) = (r.left, r.top); + let corner = ((l[1] as u32 + 2 * c as u32 + t[1] as u32 + 2) >> 2) as u16; + for i in 1..n2 { + r.left[i] = ((l[i - 1] as u32 + 2 * l[i] as u32 + l[i + 1] as u32 + 2) >> 2) as u16; + r.top[i] = ((t[i - 1] as u32 + 2 * t[i] as u32 + t[i + 1] as u32 + 2) >> 2) as u16; + } + r.left[0] = corner; + r.top[0] = corner; +} + +/// Predict an n x n block with `mode` into `out` (stride `os`). `edge_filters`: DC / horizontal / +/// vertical boundary smoothing (luma blocks smaller than 32x32). +pub fn predict(r: &Refs, mode: u32, n: usize, edge_filters: bool, bit_depth: u32, out: &mut [u16], os: usize) { + let max = (1i32 << bit_depth) - 1; + match mode { + 0 => { + // planar + let log2 = n.trailing_zeros(); + let tr = r.top[1 + n] as i32; + let bl = r.left[1 + n] as i32; + for y in 0..n { + let l = r.left[1 + y] as i32; + for x in 0..n { + let t = r.top[1 + x] as i32; + let v = ((n - 1 - x) as i32 * l + (x as i32 + 1) * tr + (n - 1 - y) as i32 * t + (y as i32 + 1) * bl + n as i32) >> (log2 + 1); + out[y * os + x] = v as u16; + } + } + } + 1 => { + let log2 = n.trailing_zeros(); + let sum: u32 = r.top[1..=n].iter().map(|&v| v as u32).sum::() + r.left[1..=n].iter().map(|&v| v as u32).sum::(); + let dc = ((sum + n as u32) >> (log2 + 1)) as u16; + for y in 0..n { + out[y * os..y * os + n].fill(dc); + } + if edge_filters && n < 32 { + let d = dc as u32; + out[0] = ((r.left[1] as u32 + 2 * d + r.top[1] as u32 + 2) >> 2) as u16; + for x in 1..n { + out[x] = ((r.top[1 + x] as u32 + 3 * d + 2) >> 2) as u16; + } + for y in 1..n { + out[y * os] = ((r.left[1 + y] as u32 + 3 * d + 2) >> 2) as u16; + } + } + } + _ => { + let angle = INTRA_PRED_ANGLE[mode as usize] as i32; + // reference array ref[-n..=2n] stored at offset n + let mut refa = [0i32; 3 * 64 + 1]; + let base = n as i32; + let (main, side) = if mode >= 18 { (&r.top, &r.left) } else { (&r.left, &r.top) }; + for x in 0..=n { + refa[(base + x as i32) as usize] = main[x] as i32; // ref[x] = p[-1+x][-1] + } + if angle < 0 { + let lim = (n as i32 * angle) >> 5; + if lim < -1 { + let inv = INV_ANGLE[(mode - 11) as usize] as i32; + for x in lim..=-1 { + let idx = -1 + ((x * inv + 128) >> 8); + refa[(base + x) as usize] = side[(idx + 1) as usize] as i32; + } + } + } else { + for x in n + 1..=2 * n { + refa[(base + x as i32) as usize] = main[x] as i32; + } + } + let horizontal = mode < 18; + for j in 0..n { + let pos = (j as i32 + 1) * angle; + let idx = pos >> 5; + let fact = pos & 31; + for i in 0..n { + let k = (base + i as i32 + idx + 1) as usize; + let v = if fact != 0 { ((32 - fact) * refa[k] + fact * refa[k + 1] + 16) >> 5 } else { refa[k] }; + if horizontal { + out[i * os + j] = v as u16; + } else { + out[j * os + i] = v as u16; + } + } + } + if edge_filters && n < 32 { + if mode == 26 { + for y in 0..n { + let v = r.top[1] as i32 + ((r.left[1 + y] as i32 - r.left[0] as i32) >> 1); + out[y * os] = v.clamp(0, max) as u16; + } + } else if mode == 10 { + for x in 0..n { + let v = r.left[1] as i32 + ((r.top[1 + x] as i32 - r.top[0] as i32) >> 1); + out[x] = v.clamp(0, max) as u16; + } + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn refs(f: impl Fn(i32) -> u16) -> Refs { + let mut r = Refs::default(); + r.left[0] = f(0); + r.top[0] = f(0); + for i in 1..129 { + r.left[i] = f(-(i as i32)); + r.top[i] = f(i as i32); + } + r + } + + #[test] + fn dc_and_planar_flat() { + let r = refs(|_| 100); + let mut out = [0u16; 64]; + predict(&r, 1, 8, true, 8, &mut out, 8); + assert!(out.iter().all(|&v| v == 100)); + predict(&r, 0, 8, true, 8, &mut out, 8); + assert!(out.iter().all(|&v| v == 100)); + } + + #[test] + fn pure_vertical_and_horizontal() { + let r = refs(|i| (128 + i) as u16); + let mut out = [0u16; 16]; + predict(&r, 26, 4, false, 8, &mut out, 4); + for y in 0..4 { + for x in 0..4 { + assert_eq!(out[y * 4 + x], r.top[1 + x]); + } + } + predict(&r, 10, 4, false, 8, &mut out, 4); + for y in 0..4 { + for x in 0..4 { + assert_eq!(out[y * 4 + x], r.left[1 + y]); + } + } + // diagonal mode 34 reads top[x + y + 2] + predict(&r, 34, 4, false, 8, &mut out, 4); + assert_eq!(out[4 + 1], r.top[1 + 1 + 1 + 1]); + } + + #[test] + fn substitution() { + let mut r = Refs::default(); + let mut al = [false; 129]; + let at = [false; 129]; + r.left[3] = 77; + al[3] = true; + substitute(&mut r, &al, &at, 4, 8); + assert!(r.left[..=8].iter().all(|&v| v == 77)); + assert!(r.top[..=8].iter().all(|&v| v == 77)); + let mut r = Refs::default(); + substitute(&mut r, &[false; 129], &[false; 129], 4, 10); + assert_eq!(r.top[5], 512); + } +} diff --git a/crates/hevcenc/src/lib.rs b/crates/hevcenc/src/lib.rs index 2d1b9968b..fa3656c6f 100644 --- a/crates/hevcenc/src/lib.rs +++ b/crates/hevcenc/src/lib.rs @@ -1,6 +1,36 @@ -//! Clean-room pure-Rust HEVC (H.265) encoder — API skeleton (implementation in progress). +//! Clean-room pure-Rust HEVC (H.265) encoder, implemented from ITU-T Rec. H.265 (ISO/IEC 23008-2). //! -//! The public API below is the contract `crates/export` codes against; keep it stable. +//! Main and Main 10 profiles, 4:2:0, one slice per picture: IDR pictures and single-reference P +//! pictures (no reordering), 64x64 CTUs with rate-distortion coding-unit decisions (8x8 .. 64x64), +//! intra (35 modes, NxN at 8x8), merge / skip and AMVP inter prediction with quarter-sample motion, +//! DCT / DST transforms and CABAC. In-loop filters (deblocking, SAO) are disabled. +//! +//! ``` +//! use effectcraft_hevcenc::{Encoder, EncoderConfig, Frame}; +//! let mut enc = Encoder::new(EncoderConfig::new(64, 48, 30, 1)).unwrap(); +//! let (y, c) = (vec![128u16; 64 * 48], vec![128u16; 32 * 24]); +//! let packet = enc.encode(&Frame { y: &y, u: &c, v: &c, y_stride: 64, uv_stride: 32 }); +//! assert!(packet.keyframe && !packet.data.is_empty()); +//! let _hvcc = enc.hvcc(); +//! ``` + +// Index loops over fixed-size blocks read more clearly than iterator chains in codec code. +#![allow(clippy::needless_range_loop, clippy::too_many_arguments)] + +mod bitstream; +mod cabac; +mod frame; +mod inter; +mod intra; +mod me; +mod params; +#[rustfmt::skip] +mod tables; +mod transform; + +use bitstream::{BitWriter, NAL_IDR_N_LP, NAL_TRAIL_R, nal_unit}; +use frame::{FrameCoder, SliceKind}; +use params::{LOG2_MAX_POC_LSB, LOG2_MIN_CB, MAX_MERGE_CAND, SeqParams}; /// Profile written in the parameter sets. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] @@ -19,6 +49,12 @@ impl Profile { Profile::Main10 => 10, } } + fn idc(self) -> u8 { + match self { + Profile::Main => 1, + Profile::Main10 => 2, + } + } } /// Rate control. @@ -33,6 +69,7 @@ pub enum RateControl { /// Encoder settings. #[derive(Clone, Debug, PartialEq)] pub struct EncoderConfig { + /// Picture size; must be even (4:2:0 cropping is in units of two luma samples). pub width: u32, pub height: u32, /// Frame rate numerator / denominator (VUI timing and rate control). @@ -42,7 +79,7 @@ pub struct EncoderConfig { /// `general_level_idc` (30 × level, e.g. 93 = level 3.1); `None` picks the lowest level that fits. pub level_idc: Option, pub rate: RateControl, - /// Frames between IDR pictures (1 = all intra). + /// Frames between IDR pictures (1 = all intra, 0 = only the first picture). pub keyint: u32, /// BT.709 limited range is signalled in the VUI unless this is set. pub full_range: bool, @@ -65,6 +102,7 @@ impl EncoderConfig { } /// One input picture, 4:2:0, samples at the profile's bit depth (8-bit values 0–255 for Main). +/// Luma is `width` x `height`, chroma `width / 2` x `height / 2`. #[derive(Clone, Copy, Debug)] pub struct Frame<'a> { pub y: &'a [u16], @@ -98,9 +136,67 @@ impl std::fmt::Display for Error { impl std::error::Error for Error {} +/// Frame-level rate control for [`RateControl::Bitrate`]: a bits ≈ X / Qstep model per picture type +/// and a buffer that spreads over- / undershoot over the following pictures. +struct RateCtl { + target: f64, + /// sum(actual) - sum(target), bits. + buffer: f64, + /// Complexity X = bits * Qstep for I and P pictures. + x: [Option; 2], + last_qp: [i32; 2], + init_qp: i32, +} + +fn qstep(qp: i32) -> f64 { + 2f64.powf((qp as f64 - 4.0) / 6.0) +} + +impl RateCtl { + fn new(kbps: u32, cfg: &EncoderConfig) -> Self { + let fps = cfg.fps_num as f64 / cfg.fps_den.max(1) as f64; + let target = kbps as f64 * 1000.0 / fps.max(0.001); + let bpp = target / (cfg.width as f64 * cfg.height as f64); + let init_qp = (28.0 - 6.0 * (bpp / 0.08).log2()).round().clamp(10.0, 46.0) as i32; + RateCtl { target, buffer: 0.0, x: [None, None], last_qp: [init_qp - 3, init_qp], init_qp } + } + + fn qp(&self, idr: bool) -> i32 { + let t = idr as usize ^ 1; // 0 = I, 1 = P + let tp = (self.target - self.buffer / 12.0).clamp(0.3 * self.target, 2.0 * self.target); + let want = if idr { 4.0 * tp } else { tp }; + let x = match (self.x[t], self.x[0]) { + (Some(x), _) => x, + (None, Some(xi)) => xi / 5.0, + (None, None) => return if idr { self.init_qp - 3 } else { self.init_qp }, + }; + let qp = (4.0 + 6.0 * (x / want).log2()).round() as i32; + qp.clamp(self.last_qp[t] - 6, self.last_qp[t] + 6).clamp(1, 51) + } + + fn update(&mut self, idr: bool, qp: i32, bits: f64) { + let t = idr as usize ^ 1; + let x = bits * qstep(qp); + self.x[t] = Some(match self.x[t] { + Some(old) => 0.5 * old + 0.5 * x, + None => x, + }); + self.last_qp[t] = qp; + self.buffer += bits - self.target; + } +} + /// HEVC encoder. Pictures come out in input order (no reordering: I and P only). pub struct Encoder { cfg: EncoderConfig, + seq: SeqParams, + sets: [Vec; 3], + frames: u64, + /// POC of the previous picture (relative to the last IDR). + poc: u32, + reference: Option<[Vec; 3]>, + recon: Option<[Vec; 3]>, + rc: Option, } impl Encoder { @@ -108,21 +204,151 @@ impl Encoder { if cfg.width == 0 || cfg.height == 0 { return Err(Error::InvalidConfig("empty frame".into())); } - Ok(Encoder { cfg }) + if !cfg.width.is_multiple_of(2) || !cfg.height.is_multiple_of(2) { + return Err(Error::InvalidConfig(format!("{}x{}: 4:2:0 HEVC needs even dimensions", cfg.width, cfg.height))); + } + if cfg.width > 16384 || cfg.height > 16384 { + return Err(Error::InvalidConfig("picture larger than 16384 samples".into())); + } + if cfg.fps_num == 0 || cfg.fps_den == 0 { + return Err(Error::InvalidConfig("zero frame rate".into())); + } + match cfg.rate { + RateControl::ConstantQp(q) if q > 51 => return Err(Error::InvalidConfig(format!("QP {q} > 51"))), + RateControl::Bitrate { kbps: 0 } => return Err(Error::InvalidConfig("zero bitrate".into())), + _ => {} + } + let min_cb = 1u32 << LOG2_MIN_CB; + let (coded_width, coded_height) = (cfg.width.div_ceil(min_cb) * min_cb, cfg.height.div_ceil(min_cb) * min_cb); + let kbps = match cfg.rate { + RateControl::Bitrate { kbps } => Some(kbps), + RateControl::ConstantQp(_) => None, + }; + let level_idc = cfg.level_idc.unwrap_or_else(|| params::choose_level(coded_width, coded_height, cfg.fps_num, cfg.fps_den, kbps)); + let seq = SeqParams { + coded_width, + coded_height, + width: cfg.width, + height: cfg.height, + bit_depth: cfg.profile.bit_depth() as u32, + profile_idc: cfg.profile.idc(), + level_idc, + full_range: cfg.full_range, + fps_num: cfg.fps_num, + fps_den: cfg.fps_den, + }; + let sets = [params::vps(&seq), params::sps(&seq), params::pps()]; + let rc = kbps.map(|k| RateCtl::new(k, &cfg)); + Ok(Encoder { cfg, seq, sets, frames: 0, poc: 0, reference: None, recon: None, rc }) } + pub fn config(&self) -> &EncoderConfig { &self.cfg } - /// VPS, SPS, PPS NAL units (without start codes or length prefixes). + + /// VPS, SPS, PPS NAL units (2-byte NAL header, emulation prevention, no start codes or length + /// prefixes). pub fn parameter_sets(&self) -> [Vec; 3] { - [vec![], vec![], vec![]] + self.sets.clone() } + /// HEVCDecoderConfigurationRecord (`hvcC` payload) with the parameter sets. pub fn hvcc(&self) -> Vec { - vec![] + params::hvcc(&self.seq, &self.sets) } + + /// `general_level_idc` written in the parameter sets. + pub fn level_idc(&self) -> u8 { + self.seq.level_idc + } + + /// Reconstruction of the last encoded picture (what a decoder outputs), cropped planes + /// `[Y, Cb, Cr]` with strides `width` and `width / 2`. For tests. + #[doc(hidden)] + pub fn last_reconstruction(&self) -> Option<&[Vec; 3]> { + self.recon.as_ref() + } + + /// Copy the input into coded-size planes, replicating the right / bottom edges. + fn pad_input(&self, f: &Frame) -> [Vec; 3] { + let max = (1u32 << self.seq.bit_depth) - 1; + let mut out: [Vec; 3] = Default::default(); + for (c, plane) in out.iter_mut().enumerate() { + let (src, stride, w, h, cw, ch) = if c == 0 { + (f.y, f.y_stride, self.seq.width, self.seq.height, self.seq.coded_width, self.seq.coded_height) + } else { + let s = if c == 1 { f.u } else { f.v }; + (s, f.uv_stride, self.seq.width / 2, self.seq.height / 2, self.seq.coded_width / 2, self.seq.coded_height / 2) + }; + let (w, h, cw, ch) = (w as usize, h as usize, cw as usize, ch as usize); + assert!(stride >= w && src.len() >= (h - 1) * stride + w, "HEVC encoder: plane {c} smaller than {w}x{h} (stride {stride})"); + plane.reserve(cw * ch); + for y in 0..ch { + let row = &src[y.min(h - 1) * stride..][..w]; + plane.extend(row.iter().map(|&v| (v as u32).min(max) as u16)); + let last = *plane.last().unwrap(); + plane.extend(std::iter::repeat_n(last, cw - w)); + } + } + out + } + /// Encodes one picture. - pub fn encode(&mut self, _frame: &Frame) -> Packet { - Packet { data: vec![], keyframe: true } + pub fn encode(&mut self, frame: &Frame) -> Packet { + let src = self.pad_input(frame); + let keyint = self.cfg.keyint as u64; + let idr = self.reference.is_none() || (keyint > 0 && self.frames.is_multiple_of(keyint)); + self.poc = if idr { 0 } else { self.poc + 1 }; + let qp = match (&self.rc, self.cfg.rate) { + (Some(rc), _) => rc.qp(idr), + (None, RateControl::ConstantQp(q)) => q as i32, + (None, RateControl::Bitrate { .. }) => 28, + }; + let (w, h) = (self.seq.coded_width as usize, self.seq.coded_height as usize); + let bd = self.seq.bit_depth; + let kind = if idr { SliceKind::I } else { SliceKind::P }; + let me = if idr { None } else { self.reference.as_ref().map(|r| me::MeRef::new(&r[0], w, h, bd)) }; + let refp = if idr { None } else { self.reference.as_ref() }; + let mut fc = FrameCoder::new(bd, w, h, kind, qp, &src, refp, me.as_ref()); + + // slice_segment_header() (7.3.6.1) + let mut bw = BitWriter::new(); + bw.put(1, 1); // first_slice_segment_in_pic_flag + if idr { + bw.put(0, 1); // no_output_of_prior_pics_flag + } + bw.ue(0); // slice_pic_parameter_set_id + bw.ue(if idr { 2 } else { 1 }); // slice_type I / P + if !idr { + bw.put(self.poc & ((1 << LOG2_MAX_POC_LSB) - 1), LOG2_MAX_POC_LSB); // slice_pic_order_cnt_lsb + bw.put(1, 1); // short_term_ref_pic_set_sps_flag (the only SPS set, no index coded) + bw.put(0, 1); // num_ref_idx_active_override_flag + bw.ue(5 - MAX_MERGE_CAND as u32); // five_minus_max_num_merge_cand + } + bw.se(qp - 26); // slice_qp_delta (init_qp = 26) + // byte_alignment() + bw.put(1, 1); + bw.align_zero(); + + let ctx = cabac::init_contexts(qp, if idr { 0 } else { 1 }); + let mut enc = cabac::CabacEncoder::new(bw, ctx); + fc.encode(&mut enc); + let rbsp = enc.finish().into_bytes(); + let nal = nal_unit(if idr { NAL_IDR_N_LP } else { NAL_TRAIL_R }, &rbsp); + let mut data = Vec::with_capacity(nal.len() + 4); + data.extend_from_slice(&(nal.len() as u32).to_be_bytes()); + data.extend_from_slice(&nal); + + if let Some(rc) = &mut self.rc { + rc.update(idr, qp, data.len() as f64 * 8.0); + } + // output picture: cropped reconstruction + let rec = std::mem::take(&mut fc.rec); + let (ow, oh) = (self.seq.width as usize, self.seq.height as usize); + let crop = |p: &[u16], stride: usize, w: usize, h: usize| -> Vec { (0..h).flat_map(|y| p[y * stride..y * stride + w].iter().copied()).collect() }; + self.recon = Some([crop(&rec[0], w, ow, oh), crop(&rec[1], w / 2, ow / 2, oh / 2), crop(&rec[2], w / 2, ow / 2, oh / 2)]); + self.reference = Some(rec); + self.frames += 1; + Packet { data, keyframe: idr } } } diff --git a/crates/hevcenc/src/me.rs b/crates/hevcenc/src/me.rs new file mode 100644 index 000000000..5f4cb2ce3 --- /dev/null +++ b/crates/hevcenc/src/me.rs @@ -0,0 +1,144 @@ +//! Motion search reference: the previous reconstructed luma plane padded by edge replication, plus +//! half-sample planes interpolated with the 8-tap half-sample filter. Quarter-sample positions are +//! approximated by averaging neighbouring half-grid samples. Used for estimation only; the coded +//! prediction always uses the normative interpolation in `inter`. + +use crate::tables::LUMA_FILTER; + +const MARGIN: i32 = 80; + +pub struct MeRef { + /// [integer, horizontal half, vertical half, diagonal half] samples, each `stride` x `rows`. + planes: [Vec; 4], + stride: usize, + rows: usize, + w: i32, + h: i32, +} + +impl MeRef { + pub fn new(y: &[u16], w: usize, h: usize, bd: u32) -> MeRef { + let m = MARGIN as usize; + let (stride, rows) = (w + 2 * m, h + 2 * m); + let mut pad = vec![0u16; stride * rows]; + for r in 0..rows { + let sy = (r as i32 - MARGIN).clamp(0, h as i32 - 1) as usize; + let src = &y[sy * w..sy * w + w]; + let row = &mut pad[r * stride..r * stride + stride]; + row[..m].fill(src[0]); + row[m..m + w].copy_from_slice(src); + row[m + w..].fill(src[w - 1]); + } + let max = (1i32 << bd) - 1; + let f = &LUMA_FILTER[2]; + let filt_h = |p: &[u16]| -> Vec { + let mut o = vec![0u16; stride * rows]; + for r in 0..rows { + let row = &p[r * stride..r * stride + stride]; + for x in 0..stride { + let mut s = 0i32; + for k in 0..8 { + let xx = (x as i32 - 3 + k as i32).clamp(0, stride as i32 - 1) as usize; + s += f[k] as i32 * row[xx] as i32; + } + o[r * stride + x] = ((s + 32) >> 6).clamp(0, max) as u16; + } + } + o + }; + let filt_v = |p: &[u16]| -> Vec { + let mut o = vec![0u16; stride * rows]; + let mut acc = vec![0i32; stride]; + for r in 0..rows { + acc.fill(0); + for k in 0..8 { + let rr = (r as i32 - 3 + k as i32).clamp(0, rows as i32 - 1) as usize; + let src = &p[rr * stride..rr * stride + stride]; + let coef = f[k] as i32; + for (a, &s) in acc.iter_mut().zip(src) { + *a += coef * s as i32; + } + } + for (d, &a) in o[r * stride..r * stride + stride].iter_mut().zip(&acc) { + *d = ((a + 32) >> 6).clamp(0, max) as u16; + } + } + o + }; + let hp = filt_h(&pad); + let vp = filt_v(&pad); + let dp = filt_v(&hp); + MeRef { planes: [pad, hp, vp, dp], stride, rows, w: w as i32, h: h as i32 } + } + + /// Integer motion vector range keeping an n x n block (plus filter margin) inside the padding. + pub fn int_range(&self, x: i32, y: i32, n: usize) -> (i32, i32, i32, i32) { + let n = n as i32; + let lo_x = -MARGIN + 4 - x; + let hi_x = self.w + MARGIN - n - 4 - x; + let lo_y = -MARGIN + 4 - y; + let hi_y = self.h + MARGIN - n - 4 - y; + (lo_x.max(-256), hi_x.min(256), lo_y.max(-256), hi_y.min(256)) + } + + /// SAD of the source block against the integer-position reference block at (rx, ry). + pub fn sad(&self, src: &[u16], ss: usize, rx: i32, ry: i32, n: usize) -> u32 { + let p = &self.planes[0]; + let base = (ry + MARGIN) as usize * self.stride + (rx + MARGIN) as usize; + let mut s = 0u32; + for r in 0..n { + let a = &src[r * ss..r * ss + n]; + let b = &p[base + r * self.stride..base + r * self.stride + n]; + s += a.iter().zip(b).map(|(&x, &y)| (x as i32 - y as i32).unsigned_abs()).sum::(); + } + s + } + + /// Approximate luma prediction of an n x n block at (x, y) with quarter-sample vector `mv`. + pub fn approx_pred(&self, x: i32, y: i32, n: usize, mv: [i16; 2], out: &mut [u16]) { + // per axis: one or two (half-sample fraction, integer offset) taps + let axis = |v: i32| -> [(usize, i32); 2] { + if v & 1 == 0 { + let hv = v >> 1; + [((hv & 1) as usize, hv >> 1); 2] + } else { + let (a, b) = ((v - 1) >> 1, (v + 1) >> 1); + [((a & 1) as usize, a >> 1), ((b & 1) as usize, b >> 1)] + } + }; + let (mx, my) = (mv[0] as i32, mv[1] as i32); + let ax = axis(mx); + let ay = axis(my); + let two = mx & 1 != 0 || my & 1 != 0; + let taps = [(ax[0], ay[0]), (ax[1], ay[1])]; + let ntaps = if two { 2 } else { 1 }; + let mut first = true; + for &((fx, ox), (fy, oy)) in taps.iter().take(ntaps) { + let plane = &self.planes[fx + 2 * fy]; + let (bx, by) = (x + ox + MARGIN, y + oy + MARGIN); + let inside = bx >= 0 && by >= 0 && bx as usize + n <= self.stride && by as usize + n <= self.rows; + for r in 0..n { + let dst = &mut out[r * n..r * n + n]; + if inside { + let o = (by as usize + r) * self.stride + bx as usize; + let src = &plane[o..o + n]; + if first { + dst.copy_from_slice(src); + } else { + for (d, &s) in dst.iter_mut().zip(src) { + *d = ((*d as u32 + s as u32 + 1) >> 1) as u16; + } + } + } else { + let yy = (by + r as i32).clamp(0, self.rows as i32 - 1) as usize; + for (c, d) in dst.iter_mut().enumerate() { + let xx = (bx + c as i32).clamp(0, self.stride as i32 - 1) as usize; + let s = plane[yy * self.stride + xx]; + *d = if first { s } else { ((*d as u32 + s as u32 + 1) >> 1) as u16 }; + } + } + } + first = false; + } + } +} diff --git a/crates/hevcenc/src/params.rs b/crates/hevcenc/src/params.rs new file mode 100644 index 000000000..df3f186d9 --- /dev/null +++ b/crates/hevcenc/src/params.rs @@ -0,0 +1,260 @@ +//! Parameter sets (7.3.2.1 VPS, 7.3.2.2 SPS with E.2.1 VUI, 7.3.2.3 PPS), profile_tier_level (7.3.3), +//! level selection (Annex A, Table A.8) and the `hvcC` record (ISO/IEC 14496-15 §8.3.3.1). + +use crate::bitstream::{BitWriter, NAL_PPS, NAL_SPS, NAL_VPS, nal_unit}; + +/// Fixed coding structure of the encoder. +pub const LOG2_CTB: u32 = 6; +pub const LOG2_MIN_CB: u32 = 3; +pub const LOG2_MIN_TB: u32 = 2; +pub const LOG2_MAX_TB: u32 = 5; +pub const LOG2_MAX_POC_LSB: u32 = 8; +/// MaxNumMergeCand. +pub const MAX_MERGE_CAND: usize = 5; + +/// Sequence-level settings shared by the parameter-set writers. +#[derive(Clone, Debug)] +pub struct SeqParams { + /// Coded (padded) size, multiples of the minimum CB size. + pub coded_width: u32, + pub coded_height: u32, + /// Output (cropped) size. + pub width: u32, + pub height: u32, + pub bit_depth: u32, + pub profile_idc: u8, + pub level_idc: u8, + pub full_range: bool, + pub fps_num: u32, + pub fps_den: u32, +} + +/// (level_idc, MaxLumaPs, MaxLumaSr, MaxBR main tier in kbit/s) of Table A.8 (and A.9 for luma rates). +const LEVELS: [(u8, u64, u64, u64); 13] = [ + (30, 36_864, 552_960, 128), + (60, 122_880, 3_686_400, 1_500), + (63, 245_760, 7_372_800, 3_000), + (90, 552_960, 16_588_800, 6_000), + (93, 983_040, 33_177_600, 10_000), + (120, 2_228_224, 66_846_720, 12_000), + (123, 2_228_224, 133_693_440, 20_000), + (150, 8_912_896, 267_386_880, 25_000), + (153, 8_912_896, 534_773_760, 40_000), + (156, 8_912_896, 1_069_547_520, 60_000), + (180, 35_651_584, 1_069_547_520, 60_000), + (183, 35_651_584, 2_139_095_040, 120_000), + (186, 35_651_584, 4_278_190_080, 240_000), +]; + +/// Lowest level whose picture size, sample rate and (when known) bitrate limits fit. +pub fn choose_level(width: u32, height: u32, fps_num: u32, fps_den: u32, kbps: Option) -> u8 { + let ps = width as u64 * height as u64; + let sr = (ps as f64 * fps_num as f64 / fps_den.max(1) as f64).ceil() as u64; + for &(idc, max_ps, max_sr, max_br) in &LEVELS { + let max_dim = ((max_ps * 8) as f64).sqrt() as u64; + if ps <= max_ps && (width as u64) <= max_dim && (height as u64) <= max_dim && sr <= max_sr && kbps.is_none_or(|k| k as u64 <= max_br) { + return idc; + } + } + 186 +} + +fn profile_tier_level(w: &mut BitWriter, p: &SeqParams) { + w.put(0, 2); // general_profile_space + w.put(0, 1); // general_tier_flag (Main tier) + w.put(p.profile_idc as u32, 5); + w.put(compat_flags(p.profile_idc), 32); + w.put(1, 1); // general_progressive_source_flag + w.put(0, 1); // general_interlaced_source_flag + w.put(0, 1); // general_non_packed_constraint_flag + w.put(1, 1); // general_frame_only_constraint_flag + w.put(0, 32); // general_reserved_zero_43bits ... + w.put(0, 11); + w.put(0, 1); // general_inbld_flag / reserved + w.put(p.level_idc as u32, 8); + // sps_max_sub_layers_minus1 == 0: no sub-layer flags +} + +/// general_profile_compatibility_flag[j] as a 32-bit word (flag 0 is the MSB). A Main stream is also +/// Main 10 compatible. +fn compat_flags(profile_idc: u8) -> u32 { + let mut f = 1u32 << (31 - profile_idc as u32); + if profile_idc == 1 { + f |= 1 << (31 - 2); + } + f +} + +pub fn vps(p: &SeqParams) -> Vec { + let mut w = BitWriter::new(); + w.put(0, 4); // vps_video_parameter_set_id + w.put(1, 1); // vps_base_layer_internal_flag + w.put(1, 1); // vps_base_layer_available_flag + w.put(0, 6); // vps_max_layers_minus1 + w.put(0, 3); // vps_max_sub_layers_minus1 + w.put(1, 1); // vps_temporal_id_nesting_flag + w.put(0xffff, 16); + profile_tier_level(&mut w, p); + w.put(1, 1); // vps_sub_layer_ordering_info_present_flag + w.ue(1); // vps_max_dec_pic_buffering_minus1 + w.ue(0); // vps_max_num_reorder_pics + w.ue(0); // vps_max_latency_increase_plus1 + w.put(0, 6); // vps_max_layer_id + w.ue(0); // vps_num_layer_sets_minus1 + w.put(0, 1); // vps_timing_info_present_flag + w.put(0, 1); // vps_extension_flag + w.trailing(); + nal_unit(NAL_VPS, &w.into_bytes()) +} + +pub fn sps(p: &SeqParams) -> Vec { + let mut w = BitWriter::new(); + w.put(0, 4); // sps_video_parameter_set_id + w.put(0, 3); // sps_max_sub_layers_minus1 + w.put(1, 1); // sps_temporal_id_nesting_flag + profile_tier_level(&mut w, p); + w.ue(0); // sps_seq_parameter_set_id + w.ue(1); // chroma_format_idc 4:2:0 + w.ue(p.coded_width); + w.ue(p.coded_height); + let crop = p.coded_width != p.width || p.coded_height != p.height; + w.flag(crop); // conformance_window_flag + if crop { + // offsets in chroma sample units (SubWidthC = SubHeightC = 2) + w.ue(0); + w.ue((p.coded_width - p.width) / 2); + w.ue(0); + w.ue((p.coded_height - p.height) / 2); + } + w.ue(p.bit_depth - 8); // bit_depth_luma_minus8 + w.ue(p.bit_depth - 8); // bit_depth_chroma_minus8 + w.ue(LOG2_MAX_POC_LSB - 4); + w.put(1, 1); // sps_sub_layer_ordering_info_present_flag + w.ue(1); // sps_max_dec_pic_buffering_minus1 + w.ue(0); // sps_max_num_reorder_pics + w.ue(0); // sps_max_latency_increase_plus1 + w.ue(LOG2_MIN_CB - 3); + w.ue(LOG2_CTB - LOG2_MIN_CB); + w.ue(LOG2_MIN_TB - 2); + w.ue(LOG2_MAX_TB - LOG2_MIN_TB); + w.ue(0); // max_transform_hierarchy_depth_inter + w.ue(0); // max_transform_hierarchy_depth_intra + w.put(0, 1); // scaling_list_enabled_flag + w.put(0, 1); // amp_enabled_flag + w.put(0, 1); // sample_adaptive_offset_enabled_flag + w.put(0, 1); // pcm_enabled_flag + w.ue(1); // num_short_term_ref_pic_sets + // st_ref_pic_set(0): one negative picture at delta POC -1, used by the current picture + w.ue(1); // num_negative_pics + w.ue(0); // num_positive_pics + w.ue(0); // delta_poc_s0_minus1 + w.put(1, 1); // used_by_curr_pic_s0_flag + w.put(0, 1); // long_term_ref_pics_present_flag + w.put(0, 1); // sps_temporal_mvp_enabled_flag + w.put(1, 1); // strong_intra_smoothing_enabled_flag + w.put(1, 1); // vui_parameters_present_flag + vui(&mut w, p); + w.put(0, 1); // sps_extension_present_flag + w.trailing(); + nal_unit(NAL_SPS, &w.into_bytes()) +} + +fn vui(w: &mut BitWriter, p: &SeqParams) { + w.put(0, 1); // aspect_ratio_info_present_flag + w.put(0, 1); // overscan_info_present_flag + w.put(1, 1); // video_signal_type_present_flag + w.put(5, 3); // video_format: unspecified + w.flag(p.full_range); + w.put(1, 1); // colour_description_present_flag + w.put(1, 8); // colour_primaries BT.709 + w.put(1, 8); // transfer_characteristics BT.709 + w.put(1, 8); // matrix_coeffs BT.709 + w.put(0, 1); // chroma_loc_info_present_flag + w.put(0, 1); // neutral_chroma_indication_flag + w.put(0, 1); // field_seq_flag + w.put(0, 1); // frame_field_info_present_flag + w.put(0, 1); // default_display_window_flag + w.put(1, 1); // vui_timing_info_present_flag + w.put(p.fps_den, 32); // vui_num_units_in_tick + w.put(p.fps_num, 32); // vui_time_scale + w.put(0, 1); // vui_poc_proportional_to_timing_flag + w.put(0, 1); // vui_hrd_parameters_present_flag + w.put(0, 1); // bitstream_restriction_flag +} + +pub fn pps() -> Vec { + let mut w = BitWriter::new(); + w.ue(0); // pps_pic_parameter_set_id + w.ue(0); // pps_seq_parameter_set_id + w.put(0, 1); // dependent_slice_segments_enabled_flag + w.put(0, 1); // output_flag_present_flag + w.put(0, 3); // num_extra_slice_header_bits + w.put(0, 1); // sign_data_hiding_enabled_flag + w.put(0, 1); // cabac_init_present_flag + w.ue(0); // num_ref_idx_l0_default_active_minus1 + w.ue(0); // num_ref_idx_l1_default_active_minus1 + w.se(0); // init_qp_minus26 + w.put(0, 1); // constrained_intra_pred_flag + w.put(0, 1); // transform_skip_enabled_flag + w.put(0, 1); // cu_qp_delta_enabled_flag + w.se(0); // pps_cb_qp_offset + w.se(0); // pps_cr_qp_offset + w.put(0, 1); // pps_slice_chroma_qp_offsets_present_flag + w.put(0, 1); // weighted_pred_flag + w.put(0, 1); // weighted_bipred_flag + w.put(0, 1); // transquant_bypass_enabled_flag + w.put(0, 1); // tiles_enabled_flag + w.put(0, 1); // entropy_coding_sync_enabled_flag + w.put(0, 1); // pps_loop_filter_across_slices_enabled_flag + w.put(1, 1); // deblocking_filter_control_present_flag + w.put(0, 1); // deblocking_filter_override_enabled_flag + w.put(1, 1); // pps_deblocking_filter_disabled_flag + w.put(0, 1); // pps_scaling_list_data_present_flag + w.put(0, 1); // lists_modification_present_flag + w.ue(0); // log2_parallel_merge_level_minus2 + w.put(0, 1); // slice_segment_header_extension_present_flag + w.put(0, 1); // pps_extension_present_flag + w.trailing(); + nal_unit(NAL_PPS, &w.into_bytes()) +} + +/// HEVCDecoderConfigurationRecord with one VPS, SPS and PPS (lengthSizeMinusOne = 3). +pub fn hvcc(p: &SeqParams, sets: &[Vec; 3]) -> Vec { + let mut o = vec![1u8]; // configurationVersion + o.push(p.profile_idc); // profile_space 0, tier 0, profile_idc + o.extend_from_slice(&compat_flags(p.profile_idc).to_be_bytes()); + // general_constraint_indicator_flags (48 bits): progressive, interlaced, non_packed, frame_only + o.extend_from_slice(&[0x90, 0, 0, 0, 0, 0]); + o.push(p.level_idc); + o.extend_from_slice(&[0xf0, 0x00]); // reserved + min_spatial_segmentation_idc 0 + o.push(0xfc); // reserved + parallelismType 0 + o.push(0xfc | 1); // reserved + chroma_format_idc 1 + o.push(0xf8 | (p.bit_depth - 8) as u8); + o.push(0xf8 | (p.bit_depth - 8) as u8); + o.extend_from_slice(&[0, 0]); // avgFrameRate unspecified + // constantFrameRate 0, numTemporalLayers 1, temporalIdNested 1, lengthSizeMinusOne 3 + o.push((1 << 3) | (1 << 2) | 3); + o.push(3); // numOfArrays + for (set, ty) in sets.iter().zip([NAL_VPS, NAL_SPS, NAL_PPS]) { + o.push(0x80 | ty); // array_completeness 1 + o.extend_from_slice(&1u16.to_be_bytes()); + o.extend_from_slice(&(set.len() as u16).to_be_bytes()); + o.extend_from_slice(set); + } + o +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn levels() { + assert_eq!(choose_level(1920, 1080, 30, 1, None), 120); + assert_eq!(choose_level(1920, 1080, 60, 1, None), 123); + assert_eq!(choose_level(3840, 2160, 30, 1, None), 150); + assert_eq!(choose_level(1280, 720, 30, 1, None), 93); + assert_eq!(choose_level(64, 64, 30, 1, None), 30); + assert_eq!(choose_level(1920, 1080, 30, 1, Some(15_000)), 123); + } +} diff --git a/crates/hevcenc/src/tables.rs b/crates/hevcenc/src/tables.rs new file mode 100644 index 000000000..5b30edd41 --- /dev/null +++ b/crates/hevcenc/src/tables.rs @@ -0,0 +1,232 @@ +//! Constant tables of ITU-T H.265, mirrored from the first-party, spec-derived `filmcraft-hevc` +//! decoder (same organisation, MIT OR Apache-2.0): CABAC context layout and initValues (9.3.2.2), +//! rangeTabLPS / transIdx (Tables 9-52, 9-53), scans (6.5.3 - 6.5.5), transform matrices (8.6.4.2), +//! intra angles (Tables 8-5, 8-6), interpolation filters (8.5.3.3.3) and the chroma QP mapping (Table 8-10). + +#![allow(dead_code)] + +/// Table 9-52 (same values as H.264) rangeTabLPS[pStateIdx][qCodIRangeIdx]. +#[rustfmt::skip] +pub static RANGE_TAB_LPS: [[u8; 4]; 64] = [ + [128, 176, 208, 240], [128, 167, 197, 227], [128, 158, 187, 216], [123, 150, 178, 205], + [116, 142, 169, 195], [111, 135, 160, 185], [105, 128, 152, 175], [100, 122, 144, 166], + [95, 116, 137, 158], [90, 110, 130, 150], [85, 104, 123, 142], [81, 99, 117, 135], + [77, 94, 111, 128], [73, 89, 105, 122], [69, 85, 100, 116], [66, 80, 95, 110], + [62, 76, 90, 104], [59, 72, 86, 99], [56, 69, 81, 94], [53, 65, 77, 89], + [51, 62, 73, 85], [48, 59, 69, 80], [46, 56, 66, 76], [43, 53, 63, 72], + [41, 50, 59, 69], [39, 48, 56, 65], [37, 45, 54, 62], [35, 43, 51, 59], + [33, 41, 48, 56], [32, 39, 46, 53], [30, 37, 43, 50], [29, 35, 41, 48], + [27, 33, 39, 45], [26, 31, 37, 43], [24, 30, 35, 41], [23, 28, 33, 39], + [22, 27, 32, 37], [21, 26, 30, 35], [20, 24, 29, 33], [19, 23, 27, 31], + [18, 22, 26, 30], [17, 21, 25, 28], [16, 20, 23, 27], [15, 19, 22, 25], + [14, 18, 21, 24], [14, 17, 20, 23], [13, 16, 19, 22], [12, 15, 18, 21], + [12, 14, 17, 20], [11, 14, 16, 19], [11, 13, 15, 18], [10, 12, 15, 17], + [10, 12, 14, 16], [9, 11, 13, 15], [9, 11, 12, 14], [8, 10, 12, 14], + [8, 9, 11, 13], [7, 9, 11, 12], [7, 9, 10, 12], [7, 8, 10, 11], + [6, 8, 9, 11], [6, 7, 9, 10], [6, 7, 8, 9], [2, 2, 2, 2], +]; + +/// Table 9-53 transIdxLPS. +#[rustfmt::skip] +static TRANS_IDX_LPS: [u8; 64] = [ + 0, 0, 1, 2, 2, 4, 4, 5, 6, 7, 8, 9, 9, 11, 11, 12, + 13, 13, 15, 15, 16, 16, 18, 18, 19, 19, 21, 21, 22, 22, 23, 24, + 24, 25, 26, 26, 27, 27, 28, 29, 29, 30, 30, 30, 31, 32, 32, 33, + 33, 33, 34, 34, 35, 35, 35, 36, 36, 36, 37, 37, 37, 38, 38, 63, +]; + +/// Table 9-53 transIdxMPS. +#[rustfmt::skip] +static TRANS_IDX_MPS: [u8; 64] = [ + 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, + 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, + 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, + 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 62, 63, +]; + +/// Combined state transition: index = (pStateIdx << 1 | valMPS) * 2 + bin_is_lps. +/// Each context is stored as `pStateIdx << 1 | valMPS`. +pub static NEXT_STATE: [[u8; 2]; 128] = { + let mut t = [[0u8; 2]; 128]; + let mut s = 0; + while s < 128 { + let p = s >> 1; + let mps = s & 1; + // MPS path + t[s][0] = (TRANS_IDX_MPS[p] << 1) | mps as u8; + // LPS path + let new_mps = if p == 0 { 1 - mps } else { mps }; + t[s][1] = (TRANS_IDX_LPS[p] << 1) | new_mps as u8; + s += 1; + } + t +}; + +/// Up-right diagonal scan (6.5.3) of a `N x N` block as (x, y) pairs. +const fn diag(n: usize) -> [(u8, u8); L] { + let mut out = [(0u8, 0u8); L]; + let mut i = 0; + let mut x: i32 = 0; + let mut y: i32 = 0; + while i < L { + while y >= 0 { + if (x as usize) < n && (y as usize) < n { + out[i] = (x as u8, y as u8); + i += 1; + } + y -= 1; + x += 1; + } + y = x; + x = 0; + } + out +} + +/// Horizontal (6.5.4, `vertical == false`) or vertical (6.5.5) scan. +const fn hv(n: usize, vertical: bool) -> [(u8, u8); L] { + let mut out = [(0u8, 0u8); L]; + let mut i = 0; + while i < L { + let (a, b) = ((i % n) as u8, (i / n) as u8); + out[i] = if vertical { (b, a) } else { (a, b) }; + i += 1; + } + out +} + +pub static SCAN_1: [[(u8, u8); 1]; 3] = [[(0, 0)], [(0, 0)], [(0, 0)]]; +pub static SCAN_2: [[(u8, u8); 4]; 3] = [diag::<4>(2), hv::<4>(2, false), hv::<4>(2, true)]; +pub static SCAN_4: [[(u8, u8); 16]; 3] = [diag::<16>(4), hv::<16>(4, false), hv::<16>(4, true)]; +pub static SCAN_8: [[(u8, u8); 64]; 3] = [diag::<64>(8), hv::<64>(8, false), hv::<64>(8, true)]; + +/// ScanOrder[log2BlockSize][scanIdx] (log2BlockSize 0..=3). +#[inline] +pub fn scan_order(log2: u32, scan_idx: usize) -> &'static [(u8, u8)] { + match log2 { + 0 => &SCAN_1[scan_idx], + 1 => &SCAN_2[scan_idx], + 2 => &SCAN_4[scan_idx], + _ => &SCAN_8[scan_idx], + } +} + +/// Up-right diagonal scan for log2 size 2 or 3. +pub fn scan_diag(log2: u32) -> &'static [(u8, u8)] { + scan_order(log2, 0) +} + +/// Table 8-10: QpC as a function of qPi (ChromaArrayType 1). +#[inline] +pub fn qpc_420(qpi: i32) -> i32 { + if qpi < 30 { + qpi + } else if qpi > 43 { + qpi - 6 + } else { + QPC_30_43[(qpi - 30) as usize] as i32 + } +} + +/// Table 9-50 ctxIdxMap for 4x4 sig_coeff_flag contexts. +pub static CTX_IDX_MAP: [u8; 16] = [0, 1, 4, 5, 2, 3, 4, 5, 6, 6, 8, 8, 7, 7, 8, 8]; + +/// levelScale (8-309). +pub static LEVEL_SCALE: [i32; 6] = [40, 45, 51, 57, 64, 72]; + +/// DST-VII 4x4 matrix (8.6.4.2, trType 1). +pub static DST_MATRIX: [[i8; 4]; 4] = [[29, 55, 74, 84], [74, 74, 0, -74], [84, -29, -74, 55], [55, -84, 74, -29]]; + + +// CABAC context layout: offset of each syntax element in the context array (Table 9-4). +pub const SAO_MERGE: usize = 0; +pub const SAO_TYPE: usize = 1; +pub const SPLIT_CU: usize = 2; +pub const CU_TRANSQUANT_BYPASS: usize = 5; +pub const CU_SKIP: usize = 6; +pub const PRED_MODE: usize = 9; +pub const PART_MODE: usize = 10; +pub const PREV_INTRA_LUMA: usize = 14; +pub const INTRA_CHROMA: usize = 15; +pub const RQT_ROOT_CBF: usize = 16; +pub const MERGE_FLAG: usize = 17; +pub const MERGE_IDX: usize = 18; +pub const INTER_PRED_IDC: usize = 19; +pub const REF_IDX: usize = 24; +pub const MVP_FLAG: usize = 26; +pub const SPLIT_TRANSFORM: usize = 27; +pub const CBF_LUMA: usize = 30; +pub const CBF_CHROMA: usize = 32; +pub const ABS_MVD_GT0: usize = 37; +pub const ABS_MVD_GT1: usize = 38; +pub const CU_QP_DELTA_ABS: usize = 39; +pub const TRANSFORM_SKIP: usize = 41; +pub const LAST_X_PREFIX: usize = 43; +pub const LAST_Y_PREFIX: usize = 61; +pub const CODED_SUB_BLOCK: usize = 79; +pub const SIG_COEFF: usize = 83; +pub const GT1: usize = 127; +pub const GT2: usize = 151; +pub const EXPLICIT_RDPCM: usize = 157; +pub const EXPLICIT_RDPCM_DIR: usize = 159; +pub const CU_CHROMA_QP_OFFSET_FLAG: usize = 161; +pub const CU_CHROMA_QP_OFFSET_IDX: usize = 162; +pub const NUM_CTX: usize = 163; + +/// initValue per initType (9.3.2.2). +pub static CABAC_INIT: [[u8; NUM_CTX]; 3] = [ + [153, 200, 139, 141, 157, 154, 154, 154, 154, 154, 184, 154, 154, 154, 184, 63, 154, 154, 154, 154, 154, 154, 154, 154, 154, 154, 154, 153, 138, 138, 111, 141, 94, 138, 182, 154, 154, 154, 154, 154, 154, 139, 139, 110, 110, 124, 125, 140, 153, 125, 127, 140, 109, 111, 143, 127, 111, 79, 108, 123, 63, 110, 110, 124, 125, 140, 153, 125, 127, 140, 109, 111, 143, 127, 111, 79, 108, 123, 63, 91, 171, 134, 141, 111, 111, 125, 110, 110, 94, 124, 108, 124, 107, 125, 141, 179, 153, 125, 107, 125, 141, 179, 153, 125, 107, 125, 141, 179, 153, 125, 140, 139, 182, 182, 152, 136, 152, 136, 153, 136, 139, 111, 136, 139, 111, 141, 111, 140, 92, 137, 138, 140, 152, 138, 139, 153, 74, 149, 92, 139, 107, 122, 152, 140, 179, 166, 182, 140, 227, 122, 197, 138, 153, 136, 167, 152, 152, 154, 154, 154, 154, 154, 154], + [153, 185, 107, 139, 126, 154, 197, 185, 201, 149, 154, 139, 154, 154, 154, 152, 79, 110, 122, 95, 79, 63, 31, 31, 153, 153, 168, 124, 138, 94, 153, 111, 149, 107, 167, 154, 154, 140, 198, 154, 154, 139, 139, 125, 110, 94, 110, 95, 79, 125, 111, 110, 78, 110, 111, 111, 95, 94, 108, 123, 108, 125, 110, 94, 110, 95, 79, 125, 111, 110, 78, 110, 111, 111, 95, 94, 108, 123, 108, 121, 140, 61, 154, 155, 154, 139, 153, 139, 123, 123, 63, 153, 166, 183, 140, 136, 153, 154, 166, 183, 140, 136, 153, 154, 166, 183, 140, 136, 153, 154, 170, 153, 123, 123, 107, 121, 107, 121, 167, 151, 183, 140, 151, 183, 140, 140, 140, 154, 196, 196, 167, 154, 152, 167, 182, 182, 134, 149, 136, 153, 121, 136, 137, 169, 194, 166, 167, 154, 167, 137, 182, 107, 167, 91, 122, 107, 167, 139, 139, 139, 139, 154, 154], + [153, 160, 107, 139, 126, 154, 197, 185, 201, 134, 154, 139, 154, 154, 183, 152, 79, 154, 137, 95, 79, 63, 31, 31, 153, 153, 168, 224, 167, 122, 153, 111, 149, 92, 167, 154, 154, 169, 198, 154, 154, 139, 139, 125, 110, 124, 110, 95, 94, 125, 111, 111, 79, 125, 126, 111, 111, 79, 108, 123, 93, 125, 110, 124, 110, 95, 94, 125, 111, 111, 79, 125, 126, 111, 111, 79, 108, 123, 93, 121, 140, 61, 154, 170, 154, 139, 153, 139, 123, 123, 63, 124, 166, 183, 140, 136, 153, 154, 166, 183, 140, 136, 153, 154, 166, 183, 140, 136, 153, 154, 170, 153, 138, 138, 122, 121, 122, 121, 167, 151, 183, 140, 151, 183, 140, 140, 140, 154, 196, 167, 167, 154, 152, 167, 182, 182, 134, 149, 136, 153, 121, 136, 122, 169, 208, 166, 167, 154, 152, 167, 182, 107, 167, 91, 107, 107, 167, 139, 139, 139, 139, 154, 154], +]; + + +/// Table 8-5: intraPredAngle for predModeIntra 0..34 (0 for planar / DC). +pub static INTRA_PRED_ANGLE: [i8; 35] = [0, 0, 32, 26, 21, 17, 13, 9, 5, 2, 0, -2, -5, -9, -13, -17, -21, -26, -32, -26, -21, -17, -13, -9, -5, -2, 0, 2, 5, 9, 13, 17, 21, 26, 32]; +/// Table 8-6: invAngle for predModeIntra 11..25 (index = mode - 11). +pub static INV_ANGLE: [i16; 15] = [-4096, -1638, -910, -630, -482, -390, -315, -256, -315, -390, -482, -630, -910, -1638, -4096]; + +/// Table 8-10: QpC for qPi = 30..43 (qPi < 30: QpC = qPi; qPi > 43: QpC = qPi - 6). +pub static QPC_30_43: [u8; 14] = [29, 30, 31, 32, 33, 33, 34, 34, 35, 35, 36, 36, 37, 37]; + + +/// Luma 8-tap interpolation filter fL[xFrac][k] (8-236..8-238), taps at offsets -3..4. +pub static LUMA_FILTER: [[i8; 8]; 4] = [[0, 0, 0, 64, 0, 0, 0, 0], [-1, 4, -10, 58, 17, -5, 1, 0], [-1, 4, -11, 40, 40, -11, 4, -1], [0, 1, -5, 17, 58, -10, 4, -1]]; +/// Chroma 4-tap interpolation filter fC[xFrac][k] (8-255..8-261), taps at offsets -1..2. +pub static CHROMA_FILTER: [[i8; 4]; 8] = [[0, 64, 0, 0], [-2, 58, 10, -2], [-4, 54, 16, -2], [-6, 46, 28, -4], [-4, 36, 36, -4], [-4, 28, 46, -6], [-2, 16, 54, -4], [-2, 10, 58, -2]]; + +/// DCT matrix (8-319, 8-321): TRANS_MATRIX[n][m] = transMatrix[m][n] (row n = basis function). +pub static TRANS_MATRIX: [[i8; 32]; 32] = [ + [64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64, 64], + [90, 90, 88, 85, 82, 78, 73, 67, 61, 54, 46, 38, 31, 22, 13, 4, -4, -13, -22, -31, -38, -46, -54, -61, -67, -73, -78, -82, -85, -88, -90, -90], + [90, 87, 80, 70, 57, 43, 25, 9, -9, -25, -43, -57, -70, -80, -87, -90, -90, -87, -80, -70, -57, -43, -25, -9, 9, 25, 43, 57, 70, 80, 87, 90], + [90, 82, 67, 46, 22, -4, -31, -54, -73, -85, -90, -88, -78, -61, -38, -13, 13, 38, 61, 78, 88, 90, 85, 73, 54, 31, 4, -22, -46, -67, -82, -90], + [89, 75, 50, 18, -18, -50, -75, -89, -89, -75, -50, -18, 18, 50, 75, 89, 89, 75, 50, 18, -18, -50, -75, -89, -89, -75, -50, -18, 18, 50, 75, 89], + [88, 67, 31, -13, -54, -82, -90, -78, -46, -4, 38, 73, 90, 85, 61, 22, -22, -61, -85, -90, -73, -38, 4, 46, 78, 90, 82, 54, 13, -31, -67, -88], + [87, 57, 9, -43, -80, -90, -70, -25, 25, 70, 90, 80, 43, -9, -57, -87, -87, -57, -9, 43, 80, 90, 70, 25, -25, -70, -90, -80, -43, 9, 57, 87], + [85, 46, -13, -67, -90, -73, -22, 38, 82, 88, 54, -4, -61, -90, -78, -31, 31, 78, 90, 61, 4, -54, -88, -82, -38, 22, 73, 90, 67, 13, -46, -85], + [83, 36, -36, -83, -83, -36, 36, 83, 83, 36, -36, -83, -83, -36, 36, 83, 83, 36, -36, -83, -83, -36, 36, 83, 83, 36, -36, -83, -83, -36, 36, 83], + [82, 22, -54, -90, -61, 13, 78, 85, 31, -46, -90, -67, 4, 73, 88, 38, -38, -88, -73, -4, 67, 90, 46, -31, -85, -78, -13, 61, 90, 54, -22, -82], + [80, 9, -70, -87, -25, 57, 90, 43, -43, -90, -57, 25, 87, 70, -9, -80, -80, -9, 70, 87, 25, -57, -90, -43, 43, 90, 57, -25, -87, -70, 9, 80], + [78, -4, -82, -73, 13, 85, 67, -22, -88, -61, 31, 90, 54, -38, -90, -46, 46, 90, 38, -54, -90, -31, 61, 88, 22, -67, -85, -13, 73, 82, 4, -78], + [75, -18, -89, -50, 50, 89, 18, -75, -75, 18, 89, 50, -50, -89, -18, 75, 75, -18, -89, -50, 50, 89, 18, -75, -75, 18, 89, 50, -50, -89, -18, 75], + [73, -31, -90, -22, 78, 67, -38, -90, -13, 82, 61, -46, -88, -4, 85, 54, -54, -85, 4, 88, 46, -61, -82, 13, 90, 38, -67, -78, 22, 90, 31, -73], + [70, -43, -87, 9, 90, 25, -80, -57, 57, 80, -25, -90, -9, 87, 43, -70, -70, 43, 87, -9, -90, -25, 80, 57, -57, -80, 25, 90, 9, -87, -43, 70], + [67, -54, -78, 38, 85, -22, -90, 4, 90, 13, -88, -31, 82, 46, -73, -61, 61, 73, -46, -82, 31, 88, -13, -90, -4, 90, 22, -85, -38, 78, 54, -67], + [64, -64, -64, 64, 64, -64, -64, 64, 64, -64, -64, 64, 64, -64, -64, 64, 64, -64, -64, 64, 64, -64, -64, 64, 64, -64, -64, 64, 64, -64, -64, 64], + [61, -73, -46, 82, 31, -88, -13, 90, -4, -90, 22, 85, -38, -78, 54, 67, -67, -54, 78, 38, -85, -22, 90, 4, -90, 13, 88, -31, -82, 46, 73, -61], + [57, -80, -25, 90, -9, -87, 43, 70, -70, -43, 87, 9, -90, 25, 80, -57, -57, 80, 25, -90, 9, 87, -43, -70, 70, 43, -87, -9, 90, -25, -80, 57], + [54, -85, -4, 88, -46, -61, 82, 13, -90, 38, 67, -78, -22, 90, -31, -73, 73, 31, -90, 22, 78, -67, -38, 90, -13, -82, 61, 46, -88, 4, 85, -54], + [50, -89, 18, 75, -75, -18, 89, -50, -50, 89, -18, -75, 75, 18, -89, 50, 50, -89, 18, 75, -75, -18, 89, -50, -50, 89, -18, -75, 75, 18, -89, 50], + [46, -90, 38, 54, -90, 31, 61, -88, 22, 67, -85, 13, 73, -82, 4, 78, -78, -4, 82, -73, -13, 85, -67, -22, 88, -61, -31, 90, -54, -38, 90, -46], + [43, -90, 57, 25, -87, 70, 9, -80, 80, -9, -70, 87, -25, -57, 90, -43, -43, 90, -57, -25, 87, -70, -9, 80, -80, 9, 70, -87, 25, 57, -90, 43], + [38, -88, 73, -4, -67, 90, -46, -31, 85, -78, 13, 61, -90, 54, 22, -82, 82, -22, -54, 90, -61, -13, 78, -85, 31, 46, -90, 67, 4, -73, 88, -38], + [36, -83, 83, -36, -36, 83, -83, 36, 36, -83, 83, -36, -36, 83, -83, 36, 36, -83, 83, -36, -36, 83, -83, 36, 36, -83, 83, -36, -36, 83, -83, 36], + [31, -78, 90, -61, 4, 54, -88, 82, -38, -22, 73, -90, 67, -13, -46, 85, -85, 46, 13, -67, 90, -73, 22, 38, -82, 88, -54, -4, 61, -90, 78, -31], + [25, -70, 90, -80, 43, 9, -57, 87, -87, 57, -9, -43, 80, -90, 70, -25, -25, 70, -90, 80, -43, -9, 57, -87, 87, -57, 9, 43, -80, 90, -70, 25], + [22, -61, 85, -90, 73, -38, -4, 46, -78, 90, -82, 54, -13, -31, 67, -88, 88, -67, 31, 13, -54, 82, -90, 78, -46, 4, 38, -73, 90, -85, 61, -22], + [18, -50, 75, -89, 89, -75, 50, -18, -18, 50, -75, 89, -89, 75, -50, 18, 18, -50, 75, -89, 89, -75, 50, -18, -18, 50, -75, 89, -89, 75, -50, 18], + [13, -38, 61, -78, 88, -90, 85, -73, 54, -31, 4, 22, -46, 67, -82, 90, -90, 82, -67, 46, -22, -4, 31, -54, 73, -85, 90, -88, 78, -61, 38, -13], + [9, -25, 43, -57, 70, -80, 87, -90, 90, -87, 80, -70, 57, -43, 25, -9, -9, 25, -43, 57, -70, 80, -87, 90, -90, 87, -80, 70, -57, 43, -25, 9], + [4, -13, 22, -31, 38, -46, 54, -61, 67, -73, 78, -82, 85, -88, 90, -90, 90, -90, 88, -85, 82, -78, 73, -67, 61, -54, 46, -38, 31, -22, 13, -4], +]; diff --git a/crates/hevcenc/src/transform.rs b/crates/hevcenc/src/transform.rs new file mode 100644 index 000000000..066059f44 --- /dev/null +++ b/crates/hevcenc/src/transform.rs @@ -0,0 +1,278 @@ +//! Forward transforms (encoder side, any integer approximation is allowed), quantisation, and the +//! normative scaling (8.6.3) and inverse transforms (8.6.4.2) the decoder applies. The inverse path +//! is exact integer arithmetic: the even/odd decomposition computes the same sums as the matrix +//! products of the specification. + +use crate::tables::{DST_MATRIX, LEVEL_SCALE, TRANS_MATRIX}; + +#[inline] +fn m(n: usize, k: usize, i: usize) -> i32 { + // k-th basis function of the n-point DCT, sample i + TRANS_MATRIX[k * (32 / n)][i] as i32 +} + +/// 1-D forward DCT of `x` (length n) into `out`. +fn fdct1d(x: &[i32], out: &mut [i32], n: usize) { + if n <= 4 { + for k in 0..n { + let mut s = 0; + for i in 0..n { + s += m(n, k, i) * x[i]; + } + out[k] = s; + } + return; + } + let h = n / 2; + let mut e = [0i32; 16]; + let mut o = [0i32; 16]; + for i in 0..h { + e[i] = x[i] + x[n - 1 - i]; + o[i] = x[i] - x[n - 1 - i]; + } + let mut ev = [0i32; 16]; + fdct1d(&e[..h], &mut ev[..h], h); + for k in 0..h { + out[2 * k] = ev[k]; + let kk = 2 * k + 1; + let mut s = 0; + for i in 0..h { + s += m(n, kk, i) * o[i]; + } + out[kk] = s; + } +} + +/// 1-D inverse DCT: x[i] = sum_k M[k][i] * c[k]. +fn idct1d(c: &[i32], out: &mut [i32], n: usize) { + if n <= 4 { + for i in 0..n { + let mut s = 0; + for k in 0..n { + s += m(n, k, i) * c[k]; + } + out[i] = s; + } + return; + } + let h = n / 2; + let mut ce = [0i32; 16]; + for k in 0..h { + ce[k] = c[2 * k]; + } + let mut e = [0i32; 16]; + idct1d(&ce[..h], &mut e[..h], h); + for i in 0..h { + let mut o = 0; + let mut k = 1; + while k < n { + let v = c[k]; + if v != 0 { + o += m(n, k, i) * v; + } + k += 2; + } + out[i] = e[i] + o; + out[n - 1 - i] = e[i] - o; + } +} + +fn fdst1d(x: &[i32], out: &mut [i32]) { + for k in 0..4 { + out[k] = (0..4).map(|i| DST_MATRIX[k][i] as i32 * x[i]).sum(); + } +} + +fn idst1d(c: &[i32], out: &mut [i32]) { + for i in 0..4 { + out[i] = (0..4).map(|k| DST_MATRIX[k][i] as i32 * c[k]).sum(); + } +} + +/// Forward 2-D transform of the n x n residual `r` (row-major) into `c`. +pub fn forward(r: &[i32], c: &mut [i32], n: usize, dst: bool, bit_depth: u32) { + let log2 = n.trailing_zeros(); + let shift1 = log2 + bit_depth - 9; + let shift2 = log2 + 6; + let (r1, r2) = (1i32 << (shift1 - 1), 1i32 << (shift2 - 1)); + let mut tmp = [0i32; 32 * 32]; + let mut out = [0i32; 32]; + // rows (horizontal frequencies) + for y in 0..n { + let row = &r[y * n..y * n + n]; + if dst { + fdst1d(row, &mut out); + } else { + fdct1d(row, &mut out, n); + } + for k in 0..n { + tmp[k * n + y] = (out[k] + r1) >> shift1; // transposed: tmp[kx][y] + } + } + // columns + let mut col = [0i32; 32]; + for kx in 0..n { + col[..n].copy_from_slice(&tmp[kx * n..kx * n + n]); + if dst { + fdst1d(&col, &mut out); + } else { + fdct1d(&col, &mut out, n); + } + for ky in 0..n { + c[ky * n + kx] = (out[ky] + r2) >> shift2; + } + } +} + +/// Inverse 2-D transform (8.6.4.2) of scaled coefficients `d` in place, producing residuals. +pub fn inverse(d: &mut [i32], n: usize, dst: bool, bit_depth: u32) { + let bd_shift = 20 - bit_depth; + let mut tmp = [0i32; 32 * 32]; + let mut col = [0i32; 32]; + let mut out = [0i32; 32]; + // max non-zero row, to skip empty work + let mut nz_cols = [false; 32]; + for y in 0..n { + for x in 0..n { + if d[y * n + x] != 0 { + nz_cols[x] = true; + } + } + } + // 1. columns + for x in 0..n { + if !nz_cols[x] { + continue; // tmp column stays zero + } + for j in 0..n { + col[j] = d[j * n + x]; + } + if dst { + idst1d(&col, &mut out); + } else { + idct1d(&col, &mut out, n); + } + for y in 0..n { + tmp[y * n + x] = ((out[y] + 64) >> 7).clamp(-32768, 32767); + } + } + // 2. rows + let round = 1 << (bd_shift - 1); + for y in 0..n { + let row = &tmp[y * n..y * n + n]; + if dst { + idst1d(row, &mut out); + } else { + idct1d(row, &mut out, n); + } + for x in 0..n { + d[y * n + x] = (out[x] + round) >> bd_shift; + } + } +} + +const QUANT_SCALE: [i64; 6] = [26214, 23302, 20560, 18396, 16384, 14564]; + +/// Quantise transform coefficients `c` into levels `l` with a dead-zone rounding offset +/// (`intra`: 1/3, otherwise 1/6). `qp` includes QpBdOffset. Returns the number of non-zero levels. +pub fn quantize(c: &[i32], l: &mut [i32], n: usize, qp: i32, intra: bool, bit_depth: u32) -> usize { + let log2 = n.trailing_zeros() as i32; + let tshift = 15 - bit_depth as i32 - log2; + let qbits = 14 + qp / 6 + tshift; + let scale = QUANT_SCALE[(qp % 6) as usize]; + let offset = (if intra { 171i64 } else { 85 }) << (qbits - 9); + let mut nz = 0; + for i in 0..n * n { + let v = c[i] as i64; + let a = ((v.abs() * scale + offset) >> qbits).min(32767); + if a != 0 { + nz += 1; + } + l[i] = if v < 0 { -(a as i32) } else { a as i32 }; + } + nz +} + +/// Scaling process (8.6.2 / 8.6.3, flat scaling matrix): levels to scaled coefficients. +pub fn dequantize(l: &[i32], d: &mut [i32], n: usize, qp: i32, bit_depth: u32) { + let log2 = n.trailing_zeros(); + let bd_shift = bit_depth + log2 - 5; + let round = 1i64 << (bd_shift - 1); + let scale = (LEVEL_SCALE[(qp % 6) as usize] as i64) << (qp / 6); + for i in 0..n * n { + let v = l[i]; + d[i] = if v == 0 { 0 } else { ((v as i64 * 16 * scale + round) >> bd_shift).clamp(-32768, 32767) as i32 }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Direct matrix form of the inverse transform, as written in the specification. + fn inverse_ref(c: &mut [i32], n: usize, dst: bool, bd: u32) { + let mut tmp = vec![0i32; n * n]; + let step = 32 / n; + for x in 0..n { + for y in 0..n { + let mut s = 0i32; + for j in 0..n { + let m = if dst { DST_MATRIX[j][y] as i32 } else { TRANS_MATRIX[j * step][y] as i32 }; + s += m * c[j * n + x]; + } + tmp[y * n + x] = ((s + 64) >> 7).clamp(-32768, 32767); + } + } + let bd_shift = 20 - bd; + for y in 0..n { + for x in 0..n { + let mut s = 0i32; + for j in 0..n { + let m = if dst { DST_MATRIX[j][x] as i32 } else { TRANS_MATRIX[j * step][x] as i32 }; + s += m * tmp[y * n + j]; + } + c[y * n + x] = (s + (1 << (bd_shift - 1))) >> bd_shift; + } + } + } + + #[test] + fn inverse_matches_matrix_form() { + let mut seed = 7u32; + for &(n, dst) in &[(4usize, true), (4, false), (8, false), (16, false), (32, false)] { + for _ in 0..20 { + let mut c = vec![0i32; n * n]; + for v in c.iter_mut() { + seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + if (seed >> 24).is_multiple_of(3) { + *v = ((seed >> 8) % 2001) as i32 - 1000; + } + } + let mut a = c.clone(); + let mut b = c.clone(); + inverse(&mut a, n, dst, 8); + inverse_ref(&mut b, n, dst, 8); + assert_eq!(a, b, "n {n} dst {dst}"); + } + } + } + + #[test] + fn round_trip_is_close() { + for &(n, dst) in &[(4usize, true), (8, false), (32, false)] { + for bd in [8u32, 10] { + let r: Vec = (0..n * n).map(|i| ((i * 37 + i / n * 11) % 61) as i32 - 30).collect(); + let mut c = vec![0; n * n]; + forward(&r, &mut c, n, dst, bd); + let mut l = vec![0; n * n]; + let qp = 4; + quantize(&c, &mut l, n, qp, true, bd); + let mut d = vec![0; n * n]; + dequantize(&l, &mut d, n, qp, bd); + inverse(&mut d, n, dst, bd); + let err: i32 = r.iter().zip(&d).map(|(a, b)| (a - b).abs()).max().unwrap(); + assert!(err <= 2, "n {n} bd {bd} err {err}"); + } + } + } +} diff --git a/crates/hevcenc/tests/roundtrip.rs b/crates/hevcenc/tests/roundtrip.rs new file mode 100644 index 000000000..d2c8a3cc3 --- /dev/null +++ b/crates/hevcenc/tests/roundtrip.rs @@ -0,0 +1,273 @@ +//! Round trips through the filmcraft-hevc decoder and (when installed) ffmpeg: decoded pictures must +//! equal the encoder's own reconstruction bit-exactly. + +use effectcraft_hevcenc::{Encoder, EncoderConfig, Frame, Profile, RateControl}; +use std::process::Command; + +struct Pic { + w: usize, + y: Vec, + u: Vec, + v: Vec, +} + +impl Pic { + fn frame(&self) -> Frame<'_> { + Frame { y: &self.y, u: &self.u, v: &self.v, y_stride: self.w, uv_stride: self.w / 2 } + } +} + +/// Smooth gradients with a few hard-edged shapes, translated by (dx, dy) samples. +fn synth(w: usize, h: usize, bd: u32, dx: i32, dy: i32) -> Pic { + let max = ((1u32 << bd) - 1) as f64; + let mut y = vec![0u16; w * h]; + let mut u = vec![0u16; w * h / 4]; + let mut v = vec![0u16; w * h / 4]; + let luma = |x: i32, yy: i32| -> f64 { + let (fx, fy) = ((x - dx) as f64, (yy - dy) as f64); + let mut l = 0.25 + 0.5 * (fx / 97.0).sin().abs() * 0.6 + 0.3 * (fy / 61.0).cos().powi(2); + // a disc and a rectangle + if (fx - 40.0).powi(2) + (fy - 30.0).powi(2) < 18.0f64.powi(2) { + l = 0.85; + } + if (70.0..110.0).contains(&fx) && (50.0..64.0).contains(&fy) { + l = 0.12; + } + // fine texture + l += 0.03 * ((fx * 0.9).sin() * (fy * 0.7).cos()); + l.clamp(0.0, 1.0) + }; + for yy in 0..h { + for x in 0..w { + y[yy * w + x] = (luma(x as i32, yy as i32) * max).round() as u16; + } + } + for yy in 0..h / 2 { + for x in 0..w / 2 { + let (fx, fy) = ((2 * x) as i32 - dx, (2 * yy) as i32 - dy); + let cu = 0.5 + 0.2 * ((fx as f64) / 53.0).sin(); + let cv = 0.5 + 0.2 * ((fy as f64) / 37.0).cos() * if (fx / 16 + fy / 16) % 2 == 0 { 1.0 } else { -0.5 }; + u[yy * w / 2 + x] = (cu * max).round() as u16; + v[yy * w / 2 + x] = (cv * max).round() as u16; + } + } + Pic { w, y, u, v } +} + +fn psnr(a: &[u16], b: &[u16], bd: u32) -> f64 { + let max = ((1u32 << bd) - 1) as f64; + let mse = a.iter().zip(b).map(|(&x, &y)| (x as f64 - y as f64).powi(2)).sum::() / a.len() as f64; + if mse == 0.0 { 99.0 } else { 10.0 * (max * max / mse).log10() } +} + +struct Run { + packets: Vec>, + keys: Vec, + recons: Vec<[Vec; 3]>, + hvcc: Vec, + sets: [Vec; 3], + psnr_y: Vec, +} + +fn encode(cfg: EncoderConfig, pics: &[Pic]) -> Run { + let bd = cfg.profile.bit_depth() as u32; + let mut enc = Encoder::new(cfg).unwrap(); + let mut run = Run { packets: vec![], keys: vec![], recons: vec![], hvcc: enc.hvcc(), sets: enc.parameter_sets(), psnr_y: vec![] }; + for p in pics { + let pkt = enc.encode(&p.frame()); + let rec = enc.last_reconstruction().unwrap().clone(); + run.psnr_y.push(psnr(&p.y, &rec[0], bd)); + run.packets.push(pkt.data); + run.keys.push(pkt.keyframe); + run.recons.push(rec); + } + run +} + +/// Decode with filmcraft-hevc and compare with the reconstructions. +fn check_filmcraft(run: &Run, w: usize, h: usize) { + let mut dec = filmcraft_hevc::Decoder::from_hvcc(&run.hvcc).unwrap(); + let mut pics = Vec::new(); + for (i, p) in run.packets.iter().enumerate() { + pics.extend(dec.decode(p, i as i64).unwrap()); + } + pics.extend(dec.flush()); + assert_eq!(pics.len(), run.packets.len(), "decoded picture count"); + for (i, pic) in pics.iter().enumerate() { + assert_eq!((pic.width as usize, pic.height as usize), (w, h)); + assert_eq!(pic.pts, i as i64); + let rec = &run.recons[i]; + for (c, plane) in [&pic.y, &pic.u, &pic.v].into_iter().enumerate() { + let (pw, ph, stride) = if c == 0 { (w, h, pic.y_stride) } else { (w / 2, h / 2, pic.uv_stride) }; + for yy in 0..ph { + for x in 0..pw { + let got = plane.get(yy * stride + x); + let want = rec[c][yy * pw + x]; + assert_eq!(got, want, "picture {i} plane {c} at ({x}, {yy})"); + } + } + } + } +} + +fn annexb(run: &Run) -> Vec { + let mut out = Vec::new(); + for s in &run.sets { + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(s); + } + for p in &run.packets { + let mut i = 0; + while i < p.len() { + let n = u32::from_be_bytes([p[i], p[i + 1], p[i + 2], p[i + 3]]) as usize; + out.extend_from_slice(&[0, 0, 0, 1]); + out.extend_from_slice(&p[i + 4..i + 4 + n]); + i += 4 + n; + } + } + out +} + +fn ffmpeg_available() -> bool { + Command::new("ffmpeg").arg("-version").output().map(|o| o.status.success()).unwrap_or(false) +} + +/// Decode the Annex-B stream with ffmpeg and require bit-exact equality and an empty stderr. +fn check_ffmpeg(run: &Run, w: usize, h: usize, bd: u32, tag: &str) { + if !ffmpeg_available() { + eprintln!("ffmpeg not found; skipping the ffmpeg oracle"); + return; + } + let dir = std::env::temp_dir(); + let uniq = format!("{}-{}-{}", tag, std::process::id(), std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap().as_nanos()); + let inp = dir.join(format!("hevcenc-{uniq}.hevc")); + let out = dir.join(format!("hevcenc-{uniq}.yuv")); + std::fs::write(&inp, annexb(run)).unwrap(); + let pix = if bd == 8 { "yuv420p" } else { "yuv420p10le" }; + let o = Command::new("ffmpeg") + .args(["-v", "error", "-xerror", "-y", "-f", "hevc", "-i"]) + .arg(&inp) + .args(["-f", "rawvideo", "-pix_fmt", pix]) + .arg(&out) + .output() + .unwrap(); + let stderr = String::from_utf8_lossy(&o.stderr).to_string(); + let data = std::fs::read(&out).unwrap_or_default(); + let _ = std::fs::remove_file(&inp); + let _ = std::fs::remove_file(&out); + assert!(o.status.success() && stderr.trim().is_empty(), "ffmpeg failed: {stderr}"); + let bps = if bd == 8 { 1 } else { 2 }; + let frame_bytes = (w * h + 2 * (w / 2) * (h / 2)) * bps; + assert_eq!(data.len(), frame_bytes * run.packets.len(), "ffmpeg output size"); + for (i, rec) in run.recons.iter().enumerate() { + let mut off = i * frame_bytes; + for (c, plane) in rec.iter().enumerate() { + let pw = if c == 0 { w } else { w / 2 }; + for (k, &want) in plane.iter().enumerate() { + let got = if bd == 8 { data[off] as u16 } else { u16::from_le_bytes([data[off], data[off + 1]]) }; + assert_eq!(got, want, "ffmpeg: picture {i} plane {c} at ({}, {})", k % pw, k / pw); + off += bps; + } + } + } +} + +fn cfg(w: u32, h: u32, qp: u8, keyint: u32, profile: Profile) -> EncoderConfig { + let mut c = EncoderConfig::new(w, h, 30, 1); + c.rate = RateControl::ConstantQp(qp); + c.keyint = keyint; + c.profile = profile; + c +} + +#[test] +fn intra_8bit_quality_and_decoders() { + let (w, h) = (192, 128); + let pics = vec![synth(w, h, 8, 0, 0)]; + let run = encode(cfg(w as u32, h as u32, 22, 1, Profile::Main), &pics); + eprintln!("intra QP22: {} bytes, PSNR-Y {:.2}", run.packets[0].len(), run.psnr_y[0]); + assert!(run.psnr_y[0] > 35.0, "PSNR {}", run.psnr_y[0]); + check_filmcraft(&run, w, h); + check_ffmpeg(&run, w, h, 8, "intra8"); +} + +#[test] +fn inter_8bit_moving_pattern() { + let (w, h) = (160, 96); + let pics: Vec = (0..6).map(|i| synth(w, h, 8, 3 * i, i)).collect(); + let run = encode(cfg(w as u32, h as u32, 27, 30, Profile::Main), &pics); + let sizes: Vec = run.packets.iter().map(|p| p.len()).collect(); + eprintln!("inter sizes {sizes:?} psnr {:?}", run.psnr_y); + assert_eq!(run.keys, vec![true, false, false, false, false, false]); + let p_avg = sizes[1..].iter().sum::() as f64 / 5.0; + assert!(p_avg < sizes[0] as f64 * 0.6, "P frames {p_avg} vs I {}", sizes[0]); + assert!(run.psnr_y.iter().all(|&p| p > 33.0)); + check_filmcraft(&run, w, h); + check_ffmpeg(&run, w, h, 8, "inter8"); +} + +#[test] +fn odd_multiple_sizes_and_cropping() { + for (w, h) in [(34usize, 18usize), (66, 38), (8, 8), (130, 72)] { + let pics: Vec = (0..3).map(|i| synth(w, h, 8, i, 2 * i)).collect(); + let run = encode(cfg(w as u32, h as u32, 30, 2, Profile::Main), &pics); + check_filmcraft(&run, w, h); + check_ffmpeg(&run, w, h, 8, &format!("size{w}x{h}")); + } +} + +#[test] +fn odd_dimensions_are_rejected() { + assert!(Encoder::new(EncoderConfig::new(33, 17, 30, 1)).is_err()); +} + +#[test] +fn main10_round_trip() { + let (w, h) = (96, 64); + let pics: Vec = (0..4).map(|i| synth(w, h, 10, 2 * i, -i)).collect(); + let run = encode(cfg(w as u32, h as u32, 24, 3, Profile::Main10), &pics); + eprintln!("main10 psnr {:?}", run.psnr_y); + assert!(run.psnr_y[0] > 35.0); + check_filmcraft(&run, w, h); + check_ffmpeg(&run, w, h, 10, "main10"); +} + +#[test] +fn qp_extremes() { + let (w, h) = (64, 64); + for qp in [0u8, 51] { + let pics: Vec = (0..2).map(|i| synth(w, h, 8, i, i)).collect(); + let run = encode(cfg(w as u32, h as u32, qp, 30, Profile::Main), &pics); + check_filmcraft(&run, w, h); + check_ffmpeg(&run, w, h, 8, &format!("qp{qp}")); + } +} + +#[test] +fn bitrate_mode_hits_target() { + let (w, h) = (192, 128); + let pics: Vec = (0..30).map(|i| synth(w, h, 8, i, i / 2)).collect(); + let mut c = EncoderConfig::new(w as u32, h as u32, 30, 1); + c.rate = RateControl::Bitrate { kbps: 400 }; + c.keyint = 15; + let run = encode(c, &pics); + let bits: usize = run.packets.iter().map(|p| p.len() * 8).sum(); + let kbps = bits as f64 / (pics.len() as f64 / 30.0) / 1000.0; + eprintln!("bitrate mode: {kbps:.1} kbit/s for a 400 kbit/s target"); + assert!(kbps > 400.0 / 2.0 && kbps < 400.0 * 2.0, "{kbps}"); + check_filmcraft(&run, w, h); +} + +#[test] +fn hvcc_layout() { + let enc = Encoder::new(cfg(64, 48, 28, 30, Profile::Main10)).unwrap(); + let h = enc.hvcc(); + assert_eq!(h[0], 1); + assert_eq!(h[1] & 31, 2); // Main 10 + assert_eq!(h[21] & 3, 3); // lengthSizeMinusOne + assert_eq!(h[22], 3); // arrays + let sets = enc.parameter_sets(); + assert_eq!(sets[0][0] >> 1, 32); + assert_eq!(sets[1][0] >> 1, 33); + assert_eq!(sets[2][0] >> 1, 34); +}