diff --git a/Cargo.lock b/Cargo.lock index efbea46b..65125f6c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1464,6 +1464,15 @@ dependencies = [ "thiserror 2.0.21", ] +[[package]] +name = "filmcraft-h264enc" +version = "0.1.0" +dependencies = [ + "filmcraft-bitstream", + "rayon", + "thiserror 2.0.21", +] + [[package]] name = "filmcraft-isobmff" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 7da9d7b9..66e04659 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ filmcraft-isobmff = { path = "crates/isobmff" } filmcraft-h264 = { path = "crates/h264" } filmcraft-prores = { path = "crates/prores" } filmcraft-aac = { path = "crates/aac" } +filmcraft-h264enc = { path = "crates/h264enc" } filmcraft-frame = { path = "crates/frame" } filmcraft-media = { path = "crates/media" } filmcraft-codecs = { path = "crates/codecs" } diff --git a/crates/h264enc/Cargo.toml b/crates/h264enc/Cargo.toml new file mode 100644 index 00000000..988eed7f --- /dev/null +++ b/crates/h264enc/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "filmcraft-h264enc" +description = "Clean-room pure-Rust H.264/AVC encoder (High/Main/Baseline, CABAC/CAVLC, B-frames, rate control)" +version.workspace = true +edition.workspace = true +license.workspace = true +rust-version.workspace = true + +[lints] +workspace = true + +[features] +default = ["threads"] +threads = ["dep:rayon"] + +[dependencies] +filmcraft-bitstream = { workspace = true } +thiserror = { workspace = true } +rayon = { workspace = true, optional = true } diff --git a/crates/h264enc/README.md b/crates/h264enc/README.md new file mode 100644 index 00000000..69622e85 --- /dev/null +++ b/crates/h264enc/README.md @@ -0,0 +1,83 @@ +# filmcraft-h264enc + +A clean-room, pure-Rust H.264/AVC encoder for FilmCraft's export path, written from the public +ITU-T H.264 specification (04/2013 edition) and general encoder literature. Layer L0: it depends only on +`filmcraft-bitstream`, `thiserror` and (optionally, feature `threads`, on by default) `rayon`. No `unsafe`, +builds for `wasm32-unknown-unknown` with or without `threads`. + +## Usage + +```rust +use filmcraft_h264enc::{Encoder, EncoderConfig, PacketFormat, Preset, Profile, RateControl, YuvFrame}; + +let mut cfg = EncoderConfig::new(1920, 1080, 30000, 1001); +cfg.profile = Profile::High; +cfg.preset = Preset::Balanced; // Speed / Balanced / Quality +cfg.rate = RateControl::Crf(20.0); // Crf(f32) | Cbr { kbps } | Vbr { target_kbps, max_kbps } | Qp(u8) +cfg.bframes = 2; // 0..=3 +cfg.keyint = 250; +cfg.format = PacketFormat::LengthPrefixed; // MP4 samples; or AnnexB (default) for .h264 +let mut enc = Encoder::new(cfg)?; +let avcc = enc.avcc(); // AVCDecoderConfigurationRecord for the MP4 `avcC` box +for (i, pic) in pictures.iter().enumerate() { + let frame = YuvFrame { y: &pic.y, u: &pic.u, v: &pic.v, y_stride: w, uv_stride: w / 2 }; + for p in enc.encode(&frame, i as i64) { /* p.data, p.pts, p.dts, p.keyframe, p.frame_type */ } +} +for p in enc.flush() { /* ... */ } +``` + +- Input: planar 4:2:0 8-bit (`YuvFrame` with strides); any even size, non-multiple-of-16 sizes are + cropped via the SPS frame cropping. `try_encode` returns an error instead of panicking on short planes. +- Output: one `Packet` per picture in decoding order with `pts` (as given), `dts` (reorder-delayed, always + `<= pts` and strictly increasing), `keyframe` (IDR) and the frame type/QP. Annex-B packets carry an AUD, + SPS/PPS on every IDR and an identification SEI on the first picture. +- Two-pass: encode once with `cfg.pass = Pass::First`, take `enc.pass_stats()` (serialisable with + `PassStats::to_text/from_text`), then encode again with `Pass::Second(stats)`. +- `set_recon_capture(true)` + `take_recon()` expose the encoder's reconstructed pictures (used by the tests to + prove encoder/decoder match). + +## Features + +| Area | Implemented | +|---|---| +| Profiles | Constrained Baseline (CAVLC, I/P), Main (CABAC, I/P/B), High (CABAC, 8x8 transform, I/P/B); 4:2:0 8-bit, frame coding | +| Parameter sets | SPS (level auto-selected from size/rate/DPB/bitrate, cropping, VUI: SAR, colour description + range, timing, bitstream restriction/reorder), PPS, AUD, user-data SEI, avcC | +| Intra | 16x16 (4 modes, SATD), 4x4 (9 modes, SATD + mode-cost, sequential reconstruction), 8x8 (9 modes with reference filtering, High), chroma (4 modes) | +| Inter P | hexagon integer search from predictors, half- then quarter-sample refinement (SATD) on precomputed 6-tap half-sample planes; 16x16/16x8/8x16/8x8 partitions (Balanced/Quality); P_Skip detection with the normative skip MV | +| Inter B | 1–3 non-reference B-frames, spatial direct / B_Skip, L0/L1/Bi 16x16 and (Balanced/Quality) 16x8/8x16 with a per-partition direction; all B-frames of a mini-GOP are encoded concurrently | +| Residual | 4x4 / 8x8 integer transforms, Intra16x16 and chroma DC Hadamards, deadzone quantisation (intra 1/3, inter 1/6), adaptive 4x4/8x8 inter transform choice, coefficient decimation for inter blocks | +| Entropy | CABAC (all I and cabac_init_idc 0 context tables, verified against the spec text) and CAVLC (all coeff_token/total_zeros/run_before tables, level escapes) | +| Loop filter | normative deblocking (bS 0–4 incl. bi-predicted rules, 8x8-transform edges, chroma QP mapping) in the reconstruction loop | +| Mode decision | SATD + λ·bits for Speed/Balanced; Quality adds RD-lite (SSD + λ·CABAC-estimated bits) between the best intra/inter/skip candidates | +| Rate control | `Qp`, `Crf` (constant quality with I/P/B offsets), 1-pass ABR with complexity model + VBV (`Cbr`, 1-pass `Vbr`), closed-loop 2-pass VBR; variance adaptive quantisation (`aq_strength`), lookahead scene-cut IDR insertion | +| Threading | slice-parallel encoding (MB-row slices, `disable_deblocking_filter_idc = 2` with >1 slice so each slice is also deblocked in parallel), parallel half-sample interpolation and lookahead | + +## Verification + +`cargo test --release -p filmcraft-h264enc` (ffmpeg/ffprobe used only as external oracles; tests skip when +absent). Every stream is decoded with `ffmpeg -ec 0 -f rawvideo` and must report no errors, produce exactly the +number of frames encoded, and be **bit-identical to the encoder's reconstruction** — this covers Baseline/Main/High, +all presets, 0–3 B-frames, 1–4 slices, odd sizes, `testsrc2` and `mandelbrot` sources, and length-prefixed output +re-wrapped with the avcC parameter sets. B-frame display order is checked with ffprobe. + +Results on the synthetic content (moving gradient, grain noise, text-like edges, sub-pixel panning): + +| Test | Result | +|---|---| +| 1080p, High, Speed, CRF 20, 2 B | mean luma PSNR 43.0 dB (IDR 39.1, P/B 43.2–44.5) | +| 640x360 CBR 1500 kbit/s | 1518 kbit/s (+1.2%) | +| 640x360 1-pass VBR 1000 kbit/s | 981 kbit/s (−1.9%) | +| 640x360 2-pass VBR 800 kbit/s | 840 kbit/s (+5.0%) | + +Development tool: `cargo run --release -p filmcraft-h264enc --example h264enc_synth -- W H FRAMES +[profile] [preset] [bframes] [qpN|crfN|cbrN|vbrN] [slices] [synth|testsrc2|mandelbrot]` encodes, decodes with +ffmpeg and reports speed, bitrate, PSNR and any encoder/decoder mismatch (`MAP=1` prints a per-MB mismatch map, +`PERF=k` cycles k source pictures for speed runs). + +## Limitations / not yet implemented + +- Single reference frame per list; no B-pyramid, no weighted prediction, no interlaced/MBAFF, no 4:2:2/4:4:4 or + high bit depth, no custom scaling matrices, no sub-8x8 partitions, no B_8x8. +- CAVLC is used only for Baseline (High/Main always use CABAC); 8x8 transform is therefore CABAC-only. +- No trellis quantisation, no psy-RD; no HRD/buffering-period SEI (VBV is enforced internally only). +- Threading is slice-based only (no frame-level pipelining); with many slices, slice edges are not deblocked. diff --git a/crates/h264enc/examples/h264enc_synth.rs b/crates/h264enc/examples/h264enc_synth.rs new file mode 100644 index 00000000..b0435185 --- /dev/null +++ b/crates/h264enc/examples/h264enc_synth.rs @@ -0,0 +1,154 @@ +//! Development tool: encode synthetic (or lavfi) content and verify against ffmpeg. +//! +//! `cargo run --release -p filmcraft-h264enc --example h264enc_synth -- W H FRAMES [profile] [preset] [bframes] [rate] [slices] [src]` +//! profile: baseline|main|high, preset: speed|balanced|quality, rate: qpN | crfN | cbrN | vbrN +#[path = "../tests/common/mod.rs"] +mod common; + +use filmcraft_h264enc::*; +use std::time::Instant; + +fn main() { + let a: Vec = std::env::args().collect(); + let arg = |i: usize, d: &str| a.get(i).cloned().unwrap_or_else(|| d.to_string()); + let w: usize = arg(1, "320").parse().unwrap(); + let h: usize = arg(2, "240").parse().unwrap(); + let n: usize = arg(3, "10").parse().unwrap(); + let mut cfg = EncoderConfig::new(w as u32, h as u32, 30, 1); + cfg.profile = match arg(4, "high").as_str() { + "baseline" => Profile::Baseline, + "main" => Profile::Main, + _ => Profile::High, + }; + cfg.preset = match arg(5, "balanced").as_str() { + "speed" => Preset::Speed, + "quality" => Preset::Quality, + _ => Preset::Balanced, + }; + cfg.bframes = arg(6, "0").parse().unwrap(); + let rate = arg(7, "qp26"); + cfg.rate = if let Some(q) = rate.strip_prefix("qp") { + RateControl::Qp(q.parse().unwrap()) + } else if let Some(q) = rate.strip_prefix("crf") { + RateControl::Crf(q.parse().unwrap()) + } else if let Some(q) = rate.strip_prefix("cbr") { + RateControl::Cbr { kbps: q.parse().unwrap() } + } else { + let k: u32 = rate.strip_prefix("vbr").unwrap().parse().unwrap(); + RateControl::Vbr { target_kbps: k, max_kbps: k * 2 } + }; + cfg.slices = arg(8, "0").parse().unwrap(); + let src_kind = arg(9, "synth"); + if let Ok(t) = std::env::var("THREADS") { + cfg.threads = t.parse().unwrap(); + } + if let Ok(k) = std::env::var("KEYINT") { + cfg.keyint = k.parse().unwrap(); + } + if let Ok(k) = std::env::var("AQ") { + cfg.aq_strength = k.parse().unwrap(); + } + // PERF=k: generate only k distinct pictures and cycle them (for long speed runs without verification). + let distinct = std::env::var("PERF").ok().map_or(n, |k| k.parse::().unwrap().min(n)); + let base: Vec = + if src_kind == "synth" { (0..distinct).map(|t| common::synth(w, h, t)).collect() } else { common::ffmpeg_source(&src_kind, w, h, distinct) }; + let perf = distinct < n; + let frames: Vec<&common::Yuv> = (0..n).map(|i| &base[i % distinct]).collect(); + let dir = common::out_dir("dev"); + let mut enc = Encoder::new(cfg).unwrap(); + enc.set_recon_capture(true); + let mut out = Vec::new(); + let mut pk = Vec::new(); + let t0 = Instant::now(); + if perf { + enc.set_recon_capture(false); + } + for (i, f) in frames.iter().enumerate() { + pk.extend(enc.encode(&f.frame(), i as i64)); + } + pk.extend(enc.flush()); + let el = t0.elapsed().as_secs_f64(); + for p in &pk { + out.extend_from_slice(&p.data); + } + let file = dir.join("out.h264"); + std::fs::write(&file, &out).unwrap(); + let mut recon = enc.take_recon(); + recon.sort_by_key(|r| r.pts); + let kbps = out.len() as f64 * 8.0 / (n as f64 / 30.0) / 1000.0; + println!("encoded {n} frames in {el:.3}s = {:.1} fps, {} bytes, {kbps:.1} kbps", n as f64 / el, out.len()); + for p in pk.iter().take(12) { + print!("{:?}(pts{} dts{} qp{} {}B) ", p.frame_type, p.pts, p.dts, p.qp, p.data.len()); + } + println!(); + if common::ffmpeg().is_none() || perf { + return; + } + let (err, dec) = common::ffmpeg_decode(&file); + if !err.is_empty() { + println!("ffmpeg stderr: {err}"); + } + let fs = w * h * 3 / 2; + println!("ffmpeg decoded {} frames (expected {n})", dec.len() / fs); + let mut src_all = Vec::new(); + let mut psnr_sum = 0.0; + for (i, r) in recon.iter().enumerate() { + let mut rb = r.y.clone(); + rb.extend_from_slice(&r.u); + rb.extend_from_slice(&r.v); + let s = frames[i].bytes(); + src_all.extend_from_slice(&s); + let p = common::psnr(&frames[i].y, &r.y); + psnr_sum += p; + if dec.len() >= (i + 1) * fs { + let d = &dec[i * fs..(i + 1) * fs]; + if d != rb.as_slice() { + let pos = d.iter().zip(&rb).position(|(a, b)| a != b).unwrap(); + let (plane, off) = if pos < w * h { + ("Y", pos) + } else if pos < w * h * 5 / 4 { + ("U", pos - w * h) + } else { + ("V", pos - w * h * 5 / 4) + }; + let pw = if plane == "Y" { w } else { w / 2 }; + let (x, y) = (off % pw, off / pw); + let mbsz = if plane == "Y" { 16 } else { 8 }; + let nd = d.iter().zip(&rb).filter(|(a, b)| a != b).count(); + println!( + "frame {i} (pts {}): MISMATCH first at {plane} ({x},{y}) mb ({},{}) dec={} rec={} ({nd} samples differ)", + r.pts, + x / mbsz, + y / mbsz, + d[pos], + rb[pos] + ); + if std::env::var("MAP").is_ok() { + let (mbw, mbh) = (w.div_ceil(16), h.div_ceil(16)); + for my in 0..mbh { + let mut line = String::new(); + for mx in 0..mbw { + let mut n = 0; + for yy in my * 16..(my * 16 + 13).min(h) { + for xx in mx * 16..(mx * 16 + 13).min(w) { + n += (d[yy * w + xx] != rb[yy * w + xx]) as usize; + } + } + line.push(if n == 0 { + '.' + } else if n < 16 { + 'e' + } else { + 'X' + }); + } + println!(" {line}"); + } + break; + } + } + } + } + println!("mean luma PSNR vs source: {:.2} dB", psnr_sum / recon.len() as f64); + std::fs::write(dir.join("src.yuv"), &src_all).unwrap(); +} diff --git a/crates/h264enc/src/cabac.rs b/crates/h264enc/src/cabac.rs new file mode 100644 index 00000000..3ffd7aac --- /dev/null +++ b/crates/h264enc/src/cabac.rs @@ -0,0 +1,339 @@ +//! CABAC arithmetic encoder (ITU-T H.264 §9.3.4.2) and a matching bit-cost estimator. + +use crate::tables::{NUM_CTX, RANGE_LPS, TRANS_LPS}; + +/// A sink for CABAC bins. Implemented by the real arithmetic coder and by the rate estimator, so the same +/// syntax-element code drives both bitstream writing and RD cost estimation. +pub trait BinSink { + fn decision(&mut self, ctx: usize, bin: u32); + fn bypass(&mut self, bin: u32); + fn terminate(&mut self, bin: u32); +} + +/// Context states: `pStateIdx << 1 | valMPS`. +#[derive(Clone)] +pub struct Contexts(pub [u8; NUM_CTX]); + +impl Contexts { + pub fn init(table: &[(i8, i8); NUM_CTX], slice_qp: i32) -> Self { + let qp = slice_qp.clamp(0, 51); + let mut s = [0u8; NUM_CTX]; + for (i, &(m, n)) in table.iter().enumerate() { + let pre = (((m as i32) * qp) >> 4) + n as i32; + let pre = pre.clamp(1, 126); + s[i] = if pre <= 63 { ((63 - pre) as u8) << 1 } else { (((pre - 64) as u8) << 1) | 1 }; + } + Contexts(s) + } +} + +/// The arithmetic encoder writing into a byte buffer. +#[derive(Clone)] +pub struct CabacEncoder { + pub ctx: Contexts, + low: u32, + range: u32, + outstanding: u32, + first_bit: bool, + buf: Vec, + acc: u64, + nbits: u32, + bins: u64, +} + +impl CabacEncoder { + /// Start encoding after a slice header; `header` must be byte aligned (cabac_alignment_one_bits written). + pub fn new(ctx: Contexts, header: Vec) -> Self { + CabacEncoder { ctx, low: 0, range: 510, outstanding: 0, first_bit: true, buf: header, acc: 0, nbits: 0, bins: 0 } + } + + #[inline] + fn write_bit_raw(&mut self, b: u32) { + self.acc = (self.acc << 1) | b as u64; + self.nbits += 1; + if self.nbits == 32 { + self.buf.extend_from_slice(&(self.acc as u32).to_be_bytes()); + self.acc = 0; + self.nbits = 0; + } + } + + #[inline] + fn put_bit(&mut self, b: u32) { + if self.first_bit { + self.first_bit = false; + } else { + self.write_bit_raw(b); + } + while self.outstanding > 0 { + self.write_bit_raw(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; + } + } + + /// Terminate with end_of_slice_flag = 1 and flush; returns the RBSP bytes (with rbsp trailing bits). + pub fn finish(mut self) -> Vec { + self.terminate(1); + // EncodeFlush wrote the stop bit as the last bit; pad with zeros to byte alignment. + while !self.nbits.is_multiple_of(8) { + self.write_bit_raw(0); + } + let nbytes = self.nbits / 8; + for i in 0..nbytes { + self.buf.push((self.acc >> (8 * (nbytes - 1 - i))) as u8); + } + self.buf + } +} + +impl BinSink for CabacEncoder { + #[inline] + fn decision(&mut self, ctx: usize, bin: u32) { + self.bins += 1; + let s = self.ctx.0[ctx]; + let state = (s >> 1) as usize; + let mps = (s & 1) as u32; + let lps = RANGE_LPS[state][((self.range >> 6) & 3) as usize] as u32; + self.range -= lps; + if bin != mps { + self.low += self.range; + self.range = lps; + let new_mps = if state == 0 { 1 - mps } else { mps }; + self.ctx.0[ctx] = (TRANS_LPS[state] << 1) | new_mps as u8; + } else if state < 62 { + self.ctx.0[ctx] = (((state + 1) as u8) << 1) | mps as u8; + } + self.renorm(); + } + + #[inline] + fn bypass(&mut self, bin: u32) { + self.bins += 1; + 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.bins += 1; + self.range -= 2; + if bin != 0 { + self.low += self.range; + // EncodeFlush + self.range = 2; + self.renorm(); + self.put_bit((self.low >> 9) & 1); + let v = ((self.low >> 7) & 3) | 1; + self.write_bit_raw((v >> 1) & 1); + self.write_bit_raw(v & 1); + } else { + self.renorm(); + } + } +} + +/// Fractional-bit cost table: `COST[state][is_mps]` in 1/256 bit units. +pub struct CostTable { + pub cost: [[u32; 2]; 64], +} + +impl CostTable { + pub fn new() -> Self { + let mut cost = [[0u32; 2]; 64]; + // Probability of the LPS for each state, estimated from rangeTabLPS relative to a mid-range interval. + for (s, c) in cost.iter_mut().enumerate() { + let avg: f64 = RANGE_LPS[s].iter().map(|&v| v as f64).sum::() / 4.0; + let p_lps = (avg / 384.0).clamp(1e-4, 0.5); + c[0] = (-(p_lps.log2()) * 256.0).round() as u32; + c[1] = (-((1.0 - p_lps).log2()) * 256.0).round() as u32; + } + CostTable { cost } + } +} + +impl Default for CostTable { + fn default() -> Self { + Self::new() + } +} + +/// Rate estimator: tracks context states like the encoder and accumulates cost in 1/256 bits. +#[derive(Clone)] +pub struct CabacEstimator<'a> { + pub ctx: Contexts, + pub cost: u64, + table: &'a CostTable, +} + +impl<'a> CabacEstimator<'a> { + pub fn new(ctx: Contexts, table: &'a CostTable) -> Self { + CabacEstimator { ctx, cost: 0, table } + } +} + +impl BinSink for CabacEstimator<'_> { + #[inline] + fn decision(&mut self, ctx: usize, bin: u32) { + let s = self.ctx.0[ctx]; + let state = (s >> 1) as usize; + let mps = (s & 1) as u32; + if bin != mps { + self.cost += self.table.cost[state][0] as u64; + let new_mps = if state == 0 { 1 - mps } else { mps }; + self.ctx.0[ctx] = (TRANS_LPS[state] << 1) | new_mps as u8; + } else { + self.cost += self.table.cost[state][1] as u64; + if state < 62 { + self.ctx.0[ctx] = (((state + 1) as u8) << 1) | mps as u8; + } + } + } + #[inline] + fn bypass(&mut self, _bin: u32) { + self.cost += 256; + } + fn terminate(&mut self, bin: u32) { + self.cost += if bin != 0 { 7 * 256 } else { 2 }; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A reference CABAC decoder (§9.3.1.2 / 9.3.3.2) used only to round-trip-test the encoder. + struct Dec<'a> { + data: &'a [u8], + pos: usize, + range: u32, + offset: u32, + ctx: Contexts, + } + impl<'a> Dec<'a> { + fn bit(&mut self) -> u32 { + let b = if self.pos / 8 < self.data.len() { (self.data[self.pos / 8] >> (7 - self.pos % 8)) & 1 } else { 0 }; + self.pos += 1; + b as u32 + } + fn new(data: &'a [u8], ctx: Contexts) -> Self { + let mut d = Dec { data, pos: 0, range: 510, offset: 0, ctx }; + for _ in 0..9 { + d.offset = (d.offset << 1) | d.bit(); + } + d + } + fn decision(&mut self, c: usize) -> u32 { + let s = self.ctx.0[c]; + let state = (s >> 1) as usize; + let mps = (s & 1) as u32; + let lps = RANGE_LPS[state][((self.range >> 6) & 3) as usize] as u32; + self.range -= lps; + let bin; + if self.offset >= self.range { + bin = 1 - mps; + self.offset -= self.range; + self.range = lps; + let new_mps = if state == 0 { 1 - mps } else { mps }; + self.ctx.0[c] = (TRANS_LPS[state] << 1) | new_mps as u8; + } else { + bin = mps; + if state < 62 { + self.ctx.0[c] = (((state + 1) as u8) << 1) | mps as u8; + } + } + while self.range < 256 { + self.range <<= 1; + self.offset = (self.offset << 1) | self.bit(); + } + bin + } + fn bypass(&mut self) -> u32 { + self.offset = (self.offset << 1) | self.bit(); + if self.offset >= self.range { + self.offset -= self.range; + 1 + } else { + 0 + } + } + fn terminate(&mut self) -> u32 { + self.range -= 2; + if self.offset >= self.range { + 1 + } else { + while self.range < 256 { + self.range <<= 1; + self.offset = (self.offset << 1) | self.bit(); + } + 0 + } + } + } + + #[test] + fn roundtrip_random_bins() { + let table = crate::tables::cabac_init_pb0(); + let ctx = Contexts::init(&table, 26); + let mut enc = CabacEncoder::new(ctx.clone(), Vec::new()); + let mut seed = 12345u32; + let mut rnd = || { + seed = seed.wrapping_mul(1664525).wrapping_add(1013904223); + seed >> 8 + }; + let mut ops = Vec::new(); + for _ in 0..20000 { + let kind = rnd() % 10; + let c = (rnd() % 400) as usize; + // skewed bins so contexts adapt + let bin = if rnd() % 100 < 80 { (c % 2) as u32 } else { 1 - (c % 2) as u32 }; + if kind < 7 { + enc.decision(c, bin); + ops.push((0, c, bin)); + } else if kind < 9 { + enc.bypass(bin); + ops.push((1, 0, bin)); + } else { + enc.terminate(0); + ops.push((2, 0, 0)); + } + } + let data = enc.finish(); + let mut dec = Dec::new(&data, ctx); + for (i, &(k, c, bin)) in ops.iter().enumerate() { + let got = match k { + 0 => dec.decision(c), + 1 => dec.bypass(), + _ => dec.terminate(), + }; + assert_eq!(got, bin, "op {i}"); + } + assert_eq!(dec.terminate(), 1); + } +} diff --git a/crates/h264enc/src/cabac_mb.rs b/crates/h264enc/src/cabac_mb.rs new file mode 100644 index 00000000..c286fbf4 --- /dev/null +++ b/crates/h264enc/src/cabac_mb.rs @@ -0,0 +1,411 @@ +//! CABAC binarisation and context selection for macroblock-layer syntax elements (§9.3.2, §9.3.3.1). + +use crate::cabac::BinSink; +use crate::mbinfo::{MbInfo, MbKind, Part}; +use crate::nal::SliceType; +use crate::syntax::{MbCode, NbCtx}; +use crate::tables::{BLK4_XY, LAST8X8, SIG8X8_FRAME}; + +pub fn write_skip_flag(s: &mut S, nb: &NbCtx, skip: bool) { + let cond = |m: Option<&MbInfo>| m.is_some_and(|m| !m.kind.is_skip()) as usize; + let base = if nb.slice_type == SliceType::B { 24 } else { 11 }; + s.decision(base + cond(nb.left) + cond(nb.top), skip as u32); +} + +fn eg_bypass(s: &mut S, mut v: u32, mut k: u32) { + loop { + if v >= (1 << k) { + s.bypass(1); + v -= 1 << k; + k += 1; + } else { + s.bypass(0); + while k > 0 { + k -= 1; + s.bypass((v >> k) & 1); + } + break; + } + } +} + +/// I16x16 / I_NxN mb_type bins. `off` = ctxIdxOffset of the (suffix) table; `i_slice` selects I-slice contexts. +fn write_intra_type(s: &mut S, nb: &NbCtx, c: &MbCode) { + let i_slice = nb.slice_type == SliceType::I; + let (off, b0ctx) = match nb.slice_type { + SliceType::I => { + let cond = |m: Option<&MbInfo>| m.is_some_and(|m| !m.kind.is_inxn()) as usize; + (3, 3 + cond(nb.left) + cond(nb.top)) + } + SliceType::P => { + s.decision(14, 1); // prefix: intra + (17, 17) + } + SliceType::B => { + // prefix 111101 + write_b_bins(s, nb, &[1, 1, 1, 1, 0, 1]); + (32, 32) + } + }; + if c.kind != MbKind::I16x16 { + s.decision(b0ctx, 0); + return; + } + s.decision(b0ctx, 1); + s.terminate(0); + let luma = (c.cbp & 15 != 0) as u32; + let chroma = (c.cbp >> 4) as u32; + let pred = c.i16_mode as u32; + if i_slice { + s.decision(off + 3, luma); + s.decision(off + 4, (chroma != 0) as u32); + if chroma != 0 { + s.decision(off + 5, (chroma == 2) as u32); + s.decision(off + 6, pred >> 1); + s.decision(off + 7, pred & 1); + } else { + s.decision(off + 6, pred >> 1); + s.decision(off + 7, pred & 1); + } + } else { + s.decision(off + 1, luma); + s.decision(off + 2, (chroma != 0) as u32); + if chroma != 0 { + s.decision(off + 2, (chroma == 2) as u32); + } + s.decision(off + 3, pred >> 1); + s.decision(off + 3, pred & 1); + } +} + +fn write_b_bins(s: &mut S, nb: &NbCtx, bins: &[u32]) { + for (i, &b) in bins.iter().enumerate() { + let ctx = match i { + 0 => { + let cond = |m: Option<&MbInfo>| m.is_some_and(|m| !matches!(m.kind, MbKind::BSkip | MbKind::BDirect)) as usize; + 27 + cond(nb.left) + cond(nb.top) + } + 1 => 27 + 3, + 2 => 27 + if bins[1] != 0 { 4 } else { 5 }, + _ => 27 + 5, + }; + s.decision(ctx, b); + } +} + +const B_BINS: [&[u32]; 23] = [ + &[0], + &[1, 0, 0], + &[1, 0, 1], + &[1, 1, 0, 0, 0, 0], + &[1, 1, 0, 0, 0, 1], + &[1, 1, 0, 0, 1, 0], + &[1, 1, 0, 0, 1, 1], + &[1, 1, 0, 1, 0, 0], + &[1, 1, 0, 1, 0, 1], + &[1, 1, 0, 1, 1, 0], + &[1, 1, 0, 1, 1, 1], + &[1, 1, 1, 1, 1, 0], + &[1, 1, 1, 0, 0, 0, 0], + &[1, 1, 1, 0, 0, 0, 1], + &[1, 1, 1, 0, 0, 1, 0], + &[1, 1, 1, 0, 0, 1, 1], + &[1, 1, 1, 0, 1, 0, 0], + &[1, 1, 1, 0, 1, 0, 1], + &[1, 1, 1, 0, 1, 1, 0], + &[1, 1, 1, 0, 1, 1, 1], + &[1, 1, 1, 1, 0, 0, 0], + &[1, 1, 1, 1, 0, 0, 1], + &[1, 1, 1, 1, 1, 1], +]; + +fn write_mvd(s: &mut S, comp: usize, v: i32, sum: u32) { + let base = if comp == 0 { 40 } else { 47 }; + let a = v.unsigned_abs(); + let inc = if sum < 3 { + 0 + } else if sum > 32 { + 2 + } else { + 1 + }; + let u = a.min(9); + s.decision(base + inc, (u > 0) as u32); + if u > 0 { + for k in 1..9u32 { + let ctx = base + [3, 4, 5, 6, 6, 6, 6, 6][(k - 1) as usize]; + if k < u { + s.decision(ctx, 1); + } else { + s.decision(ctx, 0); + break; + } + } + if a >= 9 { + eg_bypass(s, a - 9, 3); + } + s.bypass((v < 0) as u32); + } +} + +/// Sum of |mvd| of neighbours A and B for the 4x4 block at raster `r` (current MB info already holds earlier partitions). +fn mvd_sum(nb: &NbCtx, r: usize, list: usize, comp: usize) -> u32 { + let (x, y) = (r & 3, r >> 2); + let a = if x > 0 { nb.cur.mvd[list][r - 1][comp] } else { nb.left.map_or(0, |m| m.mvd[list][y * 4 + 3][comp]) }; + let b = if y > 0 { nb.cur.mvd[list][r - 4][comp] } else { nb.top.map_or(0, |m| m.mvd[list][12 + x][comp]) }; + a as u32 + b as u32 +} + +fn write_cbp(s: &mut S, nb: &NbCtx, cbp: u8) { + let luma_bit = |m: &MbInfo, b8: usize| (m.cbp >> b8) & 1; + for b8 in 0..4usize { + // (A, B) neighbour 8x8 blocks: (Some(mb), idx) where None = current MB + let cond = |m: Option<&MbInfo>, idx: usize, in_cur: bool| -> usize { + if in_cur { + ((cbp >> idx) & 1 == 0) as usize + } else { + match m { + None => 0, + Some(m) => (luma_bit(m, idx) == 0) as usize, + } + } + }; + let (ca, cb) = match b8 { + 0 => (cond(nb.left, 1, false), cond(nb.top, 2, false)), + 1 => (cond(None, 0, true), cond(nb.top, 3, false)), + 2 => (cond(nb.left, 3, false), cond(None, 0, true)), + _ => (cond(None, 2, true), cond(None, 1, true)), + }; + s.decision(73 + ca + 2 * cb, ((cbp >> b8) & 1) as u32); + } + let chroma = cbp >> 4; + let cc = |m: Option<&MbInfo>, two: bool| -> usize { + m.is_some_and(|m| { + let c = m.cbp >> 4; + !m.kind.is_skip() && if two { c == 2 } else { c != 0 } + }) as usize + }; + s.decision(77 + cc(nb.left, false) + 2 * cc(nb.top, false), (chroma != 0) as u32); + if chroma != 0 { + s.decision(77 + 4 + cc(nb.left, true) + 2 * cc(nb.top, true), (chroma == 2) as u32); + } +} + +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum Cat { + LumaDc = 0, + LumaAc = 1, + Luma4x4 = 2, + ChromaDc = 3, + ChromaAc = 4, + Luma8x8 = 5, +} + +/// Encode one residual block; `coeffs` are the block's coefficients in scan order (length = maxNumCoeff). +pub fn write_residual(s: &mut S, cat: Cat, coeffs: &[i16], cbf_inc: u32) { + const CBF_OFF: [usize; 5] = [0, 4, 8, 12, 16]; + const SIG_OFF: [usize; 5] = [0, 15, 29, 44, 47]; + const ABS_OFF: [usize; 5] = [0, 10, 20, 30, 39]; + let n = coeffs.len(); + let last = coeffs.iter().rposition(|&c| c != 0); + let ci = cat as usize; + if cat != Cat::Luma8x8 { + s.decision(85 + CBF_OFF[ci] + cbf_inc as usize, last.is_some() as u32); + } + let Some(last) = last else { return }; + let (sig_base, last_base, abs_base) = if cat == Cat::Luma8x8 { (402, 417, 426) } else { (105 + SIG_OFF[ci], 166 + SIG_OFF[ci], 227 + ABS_OFF[ci]) }; + for i in 0..n - 1 { + let (si, li) = match cat { + Cat::ChromaDc => (i.min(2), i.min(2)), + Cat::Luma8x8 => (SIG8X8_FRAME[i] as usize, LAST8X8[i] as usize), + _ => (i, i), + }; + let sig = coeffs[i] != 0; + s.decision(sig_base + si, sig as u32); + if sig { + let is_last = i == last; + s.decision(last_base + li, is_last as u32); + if is_last { + break; + } + } + } + let mut gt1 = 0usize; + let mut eq1 = 0usize; + let max_gt1 = if cat == Cat::ChromaDc { 3 } else { 4 }; + for i in (0..=last).rev() { + let c = coeffs[i]; + if c == 0 { + continue; + } + let a = c.unsigned_abs() as u32; + let v = a - 1; + let inc0 = if gt1 != 0 { 0 } else { (1 + eq1).min(4) }; + s.decision(abs_base + inc0, (v > 0) as u32); + if v > 0 { + let ctx = abs_base + 5 + gt1.min(max_gt1); + let u = v.min(14); + for _ in 1..u { + s.decision(ctx, 1); + } + if u < 14 { + s.decision(ctx, 0); + } else { + eg_bypass(s, v - 14, 0); + } + } + s.bypass((c < 0) as u32); + if a == 1 { + eq1 += 1; + } else { + gt1 += 1; + } + } +} + +fn cbf_luma(nb: &NbCtx, raster: usize) -> u32 { + let f = |v: Option| v.map_or(nb.unavail_cbf(), |n| (n != 0) as u32); + f(nb.luma_nnz(raster, true)) + 2 * f(nb.luma_nnz(raster, false)) +} + +fn cbf_chroma_ac(nb: &NbCtx, c: usize, b: usize) -> u32 { + let f = |v: Option| v.map_or(nb.unavail_cbf(), |n| (n != 0) as u32); + f(nb.chroma_nnz(c, b, true)) + 2 * f(nb.chroma_nnz(c, b, false)) +} + +fn cbf_dc(nb: &NbCtx, bit: u8) -> u32 { + let f = |m: Option<&MbInfo>| m.map_or(nb.unavail_cbf(), |m| ((m.dc_cbf >> bit) & 1) as u32); + f(nb.left) + 2 * f(nb.top) +} + +/// Write the whole macroblock layer (after mb_skip_flag = 0 in P/B slices). +pub fn write_mb(s: &mut S, nb: &NbCtx, c: &MbCode) { + let cur = nb.cur; + // ---- mb_type + match c.kind { + MbKind::I4x4 | MbKind::I8x8 | MbKind::I16x16 => write_intra_type(s, nb, c), + MbKind::PInter => { + let bins: [u32; 3] = match c.part { + Part::P16x16 => [0, 0, 0], + Part::P16x8 => [0, 1, 1], + Part::P8x16 => [0, 1, 0], + Part::P8x8 => [0, 0, 1], + }; + s.decision(14, bins[0]); + s.decision(15, bins[1]); + s.decision(if bins[1] != 1 { 16 } else { 17 }, bins[2]); + } + MbKind::BDirect | MbKind::BInter => write_b_bins(s, nb, B_BINS[c.b_type_number() as usize]), + _ => unreachable!("skip/none macroblocks are not written with write_mb"), + } + let is_inxn = c.kind.is_inxn(); + if c.kind == MbKind::PInter && c.part == Part::P8x8 { + for _ in 0..4 { + s.decision(21, 1); // P_L0_8x8 + } + } + // ---- mb_pred + if is_inxn { + if nb.t8x8_mode { + let cond = |m: Option<&MbInfo>| m.is_some_and(|m| m.t8x8) as usize; + s.decision(399 + cond(nb.left) + cond(nb.top), c.t8x8 as u32); + } + let n = if c.t8x8 { 4 } else { 16 }; + for i in 0..n { + let p = c.ipred[i]; + if p < 0 { + s.decision(68, 1); + } else { + s.decision(68, 0); + s.decision(69, (p & 1) as u32); + s.decision(69, ((p >> 1) & 1) as u32); + s.decision(69, ((p >> 2) & 1) as u32); + } + } + } + if c.kind.is_intra() { + let cond = |m: Option<&MbInfo>| m.is_some_and(|m| m.kind.is_intra() && m.chroma_mode != 0) as usize; + let m = c.chroma_mode as u32; + s.decision(64 + cond(nb.left) + cond(nb.top), (m > 0) as u32); + if m > 0 { + s.decision(64 + 3, (m > 1) as u32); + if m > 1 { + s.decision(64 + 3, (m > 2) as u32); + } + } + } + if matches!(c.kind, MbKind::PInter | MbKind::BInter) { + for list in 0..2 { + for p in 0..c.part.count() { + if !c.uses_list(p, list) { + continue; + } + let (px, py, _, _) = c.part.rect(p); + let r = (py / 4) * 4 + px / 4; + let mvd = c.mvd[list][p]; + write_mvd(s, 0, mvd.x as i32, mvd_sum(nb, r, list, 0)); + write_mvd(s, 1, mvd.y as i32, mvd_sum(nb, r, list, 1)); + } + } + } + // ---- cbp, transform size, qp delta + if c.kind != MbKind::I16x16 { + write_cbp(s, nb, c.cbp); + if c.cbp & 15 != 0 && nb.t8x8_mode && !is_inxn { + let cond = |m: Option<&MbInfo>| m.is_some_and(|m| m.t8x8) as usize; + s.decision(399 + cond(nb.left) + cond(nb.top), c.t8x8 as u32); + } + } + if c.cbp == 0 && c.kind != MbKind::I16x16 { + return; + } + { + let v = c.qp_delta; + let m = if v > 0 { 2 * v as u32 - 1 } else { (-2 * v) as u32 }; + s.decision(60 + nb.prev_dqp_nz as usize, (m > 0) as u32); + if m > 0 { + for k in 1..=m { + s.decision(if k == 1 { 62 } else { 63 }, (k < m) as u32); + } + } + } + // ---- residual + if c.kind == MbKind::I16x16 { + write_residual(s, Cat::LumaDc, &c.luma_dc, cbf_dc(nb, 0)); + if c.cbp & 15 != 0 { + for blk in 0..16 { + let (x, y) = BLK4_XY[blk]; + write_residual(s, Cat::LumaAc, &c.luma4[blk][1..16], cbf_luma(nb, y * 4 + x)); + } + } + } else { + for b8 in 0..4 { + if (c.cbp >> b8) & 1 == 0 { + continue; + } + if c.t8x8 { + write_residual(s, Cat::Luma8x8, &c.luma8[b8], 0); + } else { + for i in 0..4 { + let blk = b8 * 4 + i; + let (x, y) = BLK4_XY[blk]; + write_residual(s, Cat::Luma4x4, &c.luma4[blk], cbf_luma(nb, y * 4 + x)); + } + } + } + } + let chroma = c.cbp >> 4; + if chroma != 0 { + for comp in 0..2 { + write_residual(s, Cat::ChromaDc, &c.cdc[comp], cbf_dc(nb, 1 + comp as u8)); + } + } + if chroma == 2 { + for comp in 0..2 { + for b in 0..4 { + write_residual(s, Cat::ChromaAc, &c.cac[comp * 4 + b][1..16], cbf_chroma_ac(nb, comp, b)); + } + } + } + let _ = cur; +} diff --git a/crates/h264enc/src/cavlc.rs b/crates/h264enc/src/cavlc.rs new file mode 100644 index 00000000..cdfddc22 --- /dev/null +++ b/crates/h264enc/src/cavlc.rs @@ -0,0 +1,262 @@ +//! CAVLC residual coding (§9.2) and macroblock-layer syntax for Baseline streams. + +use filmcraft_bitstream::BitWriter; + +use crate::mbinfo::{MbKind, Part}; +use crate::nal::SliceType; +use crate::syntax::{MbCode, NbCtx}; +use crate::tables::{ + BLK4_XY, CBP_INTER_TO_CODE, CBP_INTRA_TO_CODE, COEFF_TOKEN_CDC_CODE, COEFF_TOKEN_CDC_LEN, COEFF_TOKEN_CODE, COEFF_TOKEN_LEN, RUN_BEFORE, TOTAL_ZEROS, + TOTAL_ZEROS_CDC, +}; + +/// Largest level magnitude representable with level_prefix <= 15 in every suffixLength state. +pub const MAX_LEVEL: i32 = 2047; + +/// Write one CAVLC residual block. `nc` = -1 for chroma DC. Returns TotalCoeff. +pub fn write_residual(w: &mut BitWriter, coeffs: &[i16], nc: i32) -> u8 { + let max = coeffs.len(); + // levels from highest frequency to lowest + let mut levels = [0i32; 16]; + let mut pos = [0usize; 16]; + let mut total = 0usize; + for i in (0..max).rev() { + if coeffs[i] != 0 { + levels[total] = coeffs[i] as i32; + pos[total] = i; + total += 1; + } + } + let mut t1 = 0usize; + while t1 < total && t1 < 3 && levels[t1].abs() == 1 { + t1 += 1; + } + // coeff_token + if nc == -1 { + w.write_bits(COEFF_TOKEN_CDC_CODE[total][t1] as u32, COEFF_TOKEN_CDC_LEN[total][t1] as u32); + } else if nc >= 8 { + if total == 0 { + w.write_bits(3, 6); + } else { + w.write_bits((((total - 1) << 2) | t1) as u32, 6); + } + } else { + let t = if nc < 2 { + 0 + } else if nc < 4 { + 1 + } else { + 2 + }; + w.write_bits(COEFF_TOKEN_CODE[t][total][t1] as u32, COEFF_TOKEN_LEN[t][total][t1] as u32); + } + if total == 0 { + return 0; + } + let mut suffix_len: u32 = if total > 10 && t1 < 3 { 1 } else { 0 }; + for i in 0..total { + let l = levels[i]; + if i < t1 { + w.write_bit(l < 0); + continue; + } + let mut code = if l > 0 { 2 * l - 2 } else { -2 * l - 1 }; + if i == t1 && t1 < 3 { + code -= 2; + } + let code = code as u32; + if suffix_len == 0 { + if code < 14 { + w.write_bits(1, code + 1); + } else if code < 30 { + w.write_bits(1, 15); + w.write_bits(code - 14, 4); + } else { + write_escape(w, code - 30); + } + } else if code < (15 << suffix_len) { + w.write_bits(1, (code >> suffix_len) + 1); + w.write_bits(code & ((1 << suffix_len) - 1), suffix_len); + } else { + write_escape(w, code - (15 << suffix_len)); + } + if suffix_len == 0 { + suffix_len = 1; + } + if l.unsigned_abs() > (3 << (suffix_len - 1)) && suffix_len < 6 { + suffix_len += 1; + } + } + let total_zeros = pos[0] + 1 - total; + if total < max { + let (len, code) = if nc == -1 { TOTAL_ZEROS_CDC[total - 1][total_zeros] } else { TOTAL_ZEROS[total - 1][total_zeros] }; + w.write_bits(code as u32, len as u32); + } + let mut zeros_left = total_zeros; + for i in 0..total - 1 { + if zeros_left == 0 { + break; + } + let run = pos[i] - pos[i + 1] - 1; + let (len, code) = RUN_BEFORE[zeros_left.min(7) - 1][run]; + w.write_bits(code as u32, len as u32); + zeros_left -= run; + } + total as u8 +} + +/// level_prefix >= 15 escape; `rem` = levelCode minus the prefix-15 base. +fn write_escape(w: &mut BitWriter, rem: u32) { + if rem < 4096 { + w.write_bits(1, 16); + w.write_bits(rem, 12); + } else { + // level_prefix p >= 16 (High profiles): suffix of p-3 bits, value rem + 4096 - 2^(p-3) + let mut p = 16; + while rem + 4096 >= (1 << (p - 3)) * 2 { + p += 1; + } + w.write_bits(0, p); + w.write_bit(true); + w.write_bits(rem + 4096 - (1 << (p - 3)), p - 3); + } +} + +fn nc_of(a: Option, b: Option) -> i32 { + match (a, b) { + (Some(a), Some(b)) => (a as i32 + b as i32 + 1) >> 1, + (Some(a), None) => a as i32, + (None, Some(b)) => b as i32, + _ => 0, + } +} + +/// Write the macroblock layer in CAVLC. For P/B slices the caller writes mb_skip_run first. +pub fn write_mb(w: &mut BitWriter, nb: &NbCtx, c: &MbCode) { + let intra_base = match nb.slice_type { + SliceType::I => 0, + SliceType::P => 5, + SliceType::B => 23, + }; + match c.kind { + MbKind::I4x4 | MbKind::I8x8 => w.write_ue(intra_base), + MbKind::I16x16 => { + let t = 1 + c.i16_mode as u32 + 4 * (c.cbp >> 4) as u32 + if c.cbp & 15 != 0 { 12 } else { 0 }; + w.write_ue(intra_base + t); + } + MbKind::PInter => w.write_ue(match c.part { + Part::P16x16 => 0, + Part::P16x8 => 1, + Part::P8x16 => 2, + Part::P8x8 => 3, + }), + MbKind::BDirect | MbKind::BInter => w.write_ue(c.b_type_number()), + _ => unreachable!(), + } + if c.kind == MbKind::PInter && c.part == Part::P8x8 { + for _ in 0..4 { + w.write_ue(0); + } + } + if c.kind.is_inxn() { + if nb.t8x8_mode { + w.write_bit(c.t8x8); + } + let n = if c.t8x8 { 4 } else { 16 }; + for i in 0..n { + let p = c.ipred[i]; + if p < 0 { + w.write_bit(true); + } else { + w.write_bit(false); + w.write_bits(p as u32, 3); + } + } + } + if c.kind.is_intra() { + w.write_ue(c.chroma_mode as u32); + } + if matches!(c.kind, MbKind::PInter | MbKind::BInter) { + for list in 0..2 { + for p in 0..c.part.count() { + if c.uses_list(p, list) { + w.write_se(c.mvd[list][p].x as i32); + w.write_se(c.mvd[list][p].y as i32); + } + } + } + } + if c.kind != MbKind::I16x16 { + let code = if c.kind.is_intra() { CBP_INTRA_TO_CODE[c.cbp as usize] } else { CBP_INTER_TO_CODE[c.cbp as usize] }; + let cbp_idx = (c.cbp & 15) | ((c.cbp >> 4) << 4); + debug_assert_eq!(cbp_idx, c.cbp); + w.write_ue(code as u32); + if c.cbp & 15 != 0 && nb.t8x8_mode && !c.kind.is_inxn() { + w.write_bit(c.t8x8); + } + } + if c.cbp == 0 && c.kind != MbKind::I16x16 { + return; + } + w.write_se(c.qp_delta); + let luma_nc = |r: usize| nc_of(nb.luma_nnz(r, true), nb.luma_nnz(r, false)); + if c.kind == MbKind::I16x16 { + write_residual(w, &c.luma_dc, luma_nc(0)); + if c.cbp & 15 != 0 { + for blk in 0..16 { + let (x, y) = BLK4_XY[blk]; + write_residual(w, &c.luma4[blk][1..16], luma_nc(y * 4 + x)); + } + } + } else { + for b8 in 0..4 { + if (c.cbp >> b8) & 1 == 0 { + continue; + } + for i in 0..4 { + let blk = b8 * 4 + i; + let (x, y) = BLK4_XY[blk]; + if c.t8x8 { + let mut sub = [0i16; 16]; + for k in 0..16 { + sub[k] = c.luma8[b8][4 * k + i]; + } + write_residual(w, &sub, luma_nc(y * 4 + x)); + } else { + write_residual(w, &c.luma4[blk], luma_nc(y * 4 + x)); + } + } + } + } + let chroma = c.cbp >> 4; + if chroma != 0 { + for comp in 0..2 { + write_residual(w, &c.cdc[comp], -1); + } + } + if chroma == 2 { + for comp in 0..2 { + for b in 0..4 { + let ncv = nc_of(nb.chroma_nnz(comp, b, true), nb.chroma_nnz(comp, b, false)); + write_residual(w, &c.cac[comp * 4 + b][1..16], ncv); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn simple_block_matches_spec_example() { + // scan: 0,3,-1,0,0,-1,1,0,1 -> TotalCoeff 5, T1s 3. + let coeffs: [i16; 16] = [0, 3, -1, 0, 0, -1, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0]; + let mut w = BitWriter::new(); + let t = write_residual(&mut w, &coeffs, 0); + assert_eq!(t, 5); + // coeff_token 0000100 | signs 001 | level -1: 01 | level 3: 001 0 | total_zeros(4) 110 | runs 10 11 01 1 + let bits: String = w.finish().iter().map(|b| format!("{b:08b}")).collect(); + assert!(bits.starts_with("00001000010100101101011011"), "{bits}"); + } +} diff --git a/crates/h264enc/src/deblock.rs b/crates/h264enc/src/deblock.rs new file mode 100644 index 00000000..97732ccc --- /dev/null +++ b/crates/h264enc/src/deblock.rs @@ -0,0 +1,251 @@ +//! In-loop deblocking filter (§8.7) for frame-coded 4:2:0 pictures without MBAFF. +//! +//! Operates on a band of macroblock rows that begins at a slice boundary; the top edge of the band is +//! never filtered (slices are MB-row aligned and either there is one slice, or +//! `disable_deblocking_filter_idc = 2` is used). + +use crate::mbinfo::{MbInfo, Mv}; +use crate::tables::{ALPHA, BETA, TC0, chroma_qp}; + +pub struct DeblockParams { + pub alpha_offset: i32, + pub beta_offset: i32, + pub chroma_qp_offset: i32, + /// Picture identity for (list, ref_idx 0). + pub ref_ids: [i32; 2], +} + +/// A mutable band of rows of a padded plane: sample (x, y) with y relative to the band top. +pub struct BandPlane<'a> { + pub data: &'a mut [u8], + pub stride: usize, + pub pad: usize, +} + +impl BandPlane<'_> { + #[inline] + fn idx(&self, x: usize, y: usize) -> usize { + y * self.stride + self.pad + x + } +} + +#[inline] +fn nz8(m: &MbInfo, blk: usize) -> bool { + if m.t8x8 { + let (x, y) = ((blk & 3) >> 1, (blk >> 2) >> 1); + let b = y * 8 + x * 2; + m.nnz[b] | m.nnz[b + 1] | m.nnz[b + 4] | m.nnz[b + 5] != 0 + } else { + m.nnz[blk] != 0 + } +} + +#[inline] +fn mvdiff(a: Mv, b: Mv) -> bool { + (a.x as i32 - b.x as i32).abs() >= 4 || (a.y as i32 - b.y as i32).abs() >= 4 +} + +fn bs_inter(p: &MbInfo, ip: usize, q: &MbInfo, iq: usize, ids: &[i32; 2]) -> u8 { + if nz8(p, ip) || nz8(q, iq) { + return 2; + } + let rp = [p.ref_at(0, ip), p.ref_at(1, ip)]; + let rq = [q.ref_at(0, iq), q.ref_at(1, iq)]; + // Only ref_idx 0 is ever used per list, so the list index identifies the picture. + let pic = |r: i8, l: usize| if r >= 0 { ids[l] } else { -1 }; + let pp = [pic(rp[0], 0), pic(rp[1], 1)]; + let pq = [pic(rq[0], 0), pic(rq[1], 1)]; + let np = (pp[0] >= 0) as u8 + (pp[1] >= 0) as u8; + let nq = (pq[0] >= 0) as u8 + (pq[1] >= 0) as u8; + if np != nq { + return 1; + } + if np == 1 { + let (a, ma) = if pp[0] >= 0 { (pp[0], p.mv[0][ip]) } else { (pp[1], p.mv[1][ip]) }; + let (b, mb) = if pq[0] >= 0 { (pq[0], q.mv[0][iq]) } else { (pq[1], q.mv[1][iq]) }; + if a != b { + return 1; + } + return mvdiff(ma, mb) as u8; + } + if np == 0 { + return 0; + } + // two motion vectors each + let (mp0, mp1, mq0, mq1) = (p.mv[0][ip], p.mv[1][ip], q.mv[0][iq], q.mv[1][iq]); + let same_set = (pp[0] == pq[0] && pp[1] == pq[1]) || (pp[0] == pq[1] && pp[1] == pq[0]); + if !same_set { + return 1; + } + if pp[0] != pp[1] { + // two different pictures: compare mvs referring to the same picture + if pp[0] == pq[0] { (mvdiff(mp0, mq0) || mvdiff(mp1, mq1)) as u8 } else { (mvdiff(mp0, mq1) || mvdiff(mp1, mq0)) as u8 } + } else { + ((mvdiff(mp0, mq0) || mvdiff(mp1, mq1)) && (mvdiff(mp0, mq1) || mvdiff(mp1, mq0))) as u8 + } +} + +#[inline] +fn bs_of(p: &MbInfo, ip: usize, q: &MbInfo, iq: usize, mb_edge: bool, ids: &[i32; 2]) -> u8 { + if p.kind.is_intra() || q.kind.is_intra() { + return if mb_edge { 4 } else { 3 }; + } + bs_inter(p, ip, q, iq, ids) +} + +#[inline] +fn filter_luma(d: &mut [u8], q0i: usize, step: usize, bs: u8, alpha: i32, beta: i32, tc0: i32) { + let p0 = d[q0i - step] as i32; + let q0 = d[q0i] as i32; + let p1 = d[q0i - 2 * step] as i32; + let q1 = d[q0i + step] as i32; + if (p0 - q0).abs() >= alpha || (p1 - p0).abs() >= beta || (q1 - q0).abs() >= beta { + return; + } + let p2 = d[q0i - 3 * step] as i32; + let q2 = d[q0i + 2 * step] as i32; + let ap = (p2 - p0).abs(); + let aq = (q2 - q0).abs(); + if bs < 4 { + let tc = tc0 + (ap < beta) as i32 + (aq < beta) as i32; + let delta = ((((q0 - p0) << 2) + (p1 - q1) + 4) >> 3).clamp(-tc, tc); + d[q0i - step] = (p0 + delta).clamp(0, 255) as u8; + d[q0i] = (q0 - delta).clamp(0, 255) as u8; + if ap < beta { + d[q0i - 2 * step] = (p1 + ((p2 + ((p0 + q0 + 1) >> 1) - (p1 << 1)) >> 1).clamp(-tc0, tc0)) as u8; + } + if aq < beta { + d[q0i + step] = (q1 + ((q2 + ((p0 + q0 + 1) >> 1) - (q1 << 1)) >> 1).clamp(-tc0, tc0)) as u8; + } + } else { + let p3 = d[q0i - 4 * step] as i32; + let q3 = d[q0i + 3 * step] as i32; + let strong = (p0 - q0).abs() < ((alpha >> 2) + 2); + if ap < beta && strong { + d[q0i - step] = ((p2 + 2 * p1 + 2 * p0 + 2 * q0 + q1 + 4) >> 3) as u8; + d[q0i - 2 * step] = ((p2 + p1 + p0 + q0 + 2) >> 2) as u8; + d[q0i - 3 * step] = ((2 * p3 + 3 * p2 + p1 + p0 + q0 + 4) >> 3) as u8; + } else { + d[q0i - step] = ((2 * p1 + p0 + q1 + 2) >> 2) as u8; + } + if aq < beta && strong { + d[q0i] = ((p1 + 2 * p0 + 2 * q0 + 2 * q1 + q2 + 4) >> 3) as u8; + d[q0i + step] = ((p0 + q0 + q1 + q2 + 2) >> 2) as u8; + d[q0i + 2 * step] = ((2 * q3 + 3 * q2 + q1 + q0 + p0 + 4) >> 3) as u8; + } else { + d[q0i] = ((2 * q1 + q0 + p1 + 2) >> 2) as u8; + } + } +} + +#[inline] +fn filter_chroma(d: &mut [u8], q0i: usize, step: usize, bs: u8, alpha: i32, beta: i32, tc0: i32) { + let p0 = d[q0i - step] as i32; + let q0 = d[q0i] as i32; + let p1 = d[q0i - 2 * step] as i32; + let q1 = d[q0i + step] as i32; + if (p0 - q0).abs() >= alpha || (p1 - p0).abs() >= beta || (q1 - q0).abs() >= beta { + return; + } + if bs < 4 { + let tc = tc0 + 1; + let delta = ((((q0 - p0) << 2) + (p1 - q1) + 4) >> 3).clamp(-tc, tc); + d[q0i - step] = (p0 + delta).clamp(0, 255) as u8; + d[q0i] = (q0 - delta).clamp(0, 255) as u8; + } else { + d[q0i - step] = ((2 * p1 + p0 + q1 + 2) >> 2) as u8; + d[q0i] = ((2 * q1 + q0 + p1 + 2) >> 2) as u8; + } +} + +struct EdgeQ { + alpha: i32, + beta: i32, + idx_a: usize, +} + +fn edge_q(qp_p: i32, qp_q: i32, dp: &DeblockParams) -> EdgeQ { + let qav = (qp_p + qp_q + 1) >> 1; + let ia = (qav + dp.alpha_offset).clamp(0, 51) as usize; + let ib = (qav + dp.beta_offset).clamp(0, 51) as usize; + EdgeQ { alpha: ALPHA[ia] as i32, beta: BETA[ib] as i32, idx_a: ia } +} + +/// Deblock `rows` macroblock rows of a band. `mbs` holds the band's MbInfo (row-major, `mbw` wide). +pub fn deblock_band<'a>(y: &mut BandPlane<'a>, u: &mut BandPlane<'a>, v: &mut BandPlane<'a>, mbs: &[MbInfo], mbw: usize, rows: usize, dp: &DeblockParams) { + let ids = &dp.ref_ids; + for my in 0..rows { + for mx in 0..mbw { + let q = &mbs[my * mbw + mx]; + let cqp = |m: &MbInfo| chroma_qp(m.qp as i32 + dp.chroma_qp_offset) as i32; + // ---- vertical edges (dir 0) then horizontal (dir 1) + for dir in 0..2 { + let mut bs_all = [[0u8; 4]; 4]; + for e in 0..4 { + if e == 0 && ((dir == 0 && mx == 0) || (dir == 1 && my == 0)) { + continue; + } + if e & 1 == 1 && q.t8x8 { + continue; + } + let p_mb = if e == 0 { if dir == 0 { &mbs[my * mbw + mx - 1] } else { &mbs[(my - 1) * mbw + mx] } } else { q }; + let mut bs = [0u8; 4]; + for s in 0..4 { + let (iq, ip) = if dir == 0 { + let iq = s * 4 + e; + (iq, if e == 0 { s * 4 + 3 } else { iq - 1 }) + } else { + let iq = e * 4 + s; + (iq, if e == 0 { 12 + s } else { iq - 4 }) + }; + bs[s] = bs_of(p_mb, ip, q, iq, e == 0, ids); + } + bs_all[e] = bs; + if bs == [0; 4] { + continue; + } + let eq = edge_q(p_mb.qp as i32, q.qp as i32, dp); + if eq.alpha == 0 || eq.beta == 0 { + continue; + } + for k in 0..16 { + let b = bs[k >> 2]; + if b == 0 { + continue; + } + let tc0 = if b < 4 { TC0[eq.idx_a][b as usize - 1] as i32 } else { 0 }; + let (qi, step) = if dir == 0 { (y.idx(mx * 16 + e * 4, my * 16 + k), 1) } else { (y.idx(mx * 16 + k, my * 16 + e * 4), y.stride) }; + filter_luma(y.data, qi, step, b, eq.alpha, eq.beta, tc0); + } + } + // chroma edges 0 and 2 (chroma sample offsets 0 and 4) + for e in [0usize, 2] { + if e == 0 && ((dir == 0 && mx == 0) || (dir == 1 && my == 0)) { + continue; + } + let bs = bs_all[e]; + if bs == [0; 4] { + continue; + } + let p_mb = if e == 0 { if dir == 0 { &mbs[my * mbw + mx - 1] } else { &mbs[(my - 1) * mbw + mx] } } else { q }; + let eq = edge_q(cqp(p_mb), cqp(q), dp); + if eq.alpha == 0 || eq.beta == 0 { + continue; + } + for plane in [&mut *u, &mut *v] { + for k in 0..8 { + let b = bs[k >> 1]; + if b == 0 { + continue; + } + let tc0 = if b < 4 { TC0[eq.idx_a][b as usize - 1] as i32 } else { 0 }; + let (qi, step) = + if dir == 0 { (plane.idx(mx * 8 + e * 2, my * 8 + k), 1) } else { (plane.idx(mx * 8 + k, my * 8 + e * 2), plane.stride) }; + filter_chroma(plane.data, qi, step, b, eq.alpha, eq.beta, tc0); + } + } + } + } + } + } +} diff --git a/crates/h264enc/src/dsp.rs b/crates/h264enc/src/dsp.rs new file mode 100644 index 00000000..2634701e --- /dev/null +++ b/crates/h264enc/src/dsp.rs @@ -0,0 +1,172 @@ +//! Distortion metrics: SAD, SATD (4x4 Hadamard), SSD. + +#[inline] +pub fn sad(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 { + let mut s = 0u32; + for r in 0..h { + let ra = &a[r * sa..r * sa + w]; + let rb = &b[r * sb..r * sb + w]; + s += ra.iter().zip(rb).map(|(&x, &y)| x.abs_diff(y) as u32).sum::(); + } + s +} + +/// SAD of a 16-wide block (fast path). +#[inline] +pub fn sad16(a: &[u8], sa: usize, b: &[u8], sb: usize, h: usize) -> u32 { + let mut s = 0u32; + for r in 0..h { + let ra: &[u8; 16] = a[r * sa..r * sa + 16].try_into().unwrap(); + let rb: &[u8; 16] = b[r * sb..r * sb + 16].try_into().unwrap(); + let mut acc = 0u16; + for i in 0..16 { + acc += ra[i].abs_diff(rb[i]) as u16; + } + s += acc as u32; + } + s +} + +#[cfg(test)] +fn hadamard4x4_sum(d: &[i32; 16]) -> u32 { + let mut t = [0i32; 16]; + for i in 0..4 { + let a0 = d[i * 4] + d[i * 4 + 1]; + let a1 = d[i * 4] - d[i * 4 + 1]; + let a2 = d[i * 4 + 2] + d[i * 4 + 3]; + let a3 = d[i * 4 + 2] - d[i * 4 + 3]; + t[i * 4] = a0 + a2; + t[i * 4 + 1] = a1 + a3; + t[i * 4 + 2] = a0 - a2; + t[i * 4 + 3] = a1 - a3; + } + let mut s = 0u32; + for i in 0..4 { + let a0 = t[i] + t[4 + i]; + let a1 = t[i] - t[4 + i]; + let a2 = t[8 + i] + t[12 + i]; + let a3 = t[8 + i] - t[12 + i]; + s += (a0 + a2).unsigned_abs() + (a1 + a3).unsigned_abs() + (a0 - a2).unsigned_abs() + (a1 - a3).unsigned_abs(); + } + s +} + +/// SATD of a strip of `N` columns (N/4 side-by-side 4x4 blocks) and 4 rows, computed lane-wise so the compiler +/// can vectorise it. Uses |p+q| + |p-q| = 2 max(|p|, |q|) for the last butterfly, which also absorbs the /2. +#[inline(always)] +fn satd_strip(a: &[u8], sa: usize, b: &[u8], sb: usize) -> u32 { + let mut d = [[0i16; N]; 4]; + for r in 0..4 { + let ra = &a[r * sa..r * sa + N]; + let rb = &b[r * sb..r * sb + N]; + for i in 0..N { + d[r][i] = ra[i] as i16 - rb[i] as i16; + } + } + let mut v = [[0i16; N]; 4]; + for i in 0..N { + let s01 = d[0][i] + d[1][i]; + let d01 = d[0][i] - d[1][i]; + let s23 = d[2][i] + d[3][i]; + let d23 = d[2][i] - d[3][i]; + v[0][i] = s01 + s23; + v[1][i] = s01 - s23; + v[2][i] = d01 + d23; + v[3][i] = d01 - d23; + } + let mut total = 0u32; + for row in &v { + let mut hs = [0i16; H]; + let mut hd = [0i16; H]; + for k in 0..H { + hs[k] = row[2 * k] + row[2 * k + 1]; + hd[k] = row[2 * k] - row[2 * k + 1]; + } + for k in 0..H / 2 { + total += hs[2 * k].unsigned_abs().max(hs[2 * k + 1].unsigned_abs()) as u32; + total += hd[2 * k].unsigned_abs().max(hd[2 * k + 1].unsigned_abs()) as u32; + } + } + total +} + +/// Sum of absolute 4x4 Hadamard-transformed differences (halved), for w, h multiples of 4. +pub fn satd(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 { + let mut total = 0; + for by in (0..h).step_by(4) { + let (ao, bo) = (by * sa, by * sb); + match w { + 16 => total += satd_strip::<16, 8>(&a[ao..], sa, &b[bo..], sb), + 8 => total += satd_strip::<8, 4>(&a[ao..], sa, &b[bo..], sb), + 4 => total += satd_strip::<4, 2>(&a[ao..], sa, &b[bo..], sb), + _ => { + for bx in (0..w).step_by(4) { + total += satd_strip::<4, 2>(&a[ao + bx..], sa, &b[bo + bx..], sb); + } + } + } + } + total +} + +pub fn ssd(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u64 { + let mut s = 0u64; + for r in 0..h { + let ra = &a[r * sa..r * sa + w]; + let rb = &b[r * sb..r * sb + w]; + s += ra.iter().zip(rb).map(|(&x, &y)| (x as i32 - y as i32).pow(2) as u32).sum::() as u64; + } + s +} + +/// Variance*N (sum of squares minus square of sum / N) of a block. +pub fn block_var(a: &[u8], sa: usize, w: usize, h: usize) -> u32 { + let mut s = 0u32; + let mut ss = 0u32; + for r in 0..h { + for &v in &a[r * sa..r * sa + w] { + s += v as u32; + ss += (v as u32) * (v as u32); + } + } + ss - ((s as u64 * s as u64) / (w * h) as u64) as u32 +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn satd_of_dc_difference() { + let a = [10u8; 16]; + let b = [0u8; 16]; + // DC-only difference of 10: hadamard DC = 160, halved = 80. + assert_eq!(satd(&a, 4, &b, 4, 4, 4), 80); + assert_eq!(sad(&a, 4, &b, 4, 4, 4), 160); + } + + #[test] + fn satd_matches_reference() { + let mut seed = 7u32; + let mut rnd = || { + seed = seed.wrapping_mul(1103515245).wrapping_add(12345); + (seed >> 16) as u8 + }; + let a: Vec = (0..16 * 16).map(|_| rnd()).collect(); + let b: Vec = (0..16 * 16).map(|_| rnd()).collect(); + for (w, h) in [(16, 16), (8, 8), (4, 4), (16, 8), (8, 16), (12, 4)] { + let mut want = 0; + for by in (0..h).step_by(4) { + for bx in (0..w).step_by(4) { + let mut d = [0i32; 16]; + for r in 0..4 { + for c in 0..4 { + d[r * 4 + c] = a[(by + r) * 16 + bx + c] as i32 - b[(by + r) * 16 + bx + c] as i32; + } + } + want += hadamard4x4_sum(&d); + } + } + assert_eq!(satd(&a, 16, &b, 16, w, h), want / 2, "{w}x{h}"); + } + } +} diff --git a/crates/h264enc/src/intra.rs b/crates/h264enc/src/intra.rs new file mode 100644 index 00000000..cae6fe81 --- /dev/null +++ b/crates/h264enc/src/intra.rs @@ -0,0 +1,303 @@ +//! Intra prediction (§8.3): 4x4, 8x8 (with reference filtering), 16x16 and chroma. + +/// Neighbouring samples of a block. `top` holds p[x,-1] for x = 0..2N (top then top-right), `left` p[-1,y]. +#[derive(Clone, Copy)] +pub struct Edge { + pub top: [u8; N2], + pub left: [u8; N], + pub tl: u8, + pub has_top: bool, + pub has_left: bool, + pub has_tl: bool, +} + +pub type Edge4 = Edge<8, 4>; +pub type Edge8 = Edge<16, 8>; + +/// Which modes are usable given neighbour availability. +pub fn mode_available(mode: u8, has_top: bool, has_left: bool, has_tl: bool) -> bool { + match mode { + 0 | 3 | 7 => has_top, + 1 | 8 => has_left, + 2 => true, + _ => has_top && has_left && has_tl, + } +} + +/// Generic directional prediction for an NxN block given (already filtered, for 8x8) neighbours. +/// `t(x)` = p[x,-1] for x in -1..2N, `l(y)` = p[-1,y] for y in -1..N. +fn pred_dir(mode: u8, t: &dyn Fn(i32) -> i32, l: &dyn Fn(i32) -> i32, dc: i32, out: &mut [u8]) { + let n = N as i32; + for y in 0..n { + for x in 0..n { + let v = match mode { + 0 => t(x), + 1 => l(y), + 2 => dc, + 3 => { + if x == n - 1 && y == n - 1 { + (t(2 * n - 2) + 3 * t(2 * n - 1) + 2) >> 2 + } else { + (t(x + y) + 2 * t(x + y + 1) + t(x + y + 2) + 2) >> 2 + } + } + 4 => { + if x > y { + (t(x - y - 2) + 2 * t(x - y - 1) + t(x - y) + 2) >> 2 + } else if x < y { + (l(y - x - 2) + 2 * l(y - x - 1) + l(y - x) + 2) >> 2 + } else { + (t(0) + 2 * t(-1) + l(0) + 2) >> 2 + } + } + 5 => { + let z = 2 * x - y; + if z >= 0 && z & 1 == 0 { + (t(x - (y >> 1) - 1) + t(x - (y >> 1)) + 1) >> 1 + } else if z >= 0 { + (t(x - (y >> 1) - 2) + 2 * t(x - (y >> 1) - 1) + t(x - (y >> 1)) + 2) >> 2 + } else if z == -1 { + (l(0) + 2 * l(-1) + t(0) + 2) >> 2 + } else { + (l(y - 2 * x - 1) + 2 * l(y - 2 * x - 2) + l(y - 2 * x - 3) + 2) >> 2 + } + } + 6 => { + let z = 2 * y - x; + if z >= 0 && z & 1 == 0 { + (l(y - (x >> 1) - 1) + l(y - (x >> 1)) + 1) >> 1 + } else if z >= 0 { + (l(y - (x >> 1) - 2) + 2 * l(y - (x >> 1) - 1) + l(y - (x >> 1)) + 2) >> 2 + } else if z == -1 { + (l(0) + 2 * l(-1) + t(0) + 2) >> 2 + } else { + (t(x - 2 * y - 1) + 2 * t(x - 2 * y - 2) + t(x - 2 * y - 3) + 2) >> 2 + } + } + 7 => { + if y & 1 == 0 { + (t(x + (y >> 1)) + t(x + (y >> 1) + 1) + 1) >> 1 + } else { + (t(x + (y >> 1)) + 2 * t(x + (y >> 1) + 1) + t(x + (y >> 1) + 2) + 2) >> 2 + } + } + _ => { + let z = x + 2 * y; + let lim = 2 * n - 3; + if z < lim && z & 1 == 0 { + (l(y + (x >> 1)) + l(y + (x >> 1) + 1) + 1) >> 1 + } else if z < lim { + (l(y + (x >> 1)) + 2 * l(y + (x >> 1) + 1) + l(y + (x >> 1) + 2) + 2) >> 2 + } else if z == lim { + (l(n - 2) + 3 * l(n - 1) + 2) >> 2 + } else { + l(n - 1) + } + } + }; + out[(y * n + x) as usize] = v as u8; + } + } +} + +/// Intra 4x4 prediction. `e.top[4..8]` must already be substituted with top[3] when top-right is unavailable. +pub fn pred4x4(mode: u8, e: &Edge4, out: &mut [u8; 16]) { + let dc = match (e.has_top, e.has_left) { + (true, true) => (e.top[..4].iter().map(|&v| v as i32).sum::() + e.left.iter().map(|&v| v as i32).sum::() + 4) >> 3, + (true, false) => (e.top[..4].iter().map(|&v| v as i32).sum::() + 2) >> 2, + (false, true) => (e.left.iter().map(|&v| v as i32).sum::() + 2) >> 2, + _ => 128, + }; + if mode == 2 { + out.fill(dc as u8); + return; + } + let t = |x: i32| if x < 0 { e.tl as i32 } else { e.top[x as usize] as i32 }; + let l = |y: i32| if y < 0 { e.tl as i32 } else { e.left[y as usize] as i32 }; + pred_dir::<4>(mode, &t, &l, dc, out); +} + +/// Reference sample filtering for Intra 8x8 (§8.3.2.2.1). `e.top[8..16]` must be substituted if unavailable. +pub fn filter8(e: &Edge8) -> Edge8 { + let mut f = *e; + if e.has_top { + let p = |x: usize| e.top[x] as u32; + f.top[0] = if e.has_tl { ((e.tl as u32 + 2 * p(0) + p(1) + 2) >> 2) as u8 } else { ((3 * p(0) + p(1) + 2) >> 2) as u8 }; + for x in 1..15 { + f.top[x] = ((p(x - 1) + 2 * p(x) + p(x + 1) + 2) >> 2) as u8; + } + f.top[15] = ((p(14) + 3 * p(15) + 2) >> 2) as u8; + } + if e.has_tl { + let tl = e.tl as u32; + f.tl = if !e.has_top || !e.has_left { + if e.has_top { + ((3 * tl + e.top[0] as u32 + 2) >> 2) as u8 + } else if e.has_left { + ((3 * tl + e.left[0] as u32 + 2) >> 2) as u8 + } else { + e.tl + } + } else { + ((e.top[0] as u32 + 2 * tl + e.left[0] as u32 + 2) >> 2) as u8 + }; + } + if e.has_left { + let p = |y: usize| e.left[y] as u32; + f.left[0] = if e.has_tl { ((e.tl as u32 + 2 * p(0) + p(1) + 2) >> 2) as u8 } else { ((3 * p(0) + p(1) + 2) >> 2) as u8 }; + for y in 1..7 { + f.left[y] = ((p(y - 1) + 2 * p(y) + p(y + 1) + 2) >> 2) as u8; + } + f.left[7] = ((p(6) + 3 * p(7) + 2) >> 2) as u8; + } + f +} + +/// Intra 8x8 prediction from *filtered* neighbours. +pub fn pred8x8(mode: u8, f: &Edge8, out: &mut [u8; 64]) { + let dc = match (f.has_top, f.has_left) { + (true, true) => (f.top[..8].iter().map(|&v| v as i32).sum::() + f.left.iter().map(|&v| v as i32).sum::() + 8) >> 4, + (true, false) => (f.top[..8].iter().map(|&v| v as i32).sum::() + 4) >> 3, + (false, true) => (f.left.iter().map(|&v| v as i32).sum::() + 4) >> 3, + _ => 128, + }; + if mode == 2 { + out.fill(dc as u8); + return; + } + let t = |x: i32| if x < 0 { f.tl as i32 } else { f.top[x as usize] as i32 }; + let l = |y: i32| if y < 0 { f.tl as i32 } else { f.left[y as usize] as i32 }; + pred_dir::<8>(mode, &t, &l, dc, out); +} + +/// Intra 16x16 prediction. Modes: 0 V, 1 H, 2 DC, 3 plane. +pub fn pred16x16(mode: u8, top: &[u8; 16], left: &[u8; 16], tl: u8, has_top: bool, has_left: bool, out: &mut [u8; 256]) { + match mode { + 0 => { + for y in 0..16 { + out[y * 16..y * 16 + 16].copy_from_slice(top); + } + } + 1 => { + for y in 0..16 { + out[y * 16..y * 16 + 16].fill(left[y]); + } + } + 2 => { + let st: u32 = top.iter().map(|&v| v as u32).sum(); + let sl: u32 = left.iter().map(|&v| v as u32).sum(); + let dc = match (has_top, has_left) { + (true, true) => (st + sl + 16) >> 5, + (true, false) => (st + 8) >> 4, + (false, true) => (sl + 8) >> 4, + _ => 128, + }; + out.fill(dc as u8); + } + _ => { + let p_t = |x: i32| if x < 0 { tl as i32 } else { top[x as usize] as i32 }; + let p_l = |y: i32| if y < 0 { tl as i32 } else { left[y as usize] as i32 }; + let mut h = 0; + let mut v = 0; + for i in 0..8 { + h += (i + 1) * (p_t(8 + i) - p_t(6 - i)); + v += (i + 1) * (p_l(8 + i) - p_l(6 - i)); + } + let a = 16 * (p_l(15) + p_t(15)); + let b = (5 * h + 32) >> 6; + let c = (5 * v + 32) >> 6; + for y in 0..16 { + for x in 0..16 { + out[y * 16 + x] = ((a + b * (x as i32 - 7) + c * (y as i32 - 7) + 16) >> 5).clamp(0, 255) as u8; + } + } + } + } +} + +/// Chroma 8x8 prediction (4:2:0). Modes: 0 DC, 1 H, 2 V, 3 plane. +pub fn pred_chroma(mode: u8, top: &[u8; 8], left: &[u8; 8], tl: u8, has_top: bool, has_left: bool, out: &mut [u8; 64]) { + match mode { + 0 => { + for by in 0..2 { + for bx in 0..2 { + let st: u32 = top[bx * 4..bx * 4 + 4].iter().map(|&v| v as u32).sum(); + let sl: u32 = left[by * 4..by * 4 + 4].iter().map(|&v| v as u32).sum(); + let dc = if (bx == 0 && by == 0) || (bx == 1 && by == 1) { + match (has_top, has_left) { + (true, true) => (st + sl + 4) >> 3, + (true, false) => (st + 2) >> 2, + (false, true) => (sl + 2) >> 2, + _ => 128, + } + } else if bx == 1 && by == 0 { + if has_top { + (st + 2) >> 2 + } else if has_left { + (sl + 2) >> 2 + } else { + 128 + } + } else if has_left { + (sl + 2) >> 2 + } else if has_top { + (st + 2) >> 2 + } else { + 128 + }; + for y in 0..4 { + for x in 0..4 { + out[(by * 4 + y) * 8 + bx * 4 + x] = dc as u8; + } + } + } + } + } + 1 => { + for y in 0..8 { + out[y * 8..y * 8 + 8].fill(left[y]); + } + } + 2 => { + for y in 0..8 { + out[y * 8..y * 8 + 8].copy_from_slice(top); + } + } + _ => { + let p_t = |x: i32| if x < 0 { tl as i32 } else { top[x as usize] as i32 }; + let p_l = |y: i32| if y < 0 { tl as i32 } else { left[y as usize] as i32 }; + let mut h = 0; + let mut v = 0; + for i in 0..4 { + h += (i + 1) * (p_t(4 + i) - p_t(2 - i)); + v += (i + 1) * (p_l(4 + i) - p_l(2 - i)); + } + let a = 16 * (p_l(7) + p_t(7)); + let b = (34 * h + 32) >> 6; + let c = (34 * v + 32) >> 6; + for y in 0..8 { + for x in 0..8 { + out[y * 8 + x] = ((a + b * (x as i32 - 3) + c * (y as i32 - 3) + 16) >> 5).clamp(0, 255) as u8; + } + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn dc_and_vertical() { + let e = Edge4 { top: [10, 20, 30, 40, 40, 40, 40, 40], left: [1, 2, 3, 4], tl: 5, has_top: true, has_left: true, has_tl: true }; + let mut o = [0u8; 16]; + pred4x4(0, &e, &mut o); + assert_eq!(&o[12..16], &[10, 20, 30, 40]); + pred4x4(2, &e, &mut o); + assert_eq!(o[0], ((100 + 10 + 4) >> 3) as u8); + // horizontal up saturates at bottom-left + pred4x4(8, &e, &mut o); + assert_eq!(o[15], 4); + } +} diff --git a/crates/h264enc/src/lib.rs b/crates/h264enc/src/lib.rs new file mode 100644 index 00000000..3629f171 --- /dev/null +++ b/crates/h264enc/src/lib.rs @@ -0,0 +1,903 @@ +//! Clean-room, pure-Rust H.264/AVC encoder (ITU-T H.264), written from the public specification. +//! +//! - Profiles: Constrained Baseline (CAVLC, I/P), Main (CABAC, I/P/B) and High (CABAC, 8x8 transform, I/P/B), 4:2:0 8-bit. +//! - Intra 16x16 / 4x4 / 8x8 mode decision, hexagon motion search with quarter-sample refinement, +//! 16x16/16x8/8x16/8x8 partitions, P_Skip, B-frames with spatial direct/B_Skip, in-loop deblocking. +//! - Rate control: constant QP, constant quality (CRF-like), 1-pass ABR with VBV, 2-pass VBR. +//! - Slice-parallel encoding with rayon (feature `threads`, on by default). +//! +//! ```no_run +//! use filmcraft_h264enc::{Encoder, EncoderConfig, RateControl, YuvFrame}; +//! let mut cfg = EncoderConfig::new(1280, 720, 30, 1); +//! cfg.rate = RateControl::Crf(20.0); +//! let mut enc = Encoder::new(cfg).unwrap(); +//! let (y, u, v) = (vec![128u8; 1280 * 720], vec![128u8; 640 * 360], vec![128u8; 640 * 360]); +//! let frame = YuvFrame { y: &y, u: &u, v: &v, y_stride: 1280, uv_stride: 640 }; +//! let mut packets = enc.encode(&frame, 0); +//! packets.extend(enc.flush()); +//! ``` + +mod cabac; +mod cabac_mb; +mod cavlc; +mod deblock; +mod dsp; +mod intra; +mod lookahead; +mod mbinfo; +mod nal; +mod picture; +mod ratecontrol; +mod slice; +mod syntax; +mod tables; +mod transform; + +use std::collections::VecDeque; + +pub use ratecontrol::{FrameStat, PassStats}; + +use cabac::CostTable; +use deblock::{BandPlane, DeblockParams, deblock_band}; +use lookahead::{Complexity, LowRes}; +use mbinfo::MbInfo; +use nal::{NAL_AUD, NAL_IDR, NAL_PPS, NAL_SEI, NAL_SLICE, NAL_SPS, Pps, SliceHeader, SliceType, Sps, Vui}; +use picture::{CHROMA_PAD, Frame, LUMA_PAD, RefPic, build_hpel}; +use ratecontrol::RcKind; +use slice::{Band, EncParams, FrameEnc, Lambdas, SliceEnc, SliceOut}; +use tables::NUM_CTX; +use transform::QuantTables; + +/// Encoder errors. +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("invalid dimensions {0}x{1} (must be even and non-zero)")] + InvalidDimensions(u32, u32), + #[error("invalid frame rate {0}/{1}")] + InvalidFrameRate(u32, u32), + #[error("invalid configuration: {0}")] + InvalidConfig(String), + #[error("input frame planes too small for the configured size")] + BadFrame, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum Profile { + /// Constrained Baseline: CAVLC, no B-frames. + Baseline, + /// Main: CABAC, B-frames. + Main, + /// High: CABAC, B-frames, 8x8 transform. + #[default] + High, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum Preset { + Speed, + #[default] + Balanced, + Quality, +} + +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum RateControl { + /// Constant quality; roughly the QP of P frames (0–51, lower is better). + Crf(f32), + /// Constant bitrate target (1-pass ABR with a 1-second VBV at `kbps`). + Cbr { kbps: u32 }, + /// Variable bitrate: average `target_kbps`, VBV-limited to `max_kbps`. Use [`Pass`] for 2-pass. + Vbr { target_kbps: u32, max_kbps: u32 }, + /// Constant QP (B frames use QP + 2). + Qp(u8), +} + +/// Multi-pass mode for bitrate-targeted rate control. +#[derive(Clone, Debug, Default, PartialEq)] +pub enum Pass { + #[default] + Single, + /// First pass: encode and collect [`PassStats`] via [`Encoder::pass_stats`]. + First, + /// Second pass using the statistics of the first. + Second(PassStats), +} + +/// Colour description written to the VUI (ITU-T H.273 code points). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct ColorConfig { + pub primaries: u8, + pub transfer: u8, + pub matrix: u8, + pub full_range: bool, +} + +impl Default for ColorConfig { + /// BT.709, limited range. + fn default() -> Self { + ColorConfig { primaries: 1, transfer: 1, matrix: 1, full_range: false } + } +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Default)] +pub enum PacketFormat { + /// Start-code prefixed NAL units (`00 00 00 01`), with AUD and in-band SPS/PPS on keyframes. + #[default] + AnnexB, + /// 4-byte big-endian length-prefixed NAL units (MP4 samples); SPS/PPS live in [`Encoder::avcc`]. + LengthPrefixed, +} + +#[derive(Clone, Debug)] +pub struct EncoderConfig { + pub width: u32, + pub height: u32, + pub fps_num: u32, + pub fps_den: u32, + pub profile: Profile, + pub preset: Preset, + pub rate: RateControl, + /// Maximum distance between IDR frames. + pub keyint: u32, + /// Consecutive B-frames (0–3). Ignored for Baseline. + pub bframes: u8, + pub color: ColorConfig, + /// Worker threads (0 = all available cores). Also bounds the number of slices. + pub threads: usize, + pub pass: Pass, + pub format: PacketFormat, + /// Variance adaptive quantisation strength (0 = off). + pub aq_strength: f32, + /// Sample aspect ratio (1:1 by default). + pub sar: (u16, u16), + /// Insert IDR frames at detected scene cuts. + pub scenecut: bool, + /// Number of slices per picture (0 = derived from `threads`). + pub slices: usize, + /// Access unit delimiters in Annex-B output. + pub aud: bool, +} + +impl EncoderConfig { + pub fn new(width: u32, height: u32, fps_num: u32, fps_den: u32) -> Self { + EncoderConfig { + width, + height, + fps_num, + fps_den, + profile: Profile::High, + preset: Preset::Balanced, + rate: RateControl::Crf(23.0), + keyint: 250, + bframes: 2, + color: ColorConfig::default(), + threads: 0, + pass: Pass::Single, + format: PacketFormat::AnnexB, + aq_strength: 0.6, + sar: (1, 1), + scenecut: true, + slices: 0, + aud: true, + } + } +} + +/// A planar 4:2:0 8-bit input picture of the configured size. +#[derive(Clone, Copy, Debug)] +pub struct YuvFrame<'a> { + pub y: &'a [u8], + pub u: &'a [u8], + pub v: &'a [u8], + pub y_stride: usize, + pub uv_stride: usize, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FrameType { + Idr, + I, + P, + B, +} + +/// One coded picture (access unit). +#[derive(Clone, Debug)] +pub struct Packet { + pub data: Vec, + pub pts: i64, + pub dts: i64, + pub keyframe: bool, + pub frame_type: FrameType, + /// QP used for the picture (slice QP). + pub qp: u8, +} + +/// A reconstructed (decoded) picture, available when [`Encoder::set_recon_capture`] is on. +#[derive(Clone, Debug)] +pub struct ReconFrame { + pub pts: i64, + pub y: Vec, + pub u: Vec, + pub v: Vec, +} + +struct Input { + frame: Frame, + pts: i64, + display: u64, + cplx: Complexity, + aq: Vec, +} + +struct Pool { + #[cfg(feature = "threads")] + pool: Option, +} + +impl Pool { + fn install(&self, f: impl FnOnce() -> R + Send) -> R { + #[cfg(feature = "threads")] + if let Some(p) = &self.pool { + return p.install(f); + } + f() + } +} + +pub struct Encoder { + cfg: EncoderConfig, + params: EncParams, + sps: Sps, + sps_nal: Vec, + pps_nal: Vec, + qt: QuantTables, + cost_tab: CostTable, + lam: Lambdas, + cabac_i: Box<[(i8, i8); NUM_CTX]>, + cabac_pb: Box<[(i8, i8); NUM_CTX]>, + rc: ratecontrol::RateControl, + record_pass1: bool, + pending: VecDeque, + prev_lowres: Option, + display_count: u64, + last_idr: u64, + idr_count: u32, + next_frame_num: u32, + dpb: Vec, + coded_count: u64, + display_pts: VecDeque, + dts_offset: Option, + /// First two presentation times (to derive the frame interval). + all_pts: Vec, + held: Vec, + frame_pool: Vec, + hpel_pool: Vec<[picture::Plane; 3]>, + sei_sent: bool, + slices: usize, + pool: Pool, + recon: Option>, + ref_uid: u32, + max_refs: usize, +} + +const LOG2_MAX_FRAME_NUM: u32 = 8; +const LOG2_MAX_POC_LSB: u32 = 8; + +impl Encoder { + pub fn new(mut cfg: EncoderConfig) -> Result { + if cfg.width == 0 || cfg.height == 0 || !cfg.width.is_multiple_of(2) || !cfg.height.is_multiple_of(2) || cfg.width > 16384 || cfg.height > 16384 { + return Err(Error::InvalidDimensions(cfg.width, cfg.height)); + } + if cfg.fps_num == 0 || cfg.fps_den == 0 { + return Err(Error::InvalidFrameRate(cfg.fps_num, cfg.fps_den)); + } + if cfg.bframes > 3 { + return Err(Error::InvalidConfig("bframes must be 0..=3".into())); + } + if cfg.profile == Profile::Baseline { + cfg.bframes = 0; + } + cfg.keyint = cfg.keyint.max(1); + let mbw = cfg.width.div_ceil(16) as usize; + let mbh = cfg.height.div_ceil(16) as usize; + let fps = cfg.fps_num as f64 / cfg.fps_den as f64; + let cabac = cfg.profile != Profile::Baseline; + let high = cfg.profile == Profile::High; + let params = match cfg.preset { + Preset::Speed => EncParams { + mbw, + mbh, + cabac, + t8x8: high, + i4x4: true, + i8x8: false, + partitions: false, + me_range: 8, + subpel: 2, + subpel_rounds: 0, + rd: false, + chroma_qp_offset: 0, + decimate: true, + adaptive_t8: false, + always_intra: false, + cavlc_clamp: !cabac, + }, + Preset::Balanced => EncParams { + mbw, + mbh, + cabac, + t8x8: high, + i4x4: true, + i8x8: high, + partitions: true, + me_range: 16, + subpel: 2, + subpel_rounds: 1, + rd: false, + chroma_qp_offset: 0, + decimate: true, + adaptive_t8: high, + always_intra: false, + cavlc_clamp: !cabac, + }, + Preset::Quality => EncParams { + mbw, + mbh, + cabac, + t8x8: high, + i4x4: true, + i8x8: high, + partitions: true, + me_range: 24, + subpel: 2, + subpel_rounds: 2, + rd: true, + chroma_qp_offset: 0, + decimate: true, + adaptive_t8: high, + always_intra: true, + cavlc_clamp: !cabac, + }, + }; + let max_refs = if cfg.bframes > 0 { 2 } else { 1 }; + let (kbps_for_level, rc_kind, vbv) = match (&cfg.rate, &cfg.pass) { + (RateControl::Qp(q), _) => (None, RcKind::Qp((*q).min(51)), (None, None)), + (RateControl::Crf(c), _) => (None, RcKind::Crf(c.clamp(0.0, 51.0)), (None, None)), + (RateControl::Cbr { kbps }, Pass::Second(st)) => (Some(*kbps), ratecontrol::RateControl::plan_two_pass(st, *kbps), (Some(*kbps), Some(*kbps))), + (RateControl::Cbr { kbps }, Pass::First) => (Some(*kbps), RcKind::Abr { kbps: *kbps }, (None, None)), + (RateControl::Cbr { kbps }, Pass::Single) => (Some(*kbps), RcKind::Abr { kbps: *kbps }, (Some(*kbps), Some(*kbps))), + (RateControl::Vbr { target_kbps, max_kbps }, Pass::Second(st)) => { + (Some(*max_kbps), ratecontrol::RateControl::plan_two_pass(st, *target_kbps), (Some(*max_kbps), Some(max_kbps * 2))) + } + (RateControl::Vbr { target_kbps, max_kbps }, Pass::First) => (Some(*max_kbps), RcKind::Abr { kbps: *target_kbps }, (None, None)), + (RateControl::Vbr { target_kbps, max_kbps }, Pass::Single) => { + (Some(*max_kbps), RcKind::Abr { kbps: *target_kbps }, (Some(*max_kbps), Some(max_kbps * 2))) + } + }; + if matches!(cfg.rate, RateControl::Cbr { kbps: 0 } | RateControl::Vbr { target_kbps: 0, .. }) { + return Err(Error::InvalidConfig("bitrate must be non-zero".into())); + } + let level = nal::pick_level(mbw as u32, mbh as u32, fps, max_refs as u32, kbps_for_level, high); + let (profile_idc, constraint) = match cfg.profile { + Profile::Baseline => (66u8, 0xC0u8), // constraint_set0 + set1 => Constrained Baseline + Profile::Main => (77, 0x40), + Profile::High => (100, 0x00), + }; + let reorder = if cfg.bframes > 0 { 1 } else { 0 }; + let sps = Sps { + profile_idc, + constraint_flags: constraint, + level_idc: level, + width_mbs: mbw as u32, + height_mbs: mbh as u32, + crop_right: mbw as u32 * 16 - cfg.width, + crop_bottom: mbh as u32 * 16 - cfg.height, + log2_max_frame_num: LOG2_MAX_FRAME_NUM, + log2_max_poc_lsb: LOG2_MAX_POC_LSB, + max_num_ref_frames: max_refs as u32, + vui: Vui { + sar: cfg.sar, + video_full_range: cfg.color.full_range, + colour_primaries: cfg.color.primaries, + transfer: cfg.color.transfer, + matrix: cfg.color.matrix, + num_units_in_tick: cfg.fps_den, + time_scale: cfg.fps_num.saturating_mul(2), + max_num_reorder_frames: reorder, + max_dec_frame_buffering: max_refs as u32, + }, + }; + let pps = Pps { cabac, pic_init_qp: 26, chroma_qp_offset: params.chroma_qp_offset, transform_8x8: params.t8x8, high }; + let sps_nal = nal::nal(3, NAL_SPS, &sps.rbsp()); + let pps_nal = nal::nal(3, NAL_PPS, &pps.rbsp()); + let threads = if cfg.threads == 0 { std::thread::available_parallelism().map_or(1, |n| n.get()) } else { cfg.threads }; + let slices = if cfg.slices > 0 { cfg.slices.min(mbh) } else { threads.min(mbh.div_ceil(4)).max(1) }; + #[cfg(feature = "threads")] + let pool = Pool { pool: if threads > 1 { rayon::ThreadPoolBuilder::new().num_threads(threads).build().ok() } else { None } }; + #[cfg(not(feature = "threads"))] + let pool = Pool {}; + let record_pass1 = cfg.pass == Pass::First; + let rc = ratecontrol::RateControl::new(rc_kind, fps, vbv.0, vbv.1); + Ok(Encoder { + params, + sps, + sps_nal, + pps_nal, + qt: QuantTables::new(), + cost_tab: CostTable::new(), + lam: Lambdas::new(), + cabac_i: Box::new(tables::cabac_init_i()), + cabac_pb: Box::new(tables::cabac_init_pb0()), + rc, + record_pass1, + pending: VecDeque::new(), + prev_lowres: None, + display_count: 0, + last_idr: 0, + idr_count: 0, + next_frame_num: 0, + dpb: Vec::new(), + coded_count: 0, + display_pts: VecDeque::new(), + dts_offset: None, + all_pts: Vec::new(), + held: Vec::new(), + frame_pool: Vec::new(), + hpel_pool: Vec::new(), + sei_sent: false, + slices, + pool, + recon: None, + ref_uid: 0, + max_refs, + cfg, + }) + } + + pub fn config(&self) -> &EncoderConfig { + &self.cfg + } + + /// The SPS and PPS NAL units (without start codes). + pub fn sps_pps(&self) -> (Vec, Vec) { + (self.sps_nal.clone(), self.pps_nal.clone()) + } + + /// AVCDecoderConfigurationRecord (`avcC` box payload) for MP4 muxing (4-byte NAL lengths). + pub fn avcc(&self) -> Vec { + nal::avcc(&self.sps_nal, &self.pps_nal, self.sps.profile_idc, self.sps.constraint_flags, self.sps.level_idc) + } + + /// Number of frames the output is delayed by (B-frame reordering). + pub fn delay(&self) -> usize { + self.cfg.bframes as usize + } + + /// Enable capture of reconstructed pictures (for verification). + pub fn set_recon_capture(&mut self, on: bool) { + self.recon = if on { Some(Vec::new()) } else { None }; + } + + /// Take the reconstructed pictures captured so far (coding order). + pub fn take_recon(&mut self) -> Vec { + self.recon.as_mut().map(std::mem::take).unwrap_or_default() + } + + /// First-pass statistics (only when configured with [`Pass::First`]). + pub fn pass_stats(&self) -> Option { + if !self.record_pass1 { + return None; + } + Some(PassStats { frames: self.rc.pass1.clone(), fps: self.cfg.fps_num as f64 / self.cfg.fps_den as f64 }) + } + + /// Encode one picture. Returns zero or more packets in decoding order. + /// + /// # Panics + /// If the planes are smaller than the configured size (see [`Encoder::try_encode`]). + pub fn encode(&mut self, frame: &YuvFrame, pts: i64) -> Vec { + self.try_encode(frame, pts).expect("invalid input frame") + } + + pub fn try_encode(&mut self, frame: &YuvFrame, pts: i64) -> Result, Error> { + let src = self.import(frame)?; + let use_aq = self.cfg.aq_strength > 0.0 && !matches!(self.cfg.rate, RateControl::Qp(_)); + let (mbw, mbh, strength) = (self.params.mbw, self.params.mbh, self.cfg.aq_strength); + let prev = self.prev_lowres.take(); + let (lowres, cplx, aq) = self.pool.install(|| { + let lowres = LowRes::new(&src); + let cplx = lookahead::complexity(&lowres, prev.as_ref()); + let aq = if use_aq { lookahead::aq_offsets(&src, mbw, mbh, strength) } else { Vec::new() }; + (lowres, cplx, aq) + }); + self.prev_lowres = Some(lowres); + let display = self.display_count; + self.display_count += 1; + self.display_pts.push_back(pts); + if self.all_pts.len() < 2 { + self.all_pts.push(pts); + } + let input = Input { frame: src, pts, display, cplx, aq }; + let since = display - self.last_idr; + let fps = self.cfg.fps_num as f64 / self.cfg.fps_den as f64; + let min_gap = (fps / 2.0).max(1.0) as u64; + let scenecut = self.cfg.scenecut && display > 0 && since >= min_gap && cplx.inter as f64 > 0.75 * cplx.intra as f64; + let idr = display == 0 || since >= self.cfg.keyint as u64 || scenecut; + let mut out = Vec::new(); + if idr { + self.flush_minigop(&mut out); + self.last_idr = display; + self.encode_batch(vec![(input, SliceType::I, true)], &mut out); + } else { + self.pending.push_back(input); + if self.pending.len() > self.cfg.bframes as usize { + self.flush_minigop(&mut out); + } + } + Ok(self.release(out, false)) + } + + /// Encode all buffered pictures. + pub fn flush(&mut self) -> Vec { + let mut out = Vec::new(); + self.flush_minigop(&mut out); + self.release(out, true) + } + + /// Assign decoding timestamps. The k-th coded picture gets the k-th smallest presentation time minus the + /// reordering delay (one frame interval when B-frames are enabled); packets are held back until that interval is known. + fn release(&mut self, out: Vec, flushing: bool) -> Vec { + self.held.extend(out); + if self.dts_offset.is_none() { + if self.cfg.bframes == 0 { + self.dts_offset = Some(0); + } else if self.all_pts.len() >= 2 { + let mut v = self.all_pts.clone(); + v.sort_unstable(); + self.dts_offset = Some(v[1] - v[0]); + } else if flushing { + self.dts_offset = Some(0); + } else { + return Vec::new(); + } + } + let off = self.dts_offset.unwrap_or(0); + let mut res = std::mem::take(&mut self.held); + for p in res.iter_mut() { + p.dts -= off; + } + res + } + + fn flush_minigop(&mut self, out: &mut Vec) { + let Some(p) = self.pending.pop_back() else { return }; + let bs: Vec<(Input, SliceType, bool)> = self.pending.drain(..).map(|b| (b, SliceType::B, false)).collect(); + self.encode_batch(vec![(p, SliceType::P, false)], out); + if !bs.is_empty() { + // B-frames are non-reference pictures: encode them all concurrently. + self.encode_batch(bs, out); + } + } + + fn import(&mut self, f: &YuvFrame) -> Result { + let (w, h) = (self.cfg.width as usize, self.cfg.height as usize); + let (cw, ch) = (w / 2, h / 2); + if f.y_stride < w + || f.uv_stride < cw + || f.y.len() < f.y_stride * (h - 1) + w + || f.u.len() < f.uv_stride * (ch - 1) + cw + || f.v.len() < f.uv_stride * (ch - 1) + cw + { + return Err(Error::BadFrame); + } + let mut fr = self.new_frame(); + for (pl, src, stride, pw, ph) in [(&mut fr.y, f.y, f.y_stride, w, h), (&mut fr.u, f.u, f.uv_stride, cw, ch), (&mut fr.v, f.v, f.uv_stride, cw, ch)] { + let aw = pl.w; + for y in 0..pl.h { + let sy = y.min(ph - 1); + let o = pl.idx(0, y as isize); + pl.data[o..o + pw].copy_from_slice(&src[sy * stride..sy * stride + pw]); + let last = pl.data[o + pw - 1]; + pl.data[o + pw..o + aw].fill(last); + } + } + Ok(fr) + } + + fn new_frame(&mut self) -> Frame { + self.frame_pool.pop().unwrap_or_else(|| Frame::new(self.params.mbw * 16, self.params.mbh * 16)) + } + + fn recycle(&mut self, f: Frame) { + if self.frame_pool.len() < 8 { + self.frame_pool.push(f); + } + } + + /// Encode a batch of pictures that do not reference each other (a single I/P picture, or the B-frames of a + /// mini-GOP). Slices of all pictures in the batch are encoded in parallel. + fn encode_batch(&mut self, pics: Vec<(Input, SliceType, bool)>, out: &mut Vec) { + struct Job { + input: Input, + st: SliceType, + idr: bool, + qpf: f64, + qp: u8, + cplx: f64, + aq: Vec, + hdr: SliceHeader, + nal_ref_idc: u8, + l0: Option, + l1: Option, + poc: i32, + rec: Option, + mbs: Vec, + results: Vec, + } + let mut jobs: Vec = Vec::with_capacity(pics.len()); + for (input, mut st, idr) in pics { + if idr { + let old: Vec = self.dpb.drain(..).collect(); + for r in old { + self.recycle_ref(r); + } + self.next_frame_num = 0; + } + let poc = 2 * (input.display as i64 - self.last_idr as i64) as i32; + let by_id = self.dpb.iter().enumerate().max_by_key(|(_, r)| r.id).map(|(i, _)| i); + let (mut l0, mut l1) = match st { + SliceType::I => (None, None), + SliceType::P => (by_id, None), + SliceType::B => ( + self.dpb.iter().enumerate().filter(|(_, r)| r.poc < poc).max_by_key(|(_, r)| r.poc).map(|(i, _)| i), + self.dpb.iter().enumerate().filter(|(_, r)| r.poc > poc).min_by_key(|(_, r)| r.poc).map(|(i, _)| i), + ), + }; + if st == SliceType::B && (l0.is_none() || l1.is_none()) { + st = SliceType::P; + l0 = by_id; + l1 = None; + } + if st == SliceType::P && l0.is_none() { + st = SliceType::I; + } + let is_ref = st != SliceType::B; + let cplx = match st { + SliceType::I => input.cplx.intra, + _ => input.cplx.inter, + } as f64; + let qpf = self.rc.frame_qp(st, cplx, self.coded_count as usize + jobs.len()); + let qp = qpf.round().clamp(0.0, 51.0) as u8; + let aq: Vec = if input.aq.is_empty() { + Vec::new() + } else { + input.aq.iter().map(|&o| ((qpf + o as f64).round().clamp(0.0, 51.0) as i32 - qp as i32) as i8).collect() + }; + let nal_ref_idc = if idr { + 3 + } else if is_ref { + 2 + } else { + 0 + }; + let hdr = SliceHeader { + first_mb: 0, + slice_type: st, + nal_ref_idc, + idr, + idr_pic_id: self.idr_count % 2, + frame_num: self.next_frame_num % (1 << LOG2_MAX_FRAME_NUM), + log2_max_frame_num: LOG2_MAX_FRAME_NUM, + poc_lsb: (poc as u32) % (1 << LOG2_MAX_POC_LSB), + log2_max_poc_lsb: LOG2_MAX_POC_LSB, + l0_modification: None, + cabac: self.params.cabac, + slice_qp: qp as i32, + pic_init_qp: 26, + disable_deblocking: if self.slices > 1 { 2 } else { 0 }, + alpha_offset_div2: 0, + beta_offset_div2: 0, + }; + let rec = self.new_frame(); + let mbs = vec![MbInfo::default(); self.params.mbw * self.params.mbh]; + jobs.push(Job { input, st, idr, qpf, qp, cplx, aq, hdr, nal_ref_idc, l0, l1, poc, rec: Some(rec), mbs, results: Vec::new() }); + } + + // ---- parallel slice encoding (+ deblocking of each slice band) + let (mbw, mbh) = (self.params.mbw, self.params.mbh); + let n = self.slices; + let bounds: Vec<(usize, usize)> = (0..n).map(|s| (s * mbh / n, (s + 1) * mbh / n)).filter(|(a, b)| b > a).collect(); + let capture = self.recon.is_some(); + let mut slice_results: Vec> = (0..jobs.len()).map(|_| Vec::new()).collect(); + let mut bufs: Vec<(Frame, Vec)> = jobs.iter_mut().map(|j| (j.rec.take().unwrap(), std::mem::take(&mut j.mbs))).collect(); + { + let this = &*self; + let fencs: Vec = jobs + .iter() + .map(|j| FrameEnc { + p: &this.params, + qt: &this.qt, + cost_tab: &this.cost_tab, + st: j.hdr.slice_type, + src: &j.input.frame, + l0: j.l0.map(|i| &this.dpb[i]), + l1: j.l1.map(|i| &this.dpb[i]), + qp: j.qp, + aq: &j.aq, + hdr: j.hdr.clone(), + nal_ref_idc: j.nal_ref_idc, + nal_type: if j.idr { NAL_IDR } else { NAL_SLICE }, + cabac_table: if j.hdr.slice_type == SliceType::I { &this.cabac_i } else { &this.cabac_pb }, + lam: &this.lam, + }) + .collect(); + let dps: Vec = jobs + .iter() + .map(|j| DeblockParams { + alpha_offset: 0, + beta_offset: 0, + chroma_qp_offset: this.params.chroma_qp_offset, + ref_ids: [j.l0.map_or(-1, |i| this.dpb[i].id as i32), j.l1.map_or(-1, |i| this.dpb[i].id as i32)], + }) + .collect(); + let deblock: Vec = jobs.iter().map(|j| j.hdr.slice_type != SliceType::B || capture).collect(); + type Task<'t> = (usize, usize, usize, usize, &'t mut [u8], &'t mut [u8], &'t mut [u8], &'t mut [MbInfo]); + let mut tasks: Vec = Vec::new(); + for (pi, (rec, mbs)) in bufs.iter_mut().enumerate() { + let ys = rec.y.stride; + let cs = rec.u.stride; + let mut yrest = &mut rec.y.data[LUMA_PAD * ys..]; + let mut urest = &mut rec.u.data[CHROMA_PAD * cs..]; + let mut vrest = &mut rec.v.data[CHROMA_PAD * cs..]; + let mut mrest = &mut mbs[..]; + for (k, &(r0, r1)) in bounds.iter().enumerate() { + let rows = r1 - r0; + let (a, b) = std::mem::take(&mut yrest).split_at_mut(rows * 16 * ys); + yrest = b; + let (ua, ub) = std::mem::take(&mut urest).split_at_mut(rows * 8 * cs); + urest = ub; + let (va, vb) = std::mem::take(&mut vrest).split_at_mut(rows * 8 * cs); + vrest = vb; + let (ma, mb) = std::mem::take(&mut mrest).split_at_mut(rows * mbw); + mrest = mb; + tasks.push((pi, k, r0, r1, a, ua, va, ma)); + } + } + let ys = jobs_stride_y(mbw); + let cs = jobs_stride_c(mbw); + let run = |(pi, k, r0, r1, y, u, v, m): Task| -> (usize, usize, SliceOut) { + let out = SliceEnc::new(&fencs[pi], Band { y: &mut *y, u: &mut *u, v: &mut *v, ys, cs }, &mut *m, r0, r1, k as u16).encode(); + if deblock[pi] { + let mut by = BandPlane { data: y, stride: ys, pad: LUMA_PAD }; + let mut bu = BandPlane { data: u, stride: cs, pad: CHROMA_PAD }; + let mut bv = BandPlane { data: v, stride: cs, pad: CHROMA_PAD }; + deblock_band(&mut by, &mut bu, &mut bv, m, mbw, r1 - r0, &dps[pi]); + } + (pi, k, out) + }; + #[cfg(feature = "threads")] + let results: Vec<(usize, usize, SliceOut)> = { + use rayon::prelude::*; + if tasks.len() > 1 { this.pool.install(|| tasks.into_par_iter().map(run).collect()) } else { tasks.into_iter().map(run).collect() } + }; + #[cfg(not(feature = "threads"))] + let results: Vec<(usize, usize, SliceOut)> = tasks.into_iter().map(run).collect(); + let mut results = results; + results.sort_by_key(|r| (r.0, r.1)); + for (pi, _, r) in results { + slice_results[pi].push(r); + } + } + for ((j, (rec, mbs)), res) in jobs.iter_mut().zip(bufs).zip(slice_results) { + j.rec = Some(rec); + j.mbs = mbs; + j.results = res; + } + + // ---- finalise each picture in coding order + for job in jobs { + let Job { input, st, idr, qpf, qp, cplx, hdr: _, nal_ref_idc: _, l0: _, l1: _, poc, rec, mbs, results, aq: _ } = job; + let mut rec = rec.unwrap(); + let is_ref = st != SliceType::B; + let annexb = self.cfg.format == PacketFormat::AnnexB; + let mut nals: Vec> = Vec::new(); + if annexb && self.cfg.aud { + let t = match st { + SliceType::I => 0, + SliceType::P => 1, + SliceType::B => 2, + }; + nals.push(nal::nal(0, NAL_AUD, &nal::aud_rbsp(t))); + } + if idr && annexb { + nals.push(self.sps_nal.clone()); + nals.push(self.pps_nal.clone()); + } + if !self.sei_sent { + self.sei_sent = true; + let text = format!("filmcraft-h264enc {} - clean-room H.264 encoder", env!("CARGO_PKG_VERSION")); + nals.push(nal::nal(0, NAL_SEI, &nal::sei_user_data_rbsp(&text))); + } + for r in results { + nals.push(r.nal); + } + let mut data = Vec::with_capacity(nals.iter().map(|n| n.len() + 4).sum()); + for n in &nals { + if annexb { + data.extend_from_slice(&[0, 0, 0, 1]); + } else { + data.extend_from_slice(&(n.len() as u32).to_be_bytes()); + } + data.extend_from_slice(n); + } + let total_bits = data.len() as u64 * 8; + self.rc.update(st, qpf, cplx, total_bits); + if self.record_pass1 { + self.rc.record_pass1(st, qpf, cplx, total_bits); + } + if let Some(rv) = self.recon.as_mut() { + let (w, h) = (self.cfg.width as usize, self.cfg.height as usize); + let mut y = Vec::with_capacity(w * h); + let mut u = Vec::with_capacity(w * h / 4); + let mut v = Vec::with_capacity(w * h / 4); + rec.y.copy_out(w, h, &mut y); + rec.u.copy_out(w / 2, h / 2, &mut u); + rec.v.copy_out(w / 2, h / 2, &mut v); + rv.push(ReconFrame { pts: input.pts, y, u, v }); + } + if is_ref { + let hp = self.hpel_pool.pop(); + let hpel = self.pool.install(|| { + rec.extend_edges(); + build_hpel(&rec.y, hp) + }); + self.ref_uid += 1; + let rp = RefPic { frame: rec, hpel, mbs, poc, id: self.ref_uid }; + if self.dpb.len() >= self.max_refs { + // sliding window: drop the oldest + let oldest = self.dpb.iter().enumerate().min_by_key(|(_, r)| r.id).map(|(i, _)| i).unwrap(); + let r = self.dpb.remove(oldest); + self.recycle_ref(r); + } + self.dpb.push(rp); + self.next_frame_num = (self.next_frame_num + 1) % (1 << LOG2_MAX_FRAME_NUM); + } else { + self.recycle(rec); + } + if idr { + self.idr_count += 1; + } + // decoding timestamp base: the k-th smallest presentation time (offset applied in `release`) + let pos = self.display_pts.iter().enumerate().min_by_key(|(_, p)| **p).map(|(i, _)| i).unwrap(); + let dts_base = self.display_pts.remove(pos).unwrap(); + self.coded_count += 1; + let frame_type = if idr { + FrameType::Idr + } else { + match st { + SliceType::I => FrameType::I, + SliceType::P => FrameType::P, + SliceType::B => FrameType::B, + } + }; + self.recycle(input.frame); + out.push(Packet { data, pts: input.pts, dts: dts_base, keyframe: idr, frame_type, qp }); + } + } + + fn recycle_ref(&mut self, r: RefPic) { + self.recycle(r.frame); + if self.hpel_pool.len() < 4 { + self.hpel_pool.push(r.hpel); + } + } +} + +fn jobs_stride_y(mbw: usize) -> usize { + mbw * 16 + 2 * LUMA_PAD +} + +fn jobs_stride_c(mbw: usize) -> usize { + mbw * 8 + 2 * CHROMA_PAD +} diff --git a/crates/h264enc/src/lookahead.rs b/crates/h264enc/src/lookahead.rs new file mode 100644 index 00000000..c2140b9a --- /dev/null +++ b/crates/h264enc/src/lookahead.rs @@ -0,0 +1,156 @@ +//! Cheap half-resolution analysis used for rate control, scene-cut detection and adaptive quantisation. + +use crate::dsp::{block_var, sad, satd}; +use crate::picture::Frame; + +/// Half-resolution luma (2x2 box filtered) of the macroblock-aligned picture. +pub struct LowRes { + pub w: usize, + pub h: usize, + pub data: Vec, +} + +impl LowRes { + pub fn new(f: &Frame) -> Self { + let (w, h) = (f.y.w / 2, f.y.h / 2); + let mut data = vec![0u8; w * h]; + let row = |y: usize, out: &mut [u8]| { + let r0 = f.y.idx(0, (2 * y) as isize); + let r1 = f.y.idx(0, (2 * y + 1) as isize); + for x in 0..w { + let s = f.y.data[r0 + 2 * x] as u32 + f.y.data[r0 + 2 * x + 1] as u32 + f.y.data[r1 + 2 * x] as u32 + f.y.data[r1 + 2 * x + 1] as u32; + out[x] = ((s + 2) >> 2) as u8; + } + }; + #[cfg(feature = "threads")] + { + use rayon::prelude::*; + data.par_chunks_mut(w).enumerate().for_each(|(y, r)| row(y, r)); + } + #[cfg(not(feature = "threads"))] + for (y, r) in data.chunks_mut(w).enumerate() { + row(y, r); + } + LowRes { w, h, data } + } +} + +/// Per-frame complexity estimates (sum over 8x8 low-res blocks, i.e. per macroblock). +#[derive(Clone, Copy, Debug, Default)] +pub struct Complexity { + pub intra: u64, + /// Inter cost against the previous frame (min with intra per block); equals `intra` when there is no previous frame. + pub inter: u64, +} + +fn block_rows (u64, u64) + Sync + Send>(rows: usize, f: F) -> (u64, u64) { + #[cfg(feature = "threads")] + { + use rayon::prelude::*; + (0..rows).into_par_iter().map(f).reduce(|| (0, 0), |a, b| (a.0 + b.0, a.1 + b.1)) + } + #[cfg(not(feature = "threads"))] + { + (0..rows).map(f).fold((0, 0), |a, b| (a.0 + b.0, a.1 + b.1)) + } +} + +pub fn complexity(cur: &LowRes, prev: Option<&LowRes>) -> Complexity { + let bw = cur.w / 8; + let bh = cur.h / 8; + let (intra, inter) = block_rows(bh, |by| { + let mut ci = 0u64; + let mut cp = 0u64; + let mut blk = [0u8; 64]; + for bx in 0..bw { + let (x0, y0) = (bx * 8, by * 8); + let base = y0 * cur.w + x0; + // intra: best of DC / H / V predictions from source neighbours + let mut best = u32::MAX; + for mode in 0..3 { + if (mode == 1 && x0 == 0) || (mode == 2 && y0 == 0) { + continue; + } + for y in 0..8 { + for x in 0..8 { + blk[y * 8 + x] = match mode { + 0 => { + if x0 > 0 && y0 > 0 { + ((cur.data[(y0 - 1) * cur.w + x0 + x] as u32 + cur.data[(y0 + y) * cur.w + x0 - 1] as u32) / 2) as u8 + } else { + 128 + } + } + 1 => cur.data[(y0 + y) * cur.w + x0 - 1], + _ => cur.data[(y0 - 1) * cur.w + x0 + x], + }; + } + } + best = best.min(satd(&cur.data[base..], cur.w, &blk, 8, 8, 8)); + } + let intra_c = best as u64 + 16; + ci += intra_c; + if let Some(p) = prev { + // small diamond search on SAD, then SATD of the best + let mut bmv = (0i32, 0i32); + let cost_at = |mx: i32, my: i32| -> u32 { + let x = x0 as i32 + mx; + let y = y0 as i32 + my; + if x < 0 || y < 0 || x + 8 > p.w as i32 || y + 8 > p.h as i32 { + return u32::MAX; + } + sad(&cur.data[base..], cur.w, &p.data[y as usize * p.w + x as usize..], p.w, 8, 8) + }; + let mut bc = cost_at(0, 0); + for _ in 0..8 { + let c = bmv; + for (dx, dy) in [(-1, 0), (1, 0), (0, -1), (0, 1), (-2, 0), (2, 0), (0, -2), (0, 2)] { + let k = cost_at(c.0 + dx, c.1 + dy); + if k < bc { + bc = k; + bmv = (c.0 + dx, c.1 + dy); + } + } + if bmv == c { + break; + } + } + let x = (x0 as i32 + bmv.0) as usize; + let y = (y0 as i32 + bmv.1) as usize; + let inter_c = satd(&cur.data[base..], cur.w, &p.data[y * p.w + x..], p.w, 8, 8) as u64 + 8; + cp += inter_c.min(intra_c); + } else { + cp += intra_c; + } + } + (ci, cp) + }); + Complexity { intra, inter } +} + +/// Variance-based adaptive quantisation offsets per macroblock (QP units, mean ~0). +pub fn aq_offsets(f: &Frame, mbw: usize, mbh: usize, strength: f32) -> Vec { + let mut e = vec![0f32; mbw * mbh]; + let row = |my: usize, e: &mut [f32]| { + for mx in 0..mbw { + let i = f.y.idx((mx * 16) as isize, (my * 16) as isize); + let vy = block_var(&f.y.data[i..], f.y.stride, 16, 16) as f32 / 256.0; + let iu = f.u.idx((mx * 8) as isize, (my * 8) as isize); + let vu = block_var(&f.u.data[iu..], f.u.stride, 8, 8) as f32 / 64.0; + let iv = f.v.idx((mx * 8) as isize, (my * 8) as isize); + let vv = block_var(&f.v.data[iv..], f.v.stride, 8, 8) as f32 / 64.0; + e[mx] = (vy + 0.5 * (vu + vv) + 1.0).log2(); + } + }; + #[cfg(feature = "threads")] + { + use rayon::prelude::*; + e.par_chunks_mut(mbw).enumerate().for_each(|(my, r)| row(my, r)); + } + #[cfg(not(feature = "threads"))] + for (my, r) in e.chunks_mut(mbw).enumerate() { + row(my, r); + } + let mean = e.iter().sum::() / e.len() as f32; + e.iter().map(|&v| (strength * (v - mean)).clamp(-8.0, 8.0)).collect() +} diff --git a/crates/h264enc/src/mbinfo.rs b/crates/h264enc/src/mbinfo.rs new file mode 100644 index 00000000..99264d9e --- /dev/null +++ b/crates/h264enc/src/mbinfo.rs @@ -0,0 +1,126 @@ +//! Per-macroblock state kept for neighbour context derivation, deblocking and co-located (direct) prediction. + +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default, Hash)] +pub struct Mv { + pub x: i16, + pub y: i16, +} + +impl Mv { + pub const ZERO: Mv = Mv { x: 0, y: 0 }; + #[inline] + pub const fn new(x: i32, y: i32) -> Mv { + Mv { x: x as i16, y: y as i16 } + } +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)] +pub enum MbKind { + #[default] + None, + I4x4, + I8x8, + I16x16, + PSkip, + PInter, + BSkip, + BDirect, + BInter, +} + +impl MbKind { + #[inline] + pub fn is_intra(self) -> bool { + matches!(self, MbKind::I4x4 | MbKind::I8x8 | MbKind::I16x16) + } + #[inline] + pub fn is_skip(self) -> bool { + matches!(self, MbKind::PSkip | MbKind::BSkip) + } + #[inline] + pub fn is_inxn(self) -> bool { + matches!(self, MbKind::I4x4 | MbKind::I8x8) + } +} + +/// Partition shapes for inter macroblocks. +#[derive(Clone, Copy, PartialEq, Eq, Debug, Default)] +pub enum Part { + #[default] + P16x16, + P16x8, + P8x16, + P8x8, +} + +impl Part { + pub fn count(self) -> usize { + match self { + Part::P16x16 => 1, + Part::P16x8 | Part::P8x16 => 2, + Part::P8x8 => 4, + } + } + /// (x, y, w, h) in pixels of partition `i`. + pub fn rect(self, i: usize) -> (usize, usize, usize, usize) { + match self { + Part::P16x16 => (0, 0, 16, 16), + Part::P16x8 => (0, 8 * i, 16, 8), + Part::P8x16 => (8 * i, 0, 8, 16), + Part::P8x8 => (8 * (i & 1), 8 * (i >> 1), 8, 8), + } + } +} + +#[derive(Clone, Copy, Debug)] +pub struct MbInfo { + pub kind: MbKind, + /// Slice number + 1 (0 = not yet coded). + pub slice: u16, + pub qp: u8, + /// Luma cbp in bits 0..3, chroma cbp in bits 4..5. + pub cbp: u8, + pub t8x8: bool, + pub chroma_mode: u8, + /// Non-zero coefficient counts per luma 4x4 (raster y*4+x). + pub nnz: [u8; 16], + /// Chroma AC non-zero counts: Cb raster 2x2 then Cr raster 2x2. + pub nnz_c: [u8; 8], + /// bit0 luma DC (Intra16x16), bit1 Cb DC, bit2 Cr DC. + pub dc_cbf: u8, + /// Intra 4x4/8x8 prediction modes per 4x4 (raster), -1 when the MB is not I_NxN. + pub imodes: [i8; 16], + /// Reference index per 8x8 (raster) per list; -1 = list unused. + pub ref_idx: [[i8; 4]; 2], + pub mv: [[Mv; 16]; 2], + /// |mvd| components per 4x4 per list, clamped to 127. + pub mvd: [[[u8; 2]; 16]; 2], +} + +impl Default for MbInfo { + fn default() -> Self { + MbInfo { + kind: MbKind::None, + slice: 0, + qp: 0, + cbp: 0, + t8x8: false, + chroma_mode: 0, + nnz: [0; 16], + nnz_c: [0; 8], + dc_cbf: 0, + imodes: [-1; 16], + ref_idx: [[-1; 4]; 2], + mv: [[Mv::ZERO; 16]; 2], + mvd: [[[0; 2]; 16]; 2], + } + } +} + +impl MbInfo { + #[inline] + pub fn ref_at(&self, list: usize, blk4_raster: usize) -> i8 { + let (x, y) = (blk4_raster & 3, blk4_raster >> 2); + self.ref_idx[list][(y >> 1) * 2 + (x >> 1)] + } +} diff --git a/crates/h264enc/src/nal.rs b/crates/h264enc/src/nal.rs new file mode 100644 index 00000000..8ac187b9 --- /dev/null +++ b/crates/h264enc/src/nal.rs @@ -0,0 +1,311 @@ +//! NAL unit and parameter-set/slice-header generation (§7.3), SEI and the AVCDecoderConfigurationRecord. + +use filmcraft_bitstream::{BitWriter, escape_rbsp}; + +pub const NAL_SLICE: u8 = 1; +pub const NAL_IDR: u8 = 5; +pub const NAL_SEI: u8 = 6; +pub const NAL_SPS: u8 = 7; +pub const NAL_PPS: u8 = 8; +pub const NAL_AUD: u8 = 9; + +/// Build a NAL unit (header + escaped RBSP), without start code. +pub fn nal(nal_ref_idc: u8, nal_type: u8, rbsp: &[u8]) -> Vec { + let mut out = Vec::with_capacity(rbsp.len() + rbsp.len() / 64 + 2); + out.push((nal_ref_idc << 5) | nal_type); + out.extend_from_slice(&escape_rbsp(rbsp)); + out +} + +#[derive(Clone, Debug)] +pub struct Vui { + pub sar: (u16, u16), + pub video_full_range: bool, + pub colour_primaries: u8, + pub transfer: u8, + pub matrix: u8, + pub num_units_in_tick: u32, + pub time_scale: u32, + pub max_num_reorder_frames: u32, + pub max_dec_frame_buffering: u32, +} + +#[derive(Clone, Debug)] +pub struct Sps { + pub profile_idc: u8, + pub constraint_flags: u8, + pub level_idc: u8, + pub width_mbs: u32, + pub height_mbs: u32, + pub crop_right: u32, + pub crop_bottom: u32, + pub log2_max_frame_num: u32, + pub log2_max_poc_lsb: u32, + pub max_num_ref_frames: u32, + pub vui: Vui, +} + +impl Sps { + pub fn rbsp(&self) -> Vec { + let mut w = BitWriter::new(); + w.write_bits(self.profile_idc as u32, 8); + w.write_bits(self.constraint_flags as u32, 8); + w.write_bits(self.level_idc as u32, 8); + w.write_ue(0); // seq_parameter_set_id + if matches!(self.profile_idc, 100 | 110 | 122 | 244 | 44 | 83 | 86 | 118 | 128 | 138 | 139 | 134 | 135) { + w.write_ue(1); // chroma_format_idc 4:2:0 + w.write_ue(0); // bit_depth_luma_minus8 + w.write_ue(0); // bit_depth_chroma_minus8 + w.write_bit(false); // qpprime_y_zero_transform_bypass_flag + w.write_bit(false); // seq_scaling_matrix_present_flag + } + w.write_ue(self.log2_max_frame_num - 4); + w.write_ue(0); // pic_order_cnt_type + w.write_ue(self.log2_max_poc_lsb - 4); + w.write_ue(self.max_num_ref_frames); + w.write_bit(false); // gaps_in_frame_num_value_allowed_flag + w.write_ue(self.width_mbs - 1); + w.write_ue(self.height_mbs - 1); + w.write_bit(true); // frame_mbs_only_flag + w.write_bit(true); // direct_8x8_inference_flag + let crop = self.crop_right > 0 || self.crop_bottom > 0; + w.write_bit(crop); + if crop { + w.write_ue(0); + w.write_ue(self.crop_right / 2); + w.write_ue(0); + w.write_ue(self.crop_bottom / 2); + } + w.write_bit(true); // vui_parameters_present_flag + let v = &self.vui; + let has_sar = v.sar != (0, 0); + w.write_bit(has_sar); + if has_sar { + if v.sar == (1, 1) { + w.write_bits(1, 8); + } else { + w.write_bits(255, 8); // Extended_SAR + w.write_bits(v.sar.0 as u32, 16); + w.write_bits(v.sar.1 as u32, 16); + } + } + w.write_bit(false); // overscan_info_present_flag + w.write_bit(true); // video_signal_type_present_flag + w.write_bits(5, 3); // video_format: unspecified + w.write_bit(v.video_full_range); + w.write_bit(true); // colour_description_present_flag + w.write_bits(v.colour_primaries as u32, 8); + w.write_bits(v.transfer as u32, 8); + w.write_bits(v.matrix as u32, 8); + w.write_bit(false); // chroma_loc_info_present_flag + w.write_bit(true); // timing_info_present_flag + w.write_bits(v.num_units_in_tick, 32); + w.write_bits(v.time_scale, 32); + w.write_bit(true); // fixed_frame_rate_flag + w.write_bit(false); // nal_hrd_parameters_present_flag + w.write_bit(false); // vcl_hrd_parameters_present_flag + w.write_bit(false); // pic_struct_present_flag + w.write_bit(true); // bitstream_restriction_flag + w.write_bit(true); // motion_vectors_over_pic_boundaries_flag + w.write_ue(0); // max_bytes_per_pic_denom + w.write_ue(0); // max_bits_per_mb_denom + w.write_ue(11); // log2_max_mv_length_horizontal (±512 samples) + w.write_ue(11); // log2_max_mv_length_vertical + w.write_ue(v.max_num_reorder_frames); + w.write_ue(v.max_dec_frame_buffering); + w.rbsp_trailing(); + w.finish() + } +} + +#[derive(Clone, Debug)] +pub struct Pps { + pub cabac: bool, + pub pic_init_qp: i32, + pub chroma_qp_offset: i32, + pub transform_8x8: bool, + pub high: bool, +} + +impl Pps { + pub fn rbsp(&self) -> Vec { + let mut w = BitWriter::new(); + w.write_ue(0); // pic_parameter_set_id + w.write_ue(0); // seq_parameter_set_id + w.write_bit(self.cabac); + w.write_bit(false); // bottom_field_pic_order_in_frame_present_flag + w.write_ue(0); // num_slice_groups_minus1 + w.write_ue(0); // num_ref_idx_l0_default_active_minus1 + w.write_ue(0); // num_ref_idx_l1_default_active_minus1 + w.write_bit(false); // weighted_pred_flag + w.write_bits(0, 2); // weighted_bipred_idc + w.write_se(self.pic_init_qp - 26); + w.write_se(0); // pic_init_qs_minus26 + w.write_se(self.chroma_qp_offset); + w.write_bit(true); // deblocking_filter_control_present_flag + w.write_bit(false); // constrained_intra_pred_flag + w.write_bit(false); // redundant_pic_cnt_present_flag + if self.high { + w.write_bit(self.transform_8x8); + w.write_bit(false); // pic_scaling_matrix_present_flag + w.write_se(self.chroma_qp_offset); // second_chroma_qp_index_offset + } + w.rbsp_trailing(); + w.finish() + } +} + +/// Access unit delimiter RBSP. `primary_pic_type`: 0 = I, 1 = I/P, 2 = I/P/B. +pub fn aud_rbsp(primary_pic_type: u8) -> Vec { + let mut w = BitWriter::new(); + w.write_bits(primary_pic_type as u32, 3); + w.rbsp_trailing(); + w.finish() +} + +/// user_data_unregistered SEI carrying an encoder identification string. +pub fn sei_user_data_rbsp(text: &str) -> Vec { + const UUID: [u8; 16] = [0x7a, 0x3c, 0x1f, 0x52, 0xd4, 0x0b, 0x4e, 0x8e, 0x9d, 0x61, 0x2f, 0xc4, 0x93, 0x05, 0xb7, 0x1a]; + let mut payload = UUID.to_vec(); + payload.extend_from_slice(text.as_bytes()); + let mut w = BitWriter::new(); + w.write_bits(5, 8); // payloadType user_data_unregistered + let mut size = payload.len(); + while size >= 255 { + w.write_bits(255, 8); + size -= 255; + } + w.write_bits(size as u32, 8); + w.write_bytes(&payload); + w.rbsp_trailing(); + w.finish() +} + +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +pub enum SliceType { + P = 0, + B = 1, + I = 2, +} + +#[derive(Clone, Debug)] +pub struct SliceHeader { + pub first_mb: u32, + pub slice_type: SliceType, + pub nal_ref_idc: u8, + pub idr: bool, + pub idr_pic_id: u32, + pub frame_num: u32, + pub log2_max_frame_num: u32, + pub poc_lsb: u32, + pub log2_max_poc_lsb: u32, + /// Reorder RefPicList0 so that this picNum difference (abs_diff_pic_num_minus1, subtract) comes first. + pub l0_modification: Option, + pub cabac: bool, + pub slice_qp: i32, + pub pic_init_qp: i32, + pub disable_deblocking: u32, + pub alpha_offset_div2: i32, + pub beta_offset_div2: i32, +} + +impl SliceHeader { + /// Write the slice header into `w`; for CABAC the caller must byte-align with one-bits afterwards. + pub fn write(&self, w: &mut BitWriter) { + w.write_ue(self.first_mb); + w.write_ue(self.slice_type as u32); + w.write_ue(0); // pic_parameter_set_id + w.write_bits(self.frame_num, self.log2_max_frame_num); + if self.idr { + w.write_ue(self.idr_pic_id); + } + w.write_bits(self.poc_lsb, self.log2_max_poc_lsb); + if self.slice_type == SliceType::B { + w.write_bit(true); // direct_spatial_mv_pred_flag + } + if self.slice_type != SliceType::I { + w.write_bit(false); // num_ref_idx_active_override_flag + // ref_pic_list_modification + match self.l0_modification { + Some(d) => { + w.write_bit(true); + w.write_ue(0); // modification_of_pic_nums_idc: subtract + w.write_ue(d); + w.write_ue(3); + } + None => w.write_bit(false), + } + if self.slice_type == SliceType::B { + w.write_bit(false); + } + } + if self.nal_ref_idc != 0 { + if self.idr { + w.write_bit(false); // no_output_of_prior_pics_flag + w.write_bit(false); // long_term_reference_flag + } else { + w.write_bit(false); // adaptive_ref_pic_marking_mode_flag + } + } + if self.cabac && self.slice_type != SliceType::I { + w.write_ue(0); // cabac_init_idc + } + w.write_se(self.slice_qp - self.pic_init_qp); + w.write_ue(self.disable_deblocking); + if self.disable_deblocking != 1 { + w.write_se(self.alpha_offset_div2); + w.write_se(self.beta_offset_div2); + } + } +} + +/// AVCDecoderConfigurationRecord (ISO/IEC 14496-15 §5.3.3.1) with 4-byte NAL lengths. +pub fn avcc(sps_nal: &[u8], pps_nal: &[u8], profile_idc: u8, constraint: u8, level: u8) -> Vec { + let mut v = vec![1, profile_idc, constraint, level, 0xFF, 0xE1]; + v.extend_from_slice(&(sps_nal.len() as u16).to_be_bytes()); + v.extend_from_slice(sps_nal); + v.push(1); + v.extend_from_slice(&(pps_nal.len() as u16).to_be_bytes()); + v.extend_from_slice(pps_nal); + if matches!(profile_idc, 100 | 110 | 122 | 144) { + v.push(0xFC | 1); // chroma_format 4:2:0 + v.push(0xF8); // bit_depth_luma_minus8 + v.push(0xF8); // bit_depth_chroma_minus8 + v.push(0); // numOfSequenceParameterSetExt + } + v +} + +/// Pick the lowest level satisfying frame size, macroblock rate, DPB size and (optionally) bitrate. +pub fn pick_level(width_mbs: u32, height_mbs: u32, fps: f64, dpb_frames: u32, kbps: Option, high: bool) -> u8 { + // (level_idc, MaxMBPS, MaxFS, MaxDpbMbs, MaxBR kbit/s) + const LEVELS: [(u8, u64, u64, u64, u64); 17] = [ + (10, 1485, 99, 396, 64), + (11, 3000, 396, 900, 192), + (12, 6000, 396, 2376, 384), + (13, 11880, 396, 2376, 768), + (20, 11880, 396, 2376, 2000), + (21, 19800, 792, 4752, 4000), + (22, 20250, 1620, 8100, 4000), + (30, 40500, 1620, 8100, 10000), + (31, 108000, 3600, 18000, 14000), + (32, 216000, 5120, 20480, 20000), + (40, 245760, 8192, 32768, 20000), + (41, 245760, 8192, 32768, 50000), + (42, 522240, 8704, 34816, 50000), + (50, 589824, 22080, 110400, 135000), + (51, 983040, 36864, 184320, 240000), + (52, 2073600, 36864, 184320, 240000), + (62, 16711680, 139264, 696320, 800000), + ]; + let fs = (width_mbs * height_mbs) as u64; + let mbps = (fs as f64 * fps).ceil() as u64; + for &(l, max_mbps, max_fs, max_dpb, max_br) in &LEVELS { + let br_ok = kbps.is_none_or(|k| (k as u64) * 1000 <= max_br * if high { 1250 } else { 1000 }); + let dim_ok = (width_mbs as u64).pow(2) <= 8 * max_fs && (height_mbs as u64).pow(2) <= 8 * max_fs; + if fs <= max_fs && mbps <= max_mbps && fs * dpb_frames as u64 <= max_dpb && br_ok && dim_ok { + return l; + } + } + 62 +} diff --git a/crates/h264enc/src/picture.rs b/crates/h264enc/src/picture.rs new file mode 100644 index 00000000..ee385ee9 --- /dev/null +++ b/crates/h264enc/src/picture.rs @@ -0,0 +1,259 @@ +//! Padded picture planes, reference pictures and motion-compensated prediction (§8.4.2.2). + +use crate::mbinfo::{MbInfo, Mv}; + +pub const LUMA_PAD: usize = 64; +pub const CHROMA_PAD: usize = 32; + +/// A plane with `pad` replicated samples on every side. +#[derive(Clone)] +pub struct Plane { + pub data: Vec, + pub stride: usize, + pub w: usize, + pub h: usize, + pub pad: usize, +} + +impl Plane { + pub fn new(w: usize, h: usize, pad: usize) -> Self { + let stride = w + 2 * pad; + Plane { data: vec![0; stride * (h + 2 * pad)], stride, w, h, pad } + } + #[inline] + pub fn idx(&self, x: isize, y: isize) -> usize { + ((y + self.pad as isize) as usize) * self.stride + (x + self.pad as isize) as usize + } + /// Replicate edge samples into the padding. + pub fn extend_edges(&mut self) { + let (w, h, pad, stride) = (self.w, self.h, self.pad, self.stride); + for y in 0..h { + let r = (y + pad) * stride; + let l = self.data[r + pad]; + let rr = self.data[r + pad + w - 1]; + self.data[r..r + pad].fill(l); + self.data[r + pad + w..r + stride].fill(rr); + } + let first = pad * stride; + let last = (pad + h - 1) * stride; + for y in 0..pad { + self.data.copy_within(first..first + stride, y * stride); + self.data.copy_within(last..last + stride, (pad + h + y) * stride); + } + } + /// Copy the picture area out as a tightly packed plane of `cw x ch` (cropping). + pub fn copy_out(&self, cw: usize, ch: usize, out: &mut Vec) { + for y in 0..ch { + let o = self.idx(0, y as isize); + out.extend_from_slice(&self.data[o..o + cw]); + } + } +} + +/// A 4:2:0 frame with padded planes. +#[derive(Clone)] +pub struct Frame { + pub y: Plane, + pub u: Plane, + pub v: Plane, +} + +impl Frame { + pub fn new(w: usize, h: usize) -> Self { + Frame { y: Plane::new(w, h, LUMA_PAD), u: Plane::new(w / 2, h / 2, CHROMA_PAD), v: Plane::new(w / 2, h / 2, CHROMA_PAD) } + } + pub fn extend_edges(&mut self) { + self.y.extend_edges(); + self.u.extend_edges(); + self.v.extend_edges(); + } +} + +/// A decoded reference picture with half-sample planes for fast sub-pel access. +pub struct RefPic { + pub frame: Frame, + /// Horizontal half (b), vertical half (h) and centre (j) planes, same geometry as luma. + pub hpel: [Plane; 3], + pub mbs: Vec, + pub poc: i32, + pub id: u32, +} + +#[inline] +fn clip8(v: i32) -> u8 { + v.clamp(0, 255) as u8 +} + +/// Compute the H, V and centre half-sample planes over the whole padded area (clamped at the buffer edge, +/// which equals the normative picture-edge clamping because the padding is replicated). +pub fn build_hpel(f: &Plane, reuse: Option<[Plane; 3]>) -> [Plane; 3] { + let [mut hp, mut vp, mut cp] = reuse.filter(|r| r[0].data.len() == f.data.len()).unwrap_or_else(|| { + let mk = || Plane { data: vec![0; f.data.len()], stride: f.stride, w: f.w, h: f.h, pad: f.pad }; + [mk(), mk(), mk()] + }); + let rows = f.data.len() / f.stride; + let band = 16usize; + let work = |b: usize, hrows: &mut [u8], vrows: &mut [u8], crows: &mut [u8]| { + let r0 = b * band; + let r1 = (r0 + band).min(rows); + hpel_rows(f, r0, r1, hrows, vrows, crows); + }; + let chunk = band * f.stride; + #[cfg(feature = "threads")] + { + use rayon::prelude::*; + hp.data + .par_chunks_mut(chunk) + .zip(vp.data.par_chunks_mut(chunk)) + .zip(cp.data.par_chunks_mut(chunk)) + .enumerate() + .for_each(|(b, ((h, v), c))| work(b, h, v, c)); + } + #[cfg(not(feature = "threads"))] + { + for (b, ((h, v), c)) in hp.data.chunks_mut(chunk).zip(vp.data.chunks_mut(chunk)).zip(cp.data.chunks_mut(chunk)).enumerate() { + work(b, h, v, c); + } + } + [hp, vp, cp] +} + +fn hpel_rows(f: &Plane, r0: usize, r1: usize, hout: &mut [u8], vout: &mut [u8], cout: &mut [u8]) { + let stride = f.stride; + let rows = f.data.len() / stride; + let clampr = |r: isize| r.clamp(0, rows as isize - 1) as usize; + // Horizontal intermediate (unclipped, unshifted) for rows r0-2 .. r1+3. + let nrows = r1 - r0 + 5; + let mut hs = vec![0i16; nrows * stride]; + let mut ext = vec![0u8; stride + 5]; + for k in 0..nrows { + let r = clampr(r0 as isize - 2 + k as isize); + let row = &f.data[r * stride..r * stride + stride]; + ext[0] = row[0]; + ext[1] = row[0]; + ext[2..2 + stride].copy_from_slice(row); + ext[stride + 2] = row[stride - 1]; + ext[stride + 3] = row[stride - 1]; + ext[stride + 4] = row[stride - 1]; + let o = &mut hs[k * stride..k * stride + stride]; + for x in 0..stride { + let e = &ext[x..x + 6]; + o[x] = (e[0] as i16 + e[5] as i16) - 5 * (e[1] as i16 + e[4] as i16) + 20 * (e[2] as i16 + e[3] as i16); + } + } + for r in r0..r1 { + let lo = (r - r0) * stride; + let k = r - r0 + 2; + let hrow = &hs[k * stride..k * stride + stride]; + for x in 0..stride { + hout[lo + x] = clip8((hrow[x] as i32 + 16) >> 5); + } + // vertical + let rr: [usize; 6] = std::array::from_fn(|i| clampr(r as isize - 2 + i as isize) * stride); + for x in 0..stride { + let s = f.data[rr[0] + x] as i32 + f.data[rr[5] + x] as i32 - 5 * (f.data[rr[1] + x] as i32 + f.data[rr[4] + x] as i32) + + 20 * (f.data[rr[2] + x] as i32 + f.data[rr[3] + x] as i32); + vout[lo + x] = clip8((s + 16) >> 5); + } + // centre from horizontal intermediates + let hk: [&[i16]; 6] = std::array::from_fn(|i| &hs[(k - 2 + i) * stride..(k - 1 + i) * stride]); + for x in 0..stride { + let s = hk[0][x] as i32 + hk[5][x] as i32 - 5 * (hk[1][x] as i32 + hk[4][x] as i32) + 20 * (hk[2][x] as i32 + hk[3][x] as i32); + cout[lo + x] = clip8((s + 512) >> 10); + } + } +} + +/// Plane selection per quarter-sample position: two (plane, dx, dy) sources to average (identical = no average). +/// Plane indices: 0 full, 1 horizontal half, 2 vertical half, 3 centre. +const QPEL_SRC: [[(u8, u8, u8); 2]; 16] = [ + // fy = 0 + [(0, 0, 0), (0, 0, 0)], + [(0, 0, 0), (1, 0, 0)], + [(1, 0, 0), (1, 0, 0)], + [(1, 0, 0), (0, 1, 0)], + // fy = 1 + [(0, 0, 0), (2, 0, 0)], + [(1, 0, 0), (2, 0, 0)], + [(1, 0, 0), (3, 0, 0)], + [(1, 0, 0), (2, 1, 0)], + // fy = 2 + [(2, 0, 0), (2, 0, 0)], + [(2, 0, 0), (3, 0, 0)], + [(3, 0, 0), (3, 0, 0)], + [(3, 0, 0), (2, 1, 0)], + // fy = 3 + [(2, 0, 0), (0, 0, 1)], + [(2, 0, 0), (1, 0, 1)], + [(3, 0, 0), (1, 0, 1)], + [(2, 1, 0), (1, 0, 1)], +]; + +impl RefPic { + #[inline] + fn plane(&self, p: u8) -> &Plane { + match p { + 0 => &self.frame.y, + n => &self.hpel[n as usize - 1], + } + } + + /// Luma prediction of a w x h block at picture position (x, y) displaced by `mv` (quarter samples). + pub fn mc_luma(&self, x: usize, y: usize, mv: Mv, w: usize, h: usize, out: &mut [u8], ostride: usize) { + let xi = x as isize + (mv.x as isize >> 2); + let yi = y as isize + (mv.y as isize >> 2); + let q = ((mv.y & 3) * 4 + (mv.x & 3)) as usize; + let [(pa, dxa, dya), (pb, dxb, dyb)] = QPEL_SRC[q]; + let a = self.plane(pa); + let stride = a.stride; + let ia = a.idx(xi + dxa as isize, yi + dya as isize); + if pa == pb && dxa == dxb && dya == dyb { + for r in 0..h { + out[r * ostride..r * ostride + w].copy_from_slice(&a.data[ia + r * stride..ia + r * stride + w]); + } + } else { + let b = self.plane(pb); + let ib = b.idx(xi + dxb as isize, yi + dyb as isize); + for r in 0..h { + let ra = &a.data[ia + r * stride..ia + r * stride + w]; + let rb = &b.data[ib + r * stride..ib + r * stride + w]; + let o = &mut out[r * ostride..r * ostride + w]; + for i in 0..w { + o[i] = ((ra[i] as u16 + rb[i] as u16 + 1) >> 1) as u8; + } + } + } + } + + /// Chroma prediction (both planes) of a w x h chroma block at chroma position (x, y); mv in luma quarter units. + pub fn mc_chroma(&self, x: usize, y: usize, mv: Mv, w: usize, h: usize, out_u: &mut [u8], out_v: &mut [u8], ostride: usize) { + let xi = x as isize + (mv.x as isize >> 3); + let yi = y as isize + (mv.y as isize >> 3); + let fx = (mv.x & 7) as u32; + let fy = (mv.y & 7) as u32; + let w00 = (8 - fx) * (8 - fy); + let w01 = fx * (8 - fy); + let w10 = (8 - fx) * fy; + let w11 = fx * fy; + for (p, out) in [(&self.frame.u, out_u), (&self.frame.v, out_v)] { + let s = p.stride; + let i0 = p.idx(xi, yi); + for r in 0..h { + let a = &p.data[i0 + r * s..i0 + r * s + w + 1]; + let b = &p.data[i0 + (r + 1) * s..i0 + (r + 1) * s + w + 1]; + let o = &mut out[r * ostride..r * ostride + w]; + for i in 0..w { + o[i] = ((w00 * a[i] as u32 + w01 * a[i + 1] as u32 + w10 * b[i] as u32 + w11 * b[i + 1] as u32 + 32) >> 6) as u8; + } + } + } + } +} + +/// Average two predictions in place (default bi-prediction). +#[inline] +pub fn avg_into(dst: &mut [u8], src: &[u8]) { + for (d, &s) in dst.iter_mut().zip(src) { + *d = ((*d as u16 + s as u16 + 1) >> 1) as u8; + } +} diff --git a/crates/h264enc/src/ratecontrol.rs b/crates/h264enc/src/ratecontrol.rs new file mode 100644 index 00000000..386a829c --- /dev/null +++ b/crates/h264enc/src/ratecontrol.rs @@ -0,0 +1,262 @@ +//! Frame-level rate control: constant QP, constant quality (CRF-like), 1-pass ABR with VBV, and 2-pass VBR. + +use crate::nal::SliceType; + +pub fn qp2qscale(qp: f64) -> f64 { + 0.85 * 2f64.powf((qp - 12.0) / 6.0) +} +pub fn qscale2qp(q: f64) -> f64 { + 12.0 + 6.0 * (q / 0.85).log2() +} + +const QCOMP: f64 = 0.6; + +fn type_idx(t: SliceType) -> usize { + match t { + SliceType::I => 0, + SliceType::P => 1, + SliceType::B => 2, + } +} + +/// Relative qscale factor per frame type (I frames finer, B frames coarser). +fn type_factor(t: SliceType) -> f64 { + match t { + SliceType::I => 1.0 / 1.4, + SliceType::P => 1.0, + SliceType::B => 1.3, + } +} + +/// Statistics of one frame from the first pass. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct FrameStat { + pub slice_type: SliceType, + pub qscale: f64, + pub bits: u64, + pub cplx: f64, +} + +/// First-pass statistics, serialisable as text. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct PassStats { + pub frames: Vec, + pub fps: f64, +} + +impl PassStats { + pub fn to_text(&self) -> String { + let mut s = format!("h264enc-pass1 fps={}\n", self.fps); + for f in &self.frames { + let t = match f.slice_type { + SliceType::I => 'I', + SliceType::P => 'P', + SliceType::B => 'B', + }; + s.push_str(&format!("{t} {} {} {}\n", f.qscale, f.bits, f.cplx)); + } + s + } + pub fn from_text(text: &str) -> Option { + let mut lines = text.lines(); + let head = lines.next()?; + let fps = head.strip_prefix("h264enc-pass1 fps=")?.trim().parse().ok()?; + let mut frames = Vec::new(); + for l in lines { + let mut it = l.split_whitespace(); + let t = match it.next()? { + "I" => SliceType::I, + "P" => SliceType::P, + "B" => SliceType::B, + _ => return None, + }; + frames.push(FrameStat { slice_type: t, qscale: it.next()?.parse().ok()?, bits: it.next()?.parse().ok()?, cplx: it.next()?.parse().ok()? }); + } + Some(PassStats { frames, fps }) + } +} + +#[derive(Clone, Debug)] +pub enum RcKind { + Qp(u8), + Crf(f32), + Abr { kbps: u32 }, + TwoPass { kbps: u32, plan: Vec, planned_bits: Vec }, +} + +pub struct RateControl { + kind: RcKind, + fps: f64, + /// bits ≈ coef * cplx / qscale per frame type. + coef: [f64; 3], + coef_n: [u32; 3], + total_bits: f64, + frames: u64, + sum_s: f64, + // VBV + vbv_max_bits_per_frame: f64, + vbv_size: f64, + vbv_fill: f64, + pub pass1: Vec, + model_done: f64, + inflight: std::collections::VecDeque, + last_q: f64, + prior: Option<(f64, f64)>, +} + +impl RateControl { + pub fn new(kind: RcKind, fps: f64, vbv_maxrate_kbps: Option, vbv_buf_kbits: Option) -> Self { + let max = vbv_maxrate_kbps.map_or(0.0, |k| k as f64 * 1000.0 / fps); + let size = vbv_buf_kbits.map_or(0.0, |k| k as f64 * 1000.0); + RateControl { + kind, + fps, + coef: [1.6, 1.4, 1.2], + coef_n: [0; 3], + total_bits: 0.0, + frames: 0, + sum_s: 0.0, + vbv_max_bits_per_frame: max, + vbv_size: size, + vbv_fill: size * 0.9, + pass1: Vec::new(), + model_done: 0.0, + inflight: std::collections::VecDeque::new(), + last_q: 0.0, + prior: None, + } + } + + /// Plan a second pass from first-pass statistics. + pub fn plan_two_pass(stats: &PassStats, kbps: u32) -> RcKind { + let n = stats.frames.len().max(1); + let target = kbps as f64 * 1000.0 * n as f64 / stats.fps; + let c: Vec = stats.frames.iter().map(|f| (f.bits as f64 * f.qscale).max(1.0)).collect(); + let qs = |rf: f64| -> Vec { + stats.frames.iter().zip(&c).map(|(f, &ci)| (rf * ci.powf(1.0 - QCOMP) * type_factor(f.slice_type)).clamp(qp2qscale(0.0), qp2qscale(51.0))).collect() + }; + let bits_for = |q: &[f64]| -> f64 { stats.frames.iter().zip(q).map(|(f, &qi)| f.bits as f64 * f.qscale / qi).sum() }; + // bisection on log(rf) + let (mut lo, mut hi) = (-30.0f64, 30.0f64); + for _ in 0..100 { + let mid = 0.5 * (lo + hi); + if bits_for(&qs(mid.exp())) > target { + lo = mid; + } else { + hi = mid; + } + } + let plan = qs((0.5 * (lo + hi)).exp()); + let planned_bits = stats.frames.iter().zip(&plan).map(|(f, &q)| f.bits as f64 * f.qscale / q).collect(); + RcKind::TwoPass { kbps, plan, planned_bits } + } + + /// Choose the frame QP (fractional) for frame `index` (display order) of `t` with lookahead complexity `cplx`. + pub fn frame_qp(&mut self, t: SliceType, cplx: f64, index_in_coding_order: usize) -> f64 { + let ti = type_idx(t); + let cplx = cplx.max(1.0); + let qp = match &self.kind { + RcKind::Qp(q) => { + let q = *q as f64; + return if t == SliceType::B { (q + 2.0).min(51.0) } else { q }; + } + RcKind::Crf(c) => { + let c = *c as f64; + return (c + match t { + SliceType::I => -2.0, + SliceType::P => 0.0, + SliceType::B => 1.5, + }) + .clamp(0.0, 51.0); + } + RcKind::Abr { kbps } => { + let bpf = *kbps as f64 * 1000.0 / self.fps; + // rate factor so that past frames at this rf would have hit the target + let s_here = self.coef[ti] * cplx.powf(QCOMP) / type_factor(t); + // Virtual history of one second of P frames (inter cost ~ 40% of the first frame's) so the first + // I frame gets a realistic share of the budget instead of a single frame's worth. + if self.prior.is_none() { + let n = self.fps.max(1.0); + let c0 = if t == SliceType::I { 0.4 * cplx } else { cplx }; + self.prior = Some((n, n * self.coef[1] * c0.powf(QCOMP))); + } + let (pn0, ps0) = self.prior.unwrap(); + let fade = (1.0 - self.frames as f64 / pn0).max(0.0); + let (pn, ps) = (pn0 * fade, ps0 * fade); + let s = self.sum_s + s_here + ps; + let wanted = bpf * (self.frames as f64 + 1.0 + pn); + let mut rf = s / wanted; + let abr_buffer = bpf * self.fps.max(1.0); + let overflow = (1.0 + (self.total_bits - bpf * self.frames as f64) / abr_buffer).clamp(0.5, 2.0); + rf *= overflow; + let q = rf * cplx.powf(1.0 - QCOMP) * type_factor(t); + qscale2qp(q) + } + RcKind::TwoPass { kbps, plan, planned_bits } => { + // Closed-loop second pass: `m` measures how far real frame sizes deviate from the first-pass model + // (bits ∝ 1/qscale) at the qscales actually used; the remaining plan is rescaled so its predicted + // size fits the remaining budget. + let i = index_in_coding_order.min(plan.len().saturating_sub(1)); + let bpf = *kbps as f64 * 1000.0 / self.fps; + let target_total: f64 = planned_bits.iter().sum(); + let smooth = 0.5 * bpf * self.fps.max(1.0); + let m = ((self.total_bits + smooth) / (self.model_done + smooth)).clamp(0.3, 3.0); + let inflight: f64 = self.inflight.iter().sum(); + let remaining_plan: f64 = planned_bits[i..].iter().sum(); + let remaining_budget = (target_total - self.total_bits - m * inflight).max(0.05 * remaining_plan.max(1.0)); + let f = (m * remaining_plan / remaining_budget).clamp(0.25, 4.0); + let q = plan[i] * f; + self.inflight.push_back(planned_bits[i] * plan[i] / q); + qscale2qp(q) + } + }; + let mut qp = qp.clamp(1.0, 51.0); + // VBV: make sure the predicted frame size fits in the buffer. + if self.vbv_size > 0.0 { + let fill = self.vbv_fill; + for _ in 0..60 { + let pred = self.coef[ti] * cplx / qp2qscale(qp); + if pred <= fill * 0.8 || qp >= 51.0 { + break; + } + qp += 0.5; + } + } + self.last_q = qp; + qp + } + + /// Update after coding a frame with actual `bits`. + pub fn update(&mut self, t: SliceType, qp: f64, cplx: f64, bits: u64) { + let ti = type_idx(t); + let cplx = cplx.max(1.0); + let q = qp2qscale(qp); + let c = bits as f64 * q / cplx; + let n = self.coef_n[ti]; + let w = if n == 0 { 1.0 } else { 0.35f64.max(1.0 / (n + 1) as f64) }; + self.coef[ti] = self.coef[ti] * (1.0 - w) + c * w; + self.coef_n[ti] += 1; + // If a type has never been observed, keep it proportional to observed ones. + for k in 0..3 { + if self.coef_n[k] == 0 { + self.coef[k] = self.coef[ti] * [1.2, 1.0, 0.85][k] / [1.2, 1.0, 0.85][ti]; + } + } + self.sum_s += self.coef[ti] * cplx.powf(QCOMP) / type_factor(t); + self.total_bits += bits as f64; + if let Some(p) = self.inflight.pop_front() { + self.model_done += p; + } + self.frames += 1; + if self.vbv_size > 0.0 { + self.vbv_fill = (self.vbv_fill - bits as f64 + self.vbv_max_bits_per_frame).min(self.vbv_size); + if self.vbv_fill < 0.0 { + self.vbv_fill = 0.0; + } + } + } + + pub fn record_pass1(&mut self, t: SliceType, qp: f64, cplx: f64, bits: u64) { + self.pass1.push(FrameStat { slice_type: t, qscale: qp2qscale(qp), bits, cplx }); + } +} diff --git a/crates/h264enc/src/slice.rs b/crates/h264enc/src/slice.rs new file mode 100644 index 00000000..39d40e8d --- /dev/null +++ b/crates/h264enc/src/slice.rs @@ -0,0 +1,1904 @@ +//! Slice encoding: macroblock analysis, mode decision, residual coding, reconstruction and entropy coding. + +use filmcraft_bitstream::BitWriter; + +use crate::cabac::{BinSink, CabacEncoder, CabacEstimator, Contexts, CostTable}; +use crate::dsp::{sad, sad16, satd, ssd}; +use crate::intra::{Edge4, Edge8, filter8, mode_available, pred_chroma, pred4x4, pred8x8, pred16x16}; +use crate::mbinfo::{MbInfo, MbKind, Mv, Part}; +use crate::nal::{SliceHeader, SliceType, nal}; +use crate::picture::{CHROMA_PAD, Frame, LUMA_PAD, RefPic, avg_into}; +use crate::syntax::{MbCode, NbCtx}; +use crate::tables::{BLK4_XY, NUM_CTX, RASTER_TO_BLK4, ZIGZAG4, ZIGZAG8, chroma_qp}; +use crate::transform::{ + Deadzone, QuantTables, dequant_chroma_dc, dequant_luma_dc, dequant4, dequant8, fdct4, fdct8, idct4, idct8, quant_chroma_dc, quant_luma_dc, quant4, quant8, +}; +use crate::{cabac_mb, cavlc}; + +/// Encoder tuning derived from the preset. +#[derive(Clone, Debug)] +pub struct EncParams { + pub mbw: usize, + pub mbh: usize, + pub cabac: bool, + pub t8x8: bool, + pub i4x4: bool, + pub i8x8: bool, + pub partitions: bool, + pub me_range: i32, + /// 0 = full-pel, 1 = half-pel, 2 = quarter-pel refinement. + pub subpel: u8, + /// Extra quarter-pel refinement rounds. + pub subpel_rounds: u8, + pub rd: bool, + pub chroma_qp_offset: i32, + pub decimate: bool, + /// Adaptive 4x4/8x8 transform decision for inter macroblocks. + pub adaptive_t8: bool, + /// Always try intra in P/B frames (otherwise only when inter looks poor). + pub always_intra: bool, + /// Maximum CAVLC level magnitude (Baseline). + pub cavlc_clamp: bool, +} + +/// Per-frame shared encoding state. +pub struct FrameEnc<'a> { + pub p: &'a EncParams, + pub qt: &'a QuantTables, + pub cost_tab: &'a CostTable, + pub st: SliceType, + pub src: &'a Frame, + pub l0: Option<&'a RefPic>, + pub l1: Option<&'a RefPic>, + pub qp: u8, + pub aq: &'a [i8], + pub hdr: SliceHeader, + pub nal_ref_idc: u8, + pub nal_type: u8, + pub cabac_table: &'a [(i8, i8); NUM_CTX], + pub lam: &'a Lambdas, +} + +/// Lambda tables indexed by QP. +pub struct Lambdas { + /// SATD/SAD-domain lambda, Q8. + pub sad: [u32; 52], + /// SSD-domain lambda (for RD), Q8 per 1/256 bit. + pub ssd: [f64; 52], +} + +impl Lambdas { + pub fn new() -> Self { + let mut sad = [0u32; 52]; + let mut ssd_ = [0f64; 52]; + for qp in 0..52 { + let l = 0.85 * 2f64.powf((qp as f64 - 12.0) / 3.0); + ssd_[qp] = l; + sad[qp] = ((l.sqrt()) * 256.0).round().max(16.0) as u32; + } + Lambdas { sad, ssd: ssd_ } + } +} + +impl Default for Lambdas { + fn default() -> Self { + Self::new() + } +} + +/// Mutable band of reconstructed rows belonging to one slice. +pub struct Band<'b> { + pub y: &'b mut [u8], + pub u: &'b mut [u8], + pub v: &'b mut [u8], + pub ys: usize, + pub cs: usize, +} + +#[derive(Default, Clone, Debug)] +pub struct SliceStats { + pub bits: u64, + pub intra: u32, + pub skip: u32, + pub inter: u32, + pub qp_sum: u64, + pub mbs: u32, +} + +pub struct SliceOut { + pub nal: Vec, + #[allow(dead_code)] + pub stats: SliceStats, +} + +#[derive(Clone, Copy)] +struct Avail { + left: bool, + top: bool, + tr: bool, + tl: bool, +} + +struct MbEdges { + top: [u8; 24], + left: [u8; 16], + tl: u8, + ctop: [[u8; 8]; 2], + cleft: [[u8; 8]; 2], + ctl: [u8; 2], +} + +#[derive(Clone)] +struct Cand { + code: MbCode, + info: MbInfo, + ry: [u8; 256], + ru: [u8; 64], + rv: [u8; 64], + cost: u64, +} + +impl Cand { + fn new() -> Box { + Box::new(Cand { code: MbCode::default(), info: MbInfo::default(), ry: [0; 256], ru: [0; 64], rv: [0; 64], cost: u64::MAX }) + } +} + +#[inline] +fn se_bits(v: i32) -> u32 { + let k = if v > 0 { 2 * v as u32 - 1 } else { (-2 * v) as u32 }; + 2 * (31 - (k + 1).leading_zeros()) + 1 +} + +#[inline] +fn mv_bits(mv: Mv, mvp: Mv) -> u32 { + se_bits(mv.x as i32 - mvp.x as i32) + se_bits(mv.y as i32 - mvp.y as i32) +} + +fn median(a: i16, b: i16, c: i16) -> i16 { + a.max(b).min(a.min(b).max(c)) +} + +fn decimate_score(levels: &[i16]) -> u32 { + const T: [u32; 16] = [3, 2, 2, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]; + let mut score = 0; + let mut run = 0usize; + for &l in levels { + if l == 0 { + run += 1; + continue; + } + if l.unsigned_abs() > 1 { + return 64; + } + score += T[run.min(15)]; + run = 0; + } + score +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Shape { + Normal, + Top16x8, + Bottom16x8, + Left8x16, + Right8x16, +} + +pub struct SliceEnc<'a, 'b> { + f: &'a FrameEnc<'a>, + band: Band<'b>, + mbs: &'b mut [MbInfo], + y0: usize, + y1: usize, + slice_no: u16, + last_qp: u8, + prev_dqp_nz: bool, + cabac: Option, + cavlc: Option, + skip_run: u32, + stats: SliceStats, +} + +impl<'a, 'b> SliceEnc<'a, 'b> { + pub fn new(f: &'a FrameEnc<'a>, band: Band<'b>, mbs: &'b mut [MbInfo], y0: usize, y1: usize, slice_no: u16) -> Self { + SliceEnc { f, band, mbs, y0, y1, slice_no, last_qp: f.qp, prev_dqp_nz: false, cabac: None, cavlc: None, skip_run: 0, stats: SliceStats::default() } + } + + pub fn encode(mut self) -> SliceOut { + let f = self.f; + let mbw = f.p.mbw; + let mut hdr = f.hdr.clone(); + hdr.first_mb = (self.y0 * mbw) as u32; + let mut bw = BitWriter::new(); + hdr.write(&mut bw); + if f.p.cabac { + while !bw.is_byte_aligned() { + bw.write_bit(true); + } + let ctx = Contexts::init(f.cabac_table, f.qp as i32); + self.cabac = Some(CabacEncoder::new(ctx, bw.finish())); + } else { + self.cavlc = Some(bw); + } + for my in self.y0..self.y1 { + for mx in 0..mbw { + self.encode_mb(mx, my); + if let Some(c) = self.cabac.as_mut() { + let last = my + 1 == self.y1 && mx + 1 == mbw; + if !last { + c.terminate(0); + } + } + } + } + let rbsp = if let Some(c) = self.cabac.take() { + c.finish() + } else { + let mut w = self.cavlc.take().unwrap(); + if self.skip_run > 0 { + w.write_ue(self.skip_run); + } + w.rbsp_trailing(); + w.finish() + }; + self.stats.bits = rbsp.len() as u64 * 8; + SliceOut { nal: nal(f.nal_ref_idc, f.nal_type, &rbsp), stats: self.stats } + } + + // ------------------------------------------------------------------ neighbour access + + #[inline] + fn mb(&self, mx: usize, my: usize) -> &MbInfo { + &self.mbs[(my - self.y0) * self.f.p.mbw + mx] + } + + fn avail(&self, mx: usize, my: usize) -> Avail { + let top = my > self.y0; + Avail { left: mx > 0, top, tr: top && mx + 1 < self.f.p.mbw, tl: top && mx > 0 } + } + + #[inline] + fn ry(&self, x: usize, y: usize) -> usize { + (y - self.y0 * 16) * self.band.ys + LUMA_PAD + x + } + #[inline] + fn rc(&self, x: usize, y: usize) -> usize { + (y - self.y0 * 8) * self.band.cs + CHROMA_PAD + x + } + + fn edges(&self, mx: usize, my: usize, av: Avail) -> MbEdges { + let (px, py) = (mx * 16, my * 16); + let mut e = MbEdges { top: [128; 24], left: [128; 16], tl: 128, ctop: [[128; 8]; 2], cleft: [[128; 8]; 2], ctl: [128; 2] }; + if av.top { + let i = self.ry(px, py - 1); + e.top[..16].copy_from_slice(&self.band.y[i..i + 16]); + if av.tr { + e.top[16..24].copy_from_slice(&self.band.y[i + 16..i + 24]); + } else { + let v = e.top[15]; + e.top[16..24].fill(v); + } + for (c, pl) in [&*self.band.u, &*self.band.v].into_iter().enumerate() { + let i = self.rc(px / 2, py / 2 - 1); + e.ctop[c].copy_from_slice(&pl[i..i + 8]); + } + } + if av.left { + for y in 0..16 { + e.left[y] = self.band.y[self.ry(px - 1, py + y)]; + } + for (c, pl) in [&*self.band.u, &*self.band.v].into_iter().enumerate() { + for y in 0..8 { + e.cleft[c][y] = pl[self.rc(px / 2 - 1, py / 2 + y)]; + } + } + } + if av.tl { + e.tl = self.band.y[self.ry(px - 1, py - 1)]; + e.ctl[0] = self.band.u[self.rc(px / 2 - 1, py / 2 - 1)]; + e.ctl[1] = self.band.v[self.rc(px / 2 - 1, py / 2 - 1)]; + } + e + } + + /// Neighbouring partition data at 4x4 offset (bx, by) relative to the MB (range -1..=4). + /// None = not available; Some((ref, mv)) with ref = -1 for intra / list unused. + fn nb_part(&self, cur: &MbInfo, mask: u16, mx: usize, my: usize, bx: i32, by: i32, list: usize) -> Option<(i8, Mv)> { + if (0..4).contains(&bx) && (0..4).contains(&by) { + let r = (by * 4 + bx) as usize; + if mask & (1 << r) != 0 { + return Some((cur.ref_at(list, r), cur.mv[list][r])); + } + return None; + } + let dx = if bx < 0 { + -1 + } else if bx >= 4 { + 1 + } else { + 0 + }; + let dy = if by < 0 { -1 } else { 0 }; + if by >= 4 || (dy == 0 && dx == 1) { + return None; + } + let nx = mx as i32 + dx; + let ny = my as i32 + dy; + if nx < 0 || nx >= self.f.p.mbw as i32 || ny < self.y0 as i32 { + return None; + } + let m = self.mb(nx as usize, ny as usize); + let r = (((by + 4) % 4) * 4 + (bx + 4) % 4) as usize; + if m.kind.is_intra() { + return Some((-1, Mv::ZERO)); + } + let rf = m.ref_at(list, r); + Some((rf, if rf >= 0 { m.mv[list][r] } else { Mv::ZERO })) + } + + /// Motion vector predictor (§8.4.1.3) for a partition at 4x4 coords (x, y) of size (w, h) (in 4x4 units). + fn mvp(&self, cur: &MbInfo, mask: u16, mx: usize, my: usize, x: i32, y: i32, w: i32, list: usize, rf: i8, shape: Shape) -> Mv { + let a = self.nb_part(cur, mask, mx, my, x - 1, y, list); + let b = self.nb_part(cur, mask, mx, my, x, y - 1, list); + let mut c = self.nb_part(cur, mask, mx, my, x + w, y - 1, list); + if c.is_none() { + c = self.nb_part(cur, mask, mx, my, x - 1, y - 1, list); + } + let (mut b, mut c) = (b, c); + if b.is_none() && c.is_none() && a.is_some() { + b = a; + c = a; + } + let a = a.unwrap_or((-1, Mv::ZERO)); + let b = b.unwrap_or((-1, Mv::ZERO)); + let c = c.unwrap_or((-1, Mv::ZERO)); + match shape { + Shape::Top16x8 if b.0 == rf => return b.1, + Shape::Bottom16x8 if a.0 == rf => return a.1, + Shape::Left8x16 if a.0 == rf => return a.1, + Shape::Right8x16 if c.0 == rf => return c.1, + _ => {} + } + let n = (a.0 == rf) as u32 + (b.0 == rf) as u32 + (c.0 == rf) as u32; + if n == 1 { + if a.0 == rf { + return a.1; + } + if b.0 == rf { + return b.1; + } + return c.1; + } + Mv { x: median(a.1.x, b.1.x, c.1.x), y: median(a.1.y, b.1.y, c.1.y) } + } + + fn pskip_mv(&self, mx: usize, my: usize) -> Mv { + let cur = MbInfo::default(); + let a = self.nb_part(&cur, 0, mx, my, -1, 0, 0); + let b = self.nb_part(&cur, 0, mx, my, 0, -1, 0); + match (a, b) { + (None, _) | (_, None) => Mv::ZERO, + (Some((0, m)), _) if m == Mv::ZERO => Mv::ZERO, + (_, Some((0, m))) if m == Mv::ZERO => Mv::ZERO, + _ => self.mvp(&cur, 0, mx, my, 0, 0, 4, 0, 0, Shape::Normal), + } + } + + /// Whether a luma block at picture (x, y) of size w x h with `mv` stays within the padded reference. + fn mv_ok(&self, x: usize, y: usize, w: usize, h: usize, mv: Mv) -> bool { + let pw = (self.f.p.mbw * 16) as i32; + let ph = (self.f.p.mbh * 16) as i32; + let m = LUMA_PAD as i32 - 20; + let x0 = x as i32 + (mv.x as i32 >> 2); + let y0 = y as i32 + (mv.y as i32 >> 2); + mv.x.unsigned_abs() < 2000 && mv.y.unsigned_abs() < 2000 && x0 >= -m && y0 >= -m && x0 + w as i32 <= pw + m && y0 + h as i32 <= ph + m + } + + // ------------------------------------------------------------------ macroblock + + fn encode_mb(&mut self, mx: usize, my: usize) { + let f = self.f; + let mbw = f.p.mbw; + let aq = if f.aq.is_empty() { 0 } else { f.aq[my * mbw + mx] as i32 }; + let qp = (f.qp as i32 + aq).clamp(0, 51) as u8; + let av = self.avail(mx, my); + let (px, py) = (mx * 16, my * 16); + let mut src_y = [0u8; 256]; + let mut src_u = [0u8; 64]; + let mut src_v = [0u8; 64]; + { + let s = &f.src.y; + for r in 0..16 { + let i = s.idx(px as isize, (py + r) as isize); + src_y[r * 16..r * 16 + 16].copy_from_slice(&s.data[i..i + 16]); + } + for r in 0..8 { + let i = f.src.u.idx((px / 2) as isize, (py / 2 + r) as isize); + src_u[r * 8..r * 8 + 8].copy_from_slice(&f.src.u.data[i..i + 8]); + let i = f.src.v.idx((px / 2) as isize, (py / 2 + r) as isize); + src_v[r * 8..r * 8 + 8].copy_from_slice(&f.src.v.data[i..i + 8]); + } + } + let src = Src { y: src_y, u: src_u, v: src_v }; + let edges = self.edges(mx, my, av); + let best = match f.st { + SliceType::I => self.analyse_intra(mx, my, qp, av, &edges, &src, u64::MAX), + SliceType::P => self.analyse_p(mx, my, qp, av, &edges, &src), + SliceType::B => self.analyse_b(mx, my, qp, av, &edges, &src), + }; + self.commit(mx, my, qp, *best); + } + + fn commit(&mut self, mx: usize, my: usize, qp: u8, mut c: Cand) { + let (px, py) = (mx * 16, my * 16); + // reconstruction + for r in 0..16 { + let i = self.ry(px, py + r); + self.band.y[i..i + 16].copy_from_slice(&c.ry[r * 16..r * 16 + 16]); + } + for r in 0..8 { + let i = self.rc(px / 2, py / 2 + r); + self.band.u[i..i + 8].copy_from_slice(&c.ru[r * 8..r * 8 + 8]); + self.band.v[i..i + 8].copy_from_slice(&c.rv[r * 8..r * 8 + 8]); + } + let kind = c.code.kind; + let has_dqp = !kind.is_skip() && (c.code.cbp != 0 || kind == MbKind::I16x16); + let prev_dqp_nz = self.prev_dqp_nz; + if has_dqp { + let mut d = qp as i32 - self.last_qp as i32; + if d < -26 { + d += 52; + } else if d > 25 { + d -= 52; + } + c.code.qp_delta = d; + self.prev_dqp_nz = d != 0; + self.last_qp = qp; + } else { + c.code.qp_delta = 0; + self.prev_dqp_nz = false; + } + c.info.qp = self.last_qp; + c.info.slice = self.slice_no + 1; + c.info.kind = kind; + c.info.cbp = if kind.is_skip() { 0 } else { c.code.cbp }; + self.stats.mbs += 1; + self.stats.qp_sum += c.info.qp as u64; + if kind.is_intra() { + self.stats.intra += 1; + } else if kind.is_skip() { + self.stats.skip += 1; + } else { + self.stats.inter += 1; + } + let idx = (my - self.y0) * self.f.p.mbw + mx; + self.mbs[idx] = c.info; + // entropy coding + let av = self.avail(mx, my); + let left = if av.left { Some(&self.mbs[idx - 1]) } else { None }; + let top = if av.top { Some(&self.mbs[idx - self.f.p.mbw]) } else { None }; + let nb = NbCtx { left, top, cur: &self.mbs[idx], slice_type: self.f.st, prev_dqp_nz, t8x8_mode: self.f.p.t8x8 }; + if let Some(cab) = self.cabac.as_mut() { + if nb.slice_type != SliceType::I { + cabac_mb::write_skip_flag(cab, &nb, kind.is_skip()); + } + if !kind.is_skip() { + cabac_mb::write_mb(cab, &nb, &c.code); + } + } else if let Some(w) = self.cavlc.as_mut() { + if kind.is_skip() { + self.skip_run += 1; + } else { + if nb.slice_type != SliceType::I { + w.write_ue(self.skip_run); + self.skip_run = 0; + } + cavlc::write_mb(w, &nb, &c.code); + } + } + } + + /// Estimated bits (1/256 units) of coding candidate `c` at the current entropy state. + fn est_bits(&self, mx: usize, my: usize, c: &Cand, qp: u8) -> u64 { + let av = self.avail(mx, my); + let idx = (my - self.y0) * self.f.p.mbw + mx; + let left = if av.left { Some(&self.mbs[idx - 1]) } else { None }; + let top = if av.top { Some(&self.mbs[idx - self.f.p.mbw]) } else { None }; + let mut info = c.info; + info.kind = c.code.kind; + info.cbp = c.code.cbp; + let mut code = c.code.clone(); + code.qp_delta = qp as i32 - self.last_qp as i32; + let nb = NbCtx { left, top, cur: &info, slice_type: self.f.st, prev_dqp_nz: self.prev_dqp_nz, t8x8_mode: self.f.p.t8x8 }; + if let Some(cab) = self.cabac.as_ref() { + let mut est = CabacEstimator::new(cab.ctx.clone(), self.f.cost_tab); + if nb.slice_type != SliceType::I { + cabac_mb::write_skip_flag(&mut est, &nb, code.kind.is_skip()); + } + if !code.kind.is_skip() { + cabac_mb::write_mb(&mut est, &nb, &code); + } + est.cost + } else { + if code.kind.is_skip() { + return 256; + } + let mut w = BitWriter::new(); + if nb.slice_type != SliceType::I { + w.write_ue(self.skip_run); + } + cavlc::write_mb(&mut w, &nb, &code); + w.bit_len() as u64 * 256 + } + } + + fn rd_cost(&self, mx: usize, my: usize, c: &Cand, src: &Src, qp: u8) -> u64 { + let d = ssd(&src.y, 16, &c.ry, 16, 16, 16) + ssd(&src.u, 8, &c.ru, 8, 8, 8) + ssd(&src.v, 8, &c.rv, 8, 8, 8); + let bits = self.est_bits(mx, my, c, qp); + let lam = self.f.lam.ssd[qp as usize]; + (d as f64 * 256.0 + lam * bits as f64) as u64 + } + + // ------------------------------------------------------------------ residual coding helpers + + /// Code chroma residual for prediction `pu`, `pv`; fills code/info and reconstruction. Returns chroma cbp. + fn code_chroma(&self, src: &Src, pu: &[u8; 64], pv: &[u8; 64], qp: u8, intra: bool, c: &mut Cand) -> u8 { + let qt = self.f.qt; + let qpc = chroma_qp(qp as i32 + self.f.p.chroma_qp_offset); + let dz = if intra { Deadzone::INTRA } else { Deadzone::INTER }; + let mut dc_any = false; + let mut ac_any = false; + let mut coefs = [[[0i32; 16]; 4]; 2]; + let mut dcl = [[0i32; 4]; 2]; + for comp in 0..2 { + let (s, p) = if comp == 0 { (&src.u, pu) } else { (&src.v, pv) }; + let mut dc = [0i32; 4]; + let mut ac_score = 0; + for b in 0..4 { + let (bx, by) = ((b & 1) * 4, (b >> 1) * 4); + let mut d = [0i32; 16]; + for y in 0..4 { + for x in 0..4 { + d[y * 4 + x] = s[(by + y) * 8 + bx + x] as i32 - p[(by + y) * 8 + bx + x] as i32; + } + } + let mut t = fdct4(&d); + dc[b] = t[0]; + quant4(&mut t, qt, qpc, dz, true); + if self.f.p.cavlc_clamp { + for v in t.iter_mut().skip(1) { + *v = (*v).clamp(-cavlc::MAX_LEVEL, cavlc::MAX_LEVEL); + } + } + let mut scan = [0i16; 16]; + for k in 1..16 { + scan[k] = t[ZIGZAG4[k]] as i16; + } + ac_score += decimate_score(&scan[1..]); + c.code.cac[comp * 4 + b] = scan; + coefs[comp][b] = t; + } + if !intra && self.f.p.decimate && ac_score < 7 { + for b in 0..4 { + c.code.cac[comp * 4 + b] = [0; 16]; + for k in 1..16 { + coefs[comp][b][k] = 0; + } + } + } + let (mut lv, _) = quant_chroma_dc(&dc, qt, qpc, dz); + if self.f.p.cavlc_clamp { + for v in lv.iter_mut() { + *v = (*v).clamp(-cavlc::MAX_LEVEL, cavlc::MAX_LEVEL); + } + } + dcl[comp] = lv; + for b in 0..4 { + c.code.cdc[comp][b] = lv[b] as i16; + } + dc_any |= lv.iter().any(|&v| v != 0); + ac_any |= c.code.cac[comp * 4..comp * 4 + 4].iter().any(|b| b.iter().any(|&v| v != 0)); + } + let cbp = if ac_any { + 2 + } else if dc_any { + 1 + } else { + 0 + }; + c.info.dc_cbf &= 1; + for comp in 0..2 { + let p = if comp == 0 { pu } else { pv }; + let dcq = dequant_chroma_dc(&dcl[comp], qt, qpc); + if dcl[comp].iter().any(|&v| v != 0) { + c.info.dc_cbf |= 2 << comp; + } + for b in 0..4 { + let (bx, by) = ((b & 1) * 4, (b >> 1) * 4); + let mut t = coefs[comp][b]; + let nz = if cbp == 2 { t.iter().skip(1).filter(|&&v| v != 0).count() as u8 } else { 0 }; + c.info.nnz_c[comp * 4 + b] = nz; + if cbp != 2 { + t = [0; 16]; + } else { + dequant4(&mut t, qt, qpc, true); + } + t[0] = dcq[b]; + let rec = if comp == 0 { &mut c.ru } else { &mut c.rv }; + if nz == 0 { + let dcv = (t[0] + 32) >> 6; + for y in 0..4 { + for x in 0..4 { + let i = (by + y) * 8 + bx + x; + rec[i] = (p[i] as i32 + dcv).clamp(0, 255) as u8; + } + } + } else { + let r = idct4(&t); + for y in 0..4 { + for x in 0..4 { + let i = (by + y) * 8 + bx + x; + rec[i] = (p[i] as i32 + r[y * 4 + x]).clamp(0, 255) as u8; + } + } + } + } + } + if cbp == 0 { + c.code.cdc = [[0; 4]; 2]; + } + if cbp != 2 { + c.code.cac = [[0; 16]; 8]; + } + cbp + } + + /// Inter luma residual with the 4x4 transform. Returns luma cbp. + fn code_luma4(&self, src: &[u8; 256], pred: &[u8; 256], qp: u8, dz: Deadzone, decimate: bool, c: &mut Cand) -> u8 { + let qt = self.f.qt; + let mut coefs = [[0i32; 16]; 16]; + let mut score8 = [0u32; 4]; + for blk in 0..16 { + let (bx, by) = BLK4_XY[blk]; + let mut d = [0i32; 16]; + for y in 0..4 { + for x in 0..4 { + let i = (by * 4 + y) * 16 + bx * 4 + x; + d[y * 4 + x] = src[i] as i32 - pred[i] as i32; + } + } + let mut t = fdct4(&d); + quant4(&mut t, qt, qp, dz, false); + if self.f.p.cavlc_clamp { + for v in t.iter_mut() { + *v = (*v).clamp(-cavlc::MAX_LEVEL, cavlc::MAX_LEVEL); + } + } + let mut scan = [0i16; 16]; + for k in 0..16 { + scan[k] = t[ZIGZAG4[k]] as i16; + } + score8[blk / 4] += decimate_score(&scan); + c.code.luma4[blk] = scan; + coefs[blk] = t; + } + let mut cbp = 0u8; + let total: u32 = score8.iter().sum(); + for b8 in 0..4 { + let kill = decimate && (score8[b8] < 4 || total < 6); + for i in 0..4 { + let blk = b8 * 4 + i; + if kill { + c.code.luma4[blk] = [0; 16]; + coefs[blk] = [0; 16]; + } + if c.code.luma4[blk].iter().any(|&v| v != 0) { + cbp |= 1 << b8; + } + } + } + for blk in 0..16 { + let (bx, by) = BLK4_XY[blk]; + let nz = c.code.luma4[blk].iter().filter(|&&v| v != 0).count() as u8; + c.info.nnz[by * 4 + bx] = nz; + let off = by * 4 * 16 + bx * 4; + if nz == 0 { + for y in 0..4 { + c.ry[off + y * 16..off + y * 16 + 4].copy_from_slice(&pred[off + y * 16..off + y * 16 + 4]); + } + } else { + let mut t = coefs[blk]; + dequant4(&mut t, qt, qp, false); + let r = idct4(&t); + for y in 0..4 { + for x in 0..4 { + let i = off + y * 16 + x; + c.ry[i] = (pred[i] as i32 + r[y * 4 + x]).clamp(0, 255) as u8; + } + } + } + } + cbp + } + + /// Inter luma residual with the 8x8 transform. Returns luma cbp. + fn code_luma8(&self, src: &[u8; 256], pred: &[u8; 256], qp: u8, dz: Deadzone, decimate: bool, c: &mut Cand) -> u8 { + let qt = self.f.qt; + let mut coefs = [[0i32; 64]; 4]; + let mut score = [0u32; 4]; + for b8 in 0..4 { + let (ox, oy) = ((b8 & 1) * 8, (b8 >> 1) * 8); + let mut d = [0i32; 64]; + for y in 0..8 { + for x in 0..8 { + let i = (oy + y) * 16 + ox + x; + d[y * 8 + x] = src[i] as i32 - pred[i] as i32; + } + } + let mut t = fdct8(&d); + quant8(&mut t, qt, qp, dz); + let mut scan = [0i16; 64]; + for k in 0..64 { + scan[k] = t[ZIGZAG8[k]] as i16; + } + score[b8] = decimate_score(&scan); + c.code.luma8[b8] = scan; + coefs[b8] = t; + } + let total: u32 = score.iter().sum(); + let mut cbp = 0; + for b8 in 0..4 { + let (ox, oy) = ((b8 & 1) * 8, (b8 >> 1) * 8); + if decimate && (score[b8] < 4 || total < 6) { + c.code.luma8[b8] = [0; 64]; + coefs[b8] = [0; 64]; + } + let nz = c.code.luma8[b8].iter().filter(|&&v| v != 0).count() as u8; + let (bx, by) = (ox / 4, oy / 4); + for (dx, dy) in [(0, 0), (1, 0), (0, 1), (1, 1)] { + c.info.nnz[(by + dy) * 4 + bx + dx] = nz; + } + if nz == 0 { + for y in 0..8 { + let o = (oy + y) * 16 + ox; + c.ry[o..o + 8].copy_from_slice(&pred[o..o + 8]); + } + continue; + } + cbp |= 1 << b8; + let mut t = coefs[b8]; + dequant8(&mut t, qt, qp); + let r = idct8(&t); + for y in 0..8 { + for x in 0..8 { + let i = (oy + y) * 16 + ox + x; + c.ry[i] = (pred[i] as i32 + r[y * 8 + x]).clamp(0, 255) as u8; + } + } + } + cbp + } + + /// Code an inter candidate's residual given its luma/chroma prediction. + fn code_inter(&self, src: &Src, py: &[u8; 256], pu: &[u8; 64], pv: &[u8; 64], qp: u8, c: &mut Cand) { + let t8_allowed = self.f.p.t8x8; + let use8 = if t8_allowed { + if self.f.p.adaptive_t8 { + // Compare 4x4 and 8x8 transform by SSD + lambda * (approximate) bits. + let mut c4 = c.clone(); + let cbp4 = self.code_luma4(&src.y, py, qp, Deadzone::INTER, self.f.p.decimate, &mut c4); + let mut c8 = c.clone(); + let cbp8 = self.code_luma8(&src.y, py, qp, Deadzone::INTER, self.f.p.decimate, &mut c8); + let lam = self.f.lam.ssd[qp as usize]; + let cost = |cc: &Cand, cbp: u8, eight: bool| { + let d = ssd(&src.y, 16, &cc.ry, 16, 16, 16) as f64; + let mut bits = 0f64; + if eight { + for b in 0..4 { + if (cbp >> b) & 1 != 0 { + bits += cc.code.luma8[b].iter().map(|&v| if v == 0 { 0.4 } else { 3.0 + 2.0 * (v.unsigned_abs() as f64).log2() }).sum::(); + } + } + } else { + for b in 0..16 { + if (cbp >> (b / 4)) & 1 != 0 { + bits += 1.0 + + cc.code.luma4[b].iter().map(|&v| if v == 0 { 0.4 } else { 3.0 + 2.0 * (v.unsigned_abs() as f64).log2() }).sum::(); + } + } + } + d + lam * bits + }; + let _ = (cbp4, cbp8); + cost(&c8, cbp8, true) <= cost(&c4, cbp4, false) + } else { + true + } + } else { + false + }; + let cbpl = if use8 { + self.code_luma8(&src.y, py, qp, Deadzone::INTER, self.f.p.decimate, c) + } else { + self.code_luma4(&src.y, py, qp, Deadzone::INTER, self.f.p.decimate, c) + }; + c.code.t8x8 = use8 && cbpl != 0; + c.info.t8x8 = c.code.t8x8; + let cbpc = self.code_chroma(src, pu, pv, qp, false, c); + c.code.cbp = cbpl | (cbpc << 4); + } + + // ------------------------------------------------------------------ intra + + fn pred_intra_mode(&self, mx: usize, my: usize, av: Avail, local: &[i8; 16], bx: usize, by: usize) -> u8 { + let a = if bx > 0 { + Some(local[by * 4 + bx - 1]) + } else if av.left { + Some(self.mb(mx - 1, my).imodes[by * 4 + 3]) + } else { + None + }; + let b = if by > 0 { + Some(local[(by - 1) * 4 + bx]) + } else if av.top { + Some(self.mb(mx, my - 1).imodes[12 + bx]) + } else { + None + }; + match (a, b) { + (Some(a), Some(b)) => { + let a = if a < 0 { 2 } else { a }; + let b = if b < 0 { 2 } else { b }; + a.min(b) as u8 + } + _ => 2, + } + } + + /// Luma I16x16 coding. Returns luma cbp (0 or 15). + fn code_i16(&self, src: &[u8; 256], pred: &[u8; 256], qp: u8, c: &mut Cand) -> u8 { + let qt = self.f.qt; + let mut coefs = [[0i32; 16]; 16]; + let mut dc = [0i32; 16]; + for blk in 0..16 { + let (bx, by) = BLK4_XY[blk]; + let mut d = [0i32; 16]; + for y in 0..4 { + for x in 0..4 { + let i = (by * 4 + y) * 16 + bx * 4 + x; + d[y * 4 + x] = src[i] as i32 - pred[i] as i32; + } + } + let mut t = fdct4(&d); + dc[by * 4 + bx] = t[0]; + quant4(&mut t, qt, qp, Deadzone::INTRA, true); + if self.f.p.cavlc_clamp { + for v in t.iter_mut() { + *v = (*v).clamp(-cavlc::MAX_LEVEL, cavlc::MAX_LEVEL); + } + } + coefs[blk] = t; + } + let (mut dcl, _) = quant_luma_dc(&dc, qt, qp, Deadzone::INTRA); + if self.f.p.cavlc_clamp { + for v in dcl.iter_mut() { + *v = (*v).clamp(-cavlc::MAX_LEVEL, cavlc::MAX_LEVEL); + } + } + for k in 0..16 { + c.code.luma_dc[k] = dcl[ZIGZAG4[k]] as i16; + } + let ac_any = coefs.iter().any(|t| t[1..].iter().any(|&v| v != 0)); + let dcq = dequant_luma_dc(&dcl, qt, qp); + c.info.dc_cbf = (c.info.dc_cbf & !1) | dcl.iter().any(|&v| v != 0) as u8; + for blk in 0..16 { + let (bx, by) = BLK4_XY[blk]; + let mut t = coefs[blk]; + let mut scan = [0i16; 16]; + for k in 1..16 { + scan[k] = t[ZIGZAG4[k]] as i16; + } + c.code.luma4[blk] = if ac_any { scan } else { [0; 16] }; + let nz = if ac_any { scan.iter().filter(|&&v| v != 0).count() as u8 } else { 0 }; + c.info.nnz[by * 4 + bx] = nz; + let off = by * 4 * 16 + bx * 4; + if nz == 0 { + let dcv = (dcq[by * 4 + bx] + 32) >> 6; + for y in 0..4 { + for x in 0..4 { + let i = off + y * 16 + x; + c.ry[i] = (pred[i] as i32 + dcv).clamp(0, 255) as u8; + } + } + } else { + dequant4(&mut t, qt, qp, true); + t[0] = dcq[by * 4 + bx]; + let r = idct4(&t); + for y in 0..4 { + for x in 0..4 { + let i = off + y * 16 + x; + c.ry[i] = (pred[i] as i32 + r[y * 4 + x]).clamp(0, 255) as u8; + } + } + } + } + if ac_any { 15 } else { 0 } + } + + fn best_chroma_mode(&self, e: &MbEdges, av: Avail, src: &Src, lam: u32) -> (u8, [u8; 64], [u8; 64]) { + let mut best = (u32::MAX, 0u8, [0u8; 64], [0u8; 64]); + for mode in 0..4u8 { + let ok = match mode { + 0 => true, + 1 => av.left, + 2 => av.top, + _ => av.left && av.top && av.tl, + }; + if !ok { + continue; + } + let mut pu = [0u8; 64]; + let mut pv = [0u8; 64]; + pred_chroma(mode, &e.ctop[0], &e.cleft[0], e.ctl[0], av.top, av.left, &mut pu); + pred_chroma(mode, &e.ctop[1], &e.cleft[1], e.ctl[1], av.top, av.left, &mut pv); + let cost = satd(&src.u, 8, &pu, 8, 8, 8) + satd(&src.v, 8, &pv, 8, 8, 8) + ((lam * (1 + 2 * mode as u32)) >> 8); + if cost < best.0 { + best = (cost, mode, pu, pv); + } + } + (best.1, best.2, best.3) + } + + fn edge4(&self, e: &MbEdges, av: Avail, rec: &[u8; 256], bx: usize, by: usize) -> Edge4 { + let mut top = [0u8; 8]; + let mut left = [0u8; 4]; + let has_top = by > 0 || av.top; + let has_left = bx > 0 || av.left; + let has_tl = if bx > 0 && by > 0 { + true + } else if by == 0 && bx > 0 { + av.top + } else if bx == 0 && by > 0 { + av.left + } else { + av.tl + }; + let tr_avail = + if by == 0 { if bx < 3 { av.top } else { av.tr } } else { bx < 3 && RASTER_TO_BLK4[(by - 1) * 4 + bx + 1] < RASTER_TO_BLK4[by * 4 + bx] }; + if has_top { + for i in 0..4 { + top[i] = if by > 0 { rec[(by * 4 - 1) * 16 + bx * 4 + i] } else { e.top[bx * 4 + i] }; + } + for i in 4..8 { + top[i] = if tr_avail { if by > 0 { rec[(by * 4 - 1) * 16 + bx * 4 + i] } else { e.top[bx * 4 + i] } } else { top[3] }; + } + } + if has_left { + for i in 0..4 { + left[i] = if bx > 0 { rec[(by * 4 + i) * 16 + bx * 4 - 1] } else { e.left[by * 4 + i] }; + } + } + let tl = if !has_tl { + 0 + } else if bx > 0 && by > 0 { + rec[(by * 4 - 1) * 16 + bx * 4 - 1] + } else if by == 0 && bx > 0 { + e.top[bx * 4 - 1] + } else if bx == 0 && by > 0 { + e.left[by * 4 - 1] + } else { + e.tl + }; + Edge4 { top, left, tl, has_top, has_left, has_tl } + } + + fn edge8(&self, e: &MbEdges, av: Avail, rec: &[u8; 256], b8: usize) -> Edge8 { + let (x8, y8) = (b8 & 1, b8 >> 1); + let mut top = [0u8; 16]; + let mut left = [0u8; 8]; + let has_top = y8 > 0 || av.top; + let has_left = x8 > 0 || av.left; + let has_tl = match b8 { + 0 => av.tl, + 1 => av.top, + 2 => av.left, + _ => true, + }; + let tr_avail = match b8 { + 0 => av.top, + 1 => av.tr, + 2 => true, + _ => false, + }; + if has_top { + for i in 0..8 { + top[i] = if y8 > 0 { rec[7 * 16 + x8 * 8 + i] } else { e.top[x8 * 8 + i] }; + } + for i in 8..16 { + top[i] = if tr_avail { if y8 > 0 { rec[7 * 16 + x8 * 8 + i] } else { e.top[x8 * 8 + i] } } else { top[7] }; + } + } + if has_left { + for i in 0..8 { + left[i] = if x8 > 0 { rec[(y8 * 8 + i) * 16 + 7] } else { e.left[y8 * 8 + i] }; + } + } + let tl = if !has_tl { + 0 + } else { + match b8 { + 0 => e.tl, + 1 => e.top[7], + 2 => e.left[7], + _ => rec[7 * 16 + 7], + } + }; + Edge8 { top, left, tl, has_top, has_left, has_tl } + } + + /// Intra 4x4 analysis + coding. Returns cost (SATD domain) or None if it exceeded `limit`. + fn intra4x4(&self, mx: usize, my: usize, qp: u8, av: Avail, e: &MbEdges, src: &Src, c: &mut Cand, limit: u64) -> Option { + let qt = self.f.qt; + let lam = self.f.lam.sad[qp as usize]; + let mut modes = [-1i8; 16]; + let mut total = ((lam * 24) >> 8) as u64; + for blk in 0..16 { + let (bx, by) = BLK4_XY[blk]; + let ed = self.edge4(e, av, &c.ry, bx, by); + let pm = self.pred_intra_mode(mx, my, av, &modes, bx, by); + let mut srcb = [0u8; 16]; + for y in 0..4 { + srcb[y * 4..y * 4 + 4].copy_from_slice(&src.y[(by * 4 + y) * 16 + bx * 4..(by * 4 + y) * 16 + bx * 4 + 4]); + } + let mut best = (u32::MAX, 2u8, [0u8; 16]); + for mode in 0..9u8 { + if !mode_available(mode, ed.has_top, ed.has_left, ed.has_tl) { + continue; + } + let mut p = [0u8; 16]; + pred4x4(mode, &ed, &mut p); + let bits = if mode == pm { 1 } else { 4 }; + let cost = satd(&srcb, 4, &p, 4, 4, 4) + ((lam * bits) >> 8); + if cost < best.0 { + best = (cost, mode, p); + } + } + total += best.0 as u64; + if total > limit { + return None; + } + let (mode, p) = (best.1, best.2); + modes[by * 4 + bx] = mode as i8; + c.code.ipred[blk] = if mode == pm { + -1 + } else if mode < pm { + mode as i8 + } else { + mode as i8 - 1 + }; + // code block + let mut d = [0i32; 16]; + for i in 0..16 { + d[i] = srcb[i] as i32 - p[i] as i32; + } + let mut t = fdct4(&d); + let nz = quant4(&mut t, qt, qp, Deadzone::INTRA, false); + if self.f.p.cavlc_clamp { + for v in t.iter_mut() { + *v = (*v).clamp(-cavlc::MAX_LEVEL, cavlc::MAX_LEVEL); + } + } + let mut scan = [0i16; 16]; + for k in 0..16 { + scan[k] = t[ZIGZAG4[k]] as i16; + } + c.code.luma4[blk] = scan; + c.info.nnz[by * 4 + bx] = nz as u8; + let off = by * 4 * 16 + bx * 4; + if nz == 0 { + for y in 0..4 { + c.ry[off + y * 16..off + y * 16 + 4].copy_from_slice(&p[y * 4..y * 4 + 4]); + } + } else { + dequant4(&mut t, qt, qp, false); + let r = idct4(&t); + for y in 0..4 { + for x in 0..4 { + c.ry[off + y * 16 + x] = (p[y * 4 + x] as i32 + r[y * 4 + x]).clamp(0, 255) as u8; + } + } + } + } + c.info.imodes = modes; + let mut cbp = 0; + for b8 in 0..4 { + if (0..4).any(|i| c.code.luma4[b8 * 4 + i].iter().any(|&v| v != 0)) { + cbp |= 1 << b8; + } + } + c.code.cbp = cbp; + c.code.kind = MbKind::I4x4; + c.code.t8x8 = false; + c.info.t8x8 = false; + Some(total) + } + + /// Intra 8x8 analysis + coding. + fn intra8x8(&self, mx: usize, my: usize, qp: u8, av: Avail, e: &MbEdges, src: &Src, c: &mut Cand, limit: u64) -> Option { + let qt = self.f.qt; + let lam = self.f.lam.sad[qp as usize]; + let mut modes = [-1i8; 16]; + let mut total = ((lam * 12) >> 8) as u64; + for b8 in 0..4 { + let (x8, y8) = (b8 & 1, b8 >> 1); + let ed = self.edge8(e, av, &c.ry, b8); + let fe = filter8(&ed); + let pm = self.pred_intra_mode(mx, my, av, &modes, x8 * 2, y8 * 2); + let mut srcb = [0u8; 64]; + for y in 0..8 { + srcb[y * 8..y * 8 + 8].copy_from_slice(&src.y[(y8 * 8 + y) * 16 + x8 * 8..(y8 * 8 + y) * 16 + x8 * 8 + 8]); + } + let mut best = (u32::MAX, 2u8, [0u8; 64]); + for mode in 0..9u8 { + if !mode_available(mode, ed.has_top, ed.has_left, ed.has_tl) { + continue; + } + let mut p = [0u8; 64]; + pred8x8(mode, &fe, &mut p); + let bits = if mode == pm { 1 } else { 4 }; + let cost = satd(&srcb, 8, &p, 8, 8, 8) + ((lam * bits) >> 8); + if cost < best.0 { + best = (cost, mode, p); + } + } + total += best.0 as u64; + if total > limit { + return None; + } + let (mode, p) = (best.1, best.2); + for (dx, dy) in [(0, 0), (1, 0), (0, 1), (1, 1)] { + modes[(y8 * 2 + dy) * 4 + x8 * 2 + dx] = mode as i8; + } + c.code.ipred[b8] = if mode == pm { + -1 + } else if mode < pm { + mode as i8 + } else { + mode as i8 - 1 + }; + let mut d = [0i32; 64]; + for i in 0..64 { + d[i] = srcb[i] as i32 - p[i] as i32; + } + let mut t = fdct8(&d); + let nz = quant8(&mut t, qt, qp, Deadzone::INTRA); + let mut scan = [0i16; 64]; + for k in 0..64 { + scan[k] = t[ZIGZAG8[k]] as i16; + } + c.code.luma8[b8] = scan; + for (dx, dy) in [(0, 0), (1, 0), (0, 1), (1, 1)] { + c.info.nnz[(y8 * 2 + dy) * 4 + x8 * 2 + dx] = nz as u8; + } + let off = y8 * 8 * 16 + x8 * 8; + if nz == 0 { + for y in 0..8 { + c.ry[off + y * 16..off + y * 16 + 8].copy_from_slice(&p[y * 8..y * 8 + 8]); + } + } else { + dequant8(&mut t, qt, qp); + let r = idct8(&t); + for y in 0..8 { + for x in 0..8 { + c.ry[off + y * 16 + x] = (p[y * 8 + x] as i32 + r[y * 8 + x]).clamp(0, 255) as u8; + } + } + } + } + c.info.imodes = modes; + let mut cbp = 0; + for b8 in 0..4 { + if c.code.luma8[b8].iter().any(|&v| v != 0) { + cbp |= 1 << b8; + } + } + c.code.cbp = cbp; + c.code.kind = MbKind::I8x8; + c.code.t8x8 = true; + c.info.t8x8 = true; + Some(total) + } + + /// Full intra analysis; returns the best coded intra candidate. `inter_cost` lets P/B skip expensive modes. + fn analyse_intra(&self, mx: usize, my: usize, qp: u8, av: Avail, e: &MbEdges, src: &Src, inter_cost: u64) -> Box { + let lam = self.f.lam.sad[qp as usize]; + let p = self.f.p; + // I16x16 mode + let mut top16 = [0u8; 16]; + top16.copy_from_slice(&e.top[..16]); + let mut b16 = (u32::MAX, 2u8, [0u8; 256]); + for mode in 0..4u8 { + let ok = match mode { + 0 => av.top, + 1 => av.left, + 2 => true, + _ => av.top && av.left && av.tl, + }; + if !ok { + continue; + } + let mut pr = [0u8; 256]; + pred16x16(mode, &top16, &e.left, e.tl, av.top, av.left, &mut pr); + let cost = satd(&src.y, 16, &pr, 16, 16, 16); + if cost < b16.0 { + b16 = (cost, mode, pr); + } + } + let cost16 = b16.0 as u64 + ((lam * 6) >> 8) as u64; + let (cmode, pu, pv) = self.best_chroma_mode(e, av, src, lam); + let mk16 = || { + let mut c = Cand::new(); + c.code.kind = MbKind::I16x16; + c.code.i16_mode = b16.1; + c.code.chroma_mode = cmode; + let cbpl = self.code_i16(&src.y, &b16.2, qp, &mut c); + let cbpc = self.code_chroma(src, &pu, &pv, qp, true, &mut c); + c.code.cbp = cbpl | (cbpc << 4); + c.info.chroma_mode = cmode; + c.info.imodes = [-1; 16]; + c.cost = cost16; + c + }; + let mut best: Option> = None; + let try_nxn = p.i4x4 && (inter_cost == u64::MAX || cost16 < inter_cost + inter_cost / 4 || p.always_intra); + if !try_nxn { + return mk16(); + } + let limit = cost16.min(inter_cost.saturating_add(inter_cost / 8)); + let mut limit = if p.rd { cost16 + cost16 / 4 } else { limit }; + if p.i8x8 && p.t8x8 { + let mut c = Cand::new(); + if let Some(cost) = self.intra8x8(mx, my, qp, av, e, src, &mut c, limit) { + c.cost = cost; + limit = limit.min(cost); + best = Some(c); + } + } + { + let mut c = Cand::new(); + if let Some(cost) = self.intra4x4(mx, my, qp, av, e, src, &mut c, limit) { + c.cost = cost; + if best.as_ref().is_none_or(|b| cost < b.cost) { + best = Some(c); + } + } + } + let finish = |mut c: Box| { + c.code.chroma_mode = cmode; + c.info.chroma_mode = cmode; + let cbpc = self.code_chroma(src, &pu, &pv, qp, true, &mut c); + c.code.cbp = (c.code.cbp & 15) | (cbpc << 4); + c + }; + match best { + None => mk16(), + Some(c) => { + let c = finish(c); + if p.rd { + let c16 = mk16(); + let (j1, j2) = (self.rd_cost(mx, my, &c, src, qp), self.rd_cost(mx, my, &c16, src, qp)); + let mut w = if j1 <= j2 { c } else { c16 }; + w.cost = w.cost.min(cost16); + w + } else if c.cost < cost16 { + c + } else { + mk16() + } + } + } + } + + // ------------------------------------------------------------------ motion estimation + + /// Integer + sub-pel search for a w x h block at picture (bx, by). Returns (mv, cost, satd). + fn search(&self, r: &RefPic, bx: usize, by: usize, w: usize, h: usize, src: &[u8], sstride: usize, mvp: Mv, cands: &[Mv], lam: u32) -> (Mv, u32) { + let pl = &r.frame.y; + let p = self.f.p; + let fcost = |mv: Mv| -> u32 { + let i = pl.idx(bx as isize + (mv.x as isize >> 2), by as isize + (mv.y as isize >> 2)); + let d = if w == 16 { sad16(src, sstride, &pl.data[i..], pl.stride, h) } else { sad(src, sstride, &pl.data[i..], pl.stride, w, h) }; + d + ((lam * mv_bits(mv, mvp)) >> 8) + }; + let mut best = Mv::new((mvp.x as i32 + 2) & !3, (mvp.y as i32 + 2) & !3); + if !self.mv_ok(bx, by, w, h, best) { + best = Mv::ZERO; + } + let mut bcost = fcost(best); + for &c in cands { + let c = Mv::new((c.x as i32 + 2) & !3, (c.y as i32 + 2) & !3); + if c != best && self.mv_ok(bx, by, w, h, c) { + let k = fcost(c); + if k < bcost { + bcost = k; + best = c; + } + } + } + // hexagon search (full-pel units = 4 quarter samples) + let range = p.me_range; + let centre0 = best; + const HEX: [(i32, i32); 6] = [(-2, 0), (2, 0), (-1, -2), (1, -2), (-1, 2), (1, 2)]; + for _ in 0..range { + let mut moved = false; + let c = best; + for &(dx, dy) in &HEX { + let m = Mv::new(c.x as i32 + dx * 4, c.y as i32 + dy * 4); + if ((m.x as i32 - centre0.x as i32).abs() >> 2) > range || ((m.y as i32 - centre0.y as i32).abs() >> 2) > range || !self.mv_ok(bx, by, w, h, m) + { + continue; + } + let k = fcost(m); + if k < bcost { + bcost = k; + best = m; + moved = true; + } + } + if !moved { + break; + } + } + // small diamond / square refinement + for _ in 0..2 { + let c = best; + for (dx, dy) in [(-1, 0), (1, 0), (0, -1), (0, 1), (-1, -1), (1, -1), (-1, 1), (1, 1)] { + let m = Mv::new(c.x as i32 + dx * 4, c.y as i32 + dy * 4); + if !self.mv_ok(bx, by, w, h, m) { + continue; + } + let k = fcost(m); + if k < bcost { + bcost = k; + best = m; + } + } + if best == c { + break; + } + } + // sub-pel refinement with SATD + let mut buf = [0u8; 256]; + let scost = |mv: Mv, buf: &mut [u8; 256]| -> u32 { + r.mc_luma(bx, by, mv, w, h, buf, 16); + satd(src, sstride, buf, 16, w, h) + ((lam * mv_bits(mv, mvp)) >> 8) + }; + let mut bs = scost(best, &mut buf); + if p.subpel >= 1 { + for step in [2, 1] { + if step == 1 && p.subpel < 2 { + break; + } + let rounds = if step == 1 { 1 + p.subpel_rounds as usize } else { 1 }; + for _ in 0..rounds { + let c = best; + for (dx, dy) in [(-1, 0), (1, 0), (0, -1), (0, 1), (-1, -1), (1, -1), (-1, 1), (1, 1)] { + let m = Mv::new(c.x as i32 + dx * step, c.y as i32 + dy * step); + if !self.mv_ok(bx, by, w, h, m) { + continue; + } + let k = scost(m, &mut buf); + if k < bs { + bs = k; + best = m; + } + } + if best == c { + break; + } + } + } + } + (best, bs) + } + + fn pred_block(&self, r: &RefPic, px: usize, py: usize, rect: (usize, usize, usize, usize), mv: Mv, y: &mut [u8; 256], u: &mut [u8; 64], v: &mut [u8; 64]) { + let (x, yy, w, h) = rect; + r.mc_luma(px + x, py + yy, mv, w, h, &mut y[yy * 16 + x..], 16); + let co = (yy / 2) * 8 + x / 2; + r.mc_chroma((px + x) / 2, (py + yy) / 2, mv, w / 2, h / 2, &mut u[co..], &mut v[co..], 8); + } + + fn cand_inter_info(c: &mut Cand, part: Part, refs: [[i8; 4]; 2], mvs: [[Mv; 4]; 2], mvds: [[Mv; 4]; 2]) { + for p in 0..part.count() { + let (x, y, w, h) = part.rect(p); + for by in y / 4..(y + h) / 4 { + for bx in x / 4..(x + w) / 4 { + let r = by * 4 + bx; + for l in 0..2 { + c.info.mv[l][r] = if refs[l][p] >= 0 { mvs[l][p] } else { Mv::ZERO }; + let d = if refs[l][p] >= 0 { mvds[l][p] } else { Mv::ZERO }; + c.info.mvd[l][r] = [d.x.unsigned_abs().min(127) as u8, d.y.unsigned_abs().min(127) as u8]; + } + } + } + for by8 in y / 8..(y + h).div_ceil(8) { + for bx8 in x / 8..(x + w).div_ceil(8) { + for l in 0..2 { + c.info.ref_idx[l][by8 * 2 + bx8] = refs[l][p]; + } + } + } + } + } + + fn mask_of(part: Part, upto: usize) -> u16 { + let mut m = 0u16; + for p in 0..upto { + let (x, y, w, h) = part.rect(p); + for by in y / 4..(y + h) / 4 { + for bx in x / 4..(x + w) / 4 { + m |= 1 << (by * 4 + bx); + } + } + } + m + } + + fn shape_of(part: Part, p: usize) -> Shape { + match (part, p) { + (Part::P16x8, 0) => Shape::Top16x8, + (Part::P16x8, _) => Shape::Bottom16x8, + (Part::P8x16, 0) => Shape::Left8x16, + (Part::P8x16, _) => Shape::Right8x16, + _ => Shape::Normal, + } + } + + /// Search each partition of `part` for list `list` sequentially (mvp depends on earlier partitions). + fn search_part(&self, mx: usize, my: usize, r: &RefPic, part: Part, list: usize, src: &Src, lam: u32, seed: Mv) -> ([Mv; 4], [Mv; 4], u32) { + let (px, py) = (mx * 16, my * 16); + let mut info = MbInfo::default(); + let mut mvs = [Mv::ZERO; 4]; + let mut mvps = [Mv::ZERO; 4]; + let mut total = 0; + for p in 0..part.count() { + let (x, y, w, h) = part.rect(p); + let mask = Self::mask_of(part, p); + let mvp = self.mvp(&info, mask, mx, my, (x / 4) as i32, (y / 4) as i32, (w / 4) as i32, list, 0, Self::shape_of(part, p)); + let (mv, cost) = self.search(r, px + x, py + y, w, h, &src.y[y * 16 + x..], 16, mvp, &[seed, mvs[0]], lam); + mvs[p] = mv; + mvps[p] = mvp; + total += cost; + // record in scratch info for the next partition's predictor + for by in y / 4..(y + h) / 4 { + for bx in x / 4..(x + w) / 4 { + info.mv[list][by * 4 + bx] = mv; + } + } + for by8 in y / 8..(y + h).div_ceil(8) { + for bx8 in x / 8..(x + w).div_ceil(8) { + info.ref_idx[list][by8 * 2 + bx8] = 0; + } + } + } + (mvs, mvps, total) + } + + fn analyse_p(&self, mx: usize, my: usize, qp: u8, av: Avail, e: &MbEdges, src: &Src) -> Box { + let f = self.f; + let r = f.l0.expect("P slice without reference"); + let lam = f.lam.sad[qp as usize]; + let (px, py) = (mx * 16, my * 16); + let skip_mv = self.pskip_mv(mx, my); + let skip_ok = self.mv_ok(px, py, 16, 16, skip_mv); + // --- try P_Skip early + let mut skip_pred = None; + if skip_ok { + let mut y = [0u8; 256]; + let mut u = [0u8; 64]; + let mut v = [0u8; 64]; + self.pred_block(r, px, py, (0, 0, 16, 16), skip_mv, &mut y, &mut u, &mut v); + let s = sad16(&src.y, 16, &y, 16, 16); + let qstep_thresh = (f.qt.v4[(qp % 6) as usize][0] << (qp / 6)) as u32 * 16; + if s < qstep_thresh * 2 { + let mut c = Cand::new(); + c.code.kind = MbKind::PInter; + c.code.part = Part::P16x16; + self.code_inter(src, &y, &u, &v, qp, &mut c); + if c.code.cbp == 0 { + c.code.kind = MbKind::PSkip; + Self::cand_inter_info(&mut c, Part::P16x16, [[0; 4], [-1; 4]], [[skip_mv; 4], [Mv::ZERO; 4]], [[Mv::ZERO; 4]; 2]); + c.info.nnz = [0; 16]; + c.info.nnz_c = [0; 8]; + c.info.dc_cbf = 0; + c.info.t8x8 = false; + c.ry = y; + c.ru = u; + c.rv = v; + return c; + } + } + skip_pred = Some((y, u, v)); + } + // --- 16x16 search + let mut cands = vec![skip_mv, Mv::ZERO]; + let colo = r.mbs[my * f.p.mbw + mx]; + if colo.ref_idx[0][0] >= 0 { + cands.push(colo.mv[0][0]); + } + let (mv16, mvp16, cost16) = { + let (m, p, c) = self.search_part(mx, my, r, Part::P16x16, 0, src, lam, skip_mv); + (m, p, c) + }; + let _ = cands; + let mut best_part = Part::P16x16; + let mut best_mvs = mv16; + let mut best_mvps = mvp16; + let mut best_cost = cost16 as u64 + (lam >> 8) as u64; + if f.p.partitions { + let big = cost16 as u64 > (256 * 4) as u64; + if big { + for part in [Part::P16x8, Part::P8x16, Part::P8x8] { + let (m, p, c) = self.search_part(mx, my, r, part, 0, src, lam, mv16[0]); + let extra = match part { + Part::P8x8 => 10, + _ => 3, + }; + let c = c as u64 + ((lam * extra) >> 8) as u64; + if c < best_cost { + best_cost = c; + best_part = part; + best_mvs = m; + best_mvps = p; + } + } + } + } + // --- build inter candidate + let mut ci = Cand::new(); + ci.code.kind = MbKind::PInter; + ci.code.part = best_part; + let mut py_ = [0u8; 256]; + let mut pu = [0u8; 64]; + let mut pv = [0u8; 64]; + if let (Part::P16x16, true, Some((a, b, c))) = (best_part, best_mvs[0] == skip_mv, skip_pred) { + py_ = a; + pu = b; + pv = c; + } else { + for p in 0..best_part.count() { + self.pred_block(r, px, py, best_part.rect(p), best_mvs[p], &mut py_, &mut pu, &mut pv); + } + } + let mut mvds = [[Mv::ZERO; 4]; 2]; + for p in 0..best_part.count() { + mvds[0][p] = Mv::new(best_mvs[p].x as i32 - best_mvps[p].x as i32, best_mvs[p].y as i32 - best_mvps[p].y as i32); + } + ci.code.mvd = mvds; + Self::cand_inter_info(&mut ci, best_part, [[0; 4], [-1; 4]], [best_mvs, [Mv::ZERO; 4]], mvds); + self.code_inter(src, &py_, &pu, &pv, qp, &mut ci); + ci.cost = best_cost; + if ci.code.cbp == 0 && best_part == Part::P16x16 && best_mvs[0] == skip_mv && skip_ok { + ci.code.kind = MbKind::PSkip; + ci.info.t8x8 = false; + ci.info.mvd = [[[0; 2]; 16]; 2]; + return ci; + } + // --- intra + let need_intra = f.p.always_intra || best_cost > (256 * 3 + 16 * lam as u64 / 256) * 2; + if need_intra { + let ic = self.analyse_intra(mx, my, qp, av, e, src, best_cost); + if f.p.rd { + if self.rd_cost(mx, my, &ic, src, qp) < self.rd_cost(mx, my, &ci, src, qp) { + return ic; + } + } else if ic.cost < best_cost { + return ic; + } + } + if f.p.rd && skip_ok && ci.code.cbp != 0 { + // RD check against P_Skip + if let Some((y, u, v)) = skip_pred { + let mut sc = Cand::new(); + sc.code.kind = MbKind::PSkip; + Self::cand_inter_info(&mut sc, Part::P16x16, [[0; 4], [-1; 4]], [[skip_mv; 4], [Mv::ZERO; 4]], [[Mv::ZERO; 4]; 2]); + sc.ry = y; + sc.ru = u; + sc.rv = v; + if self.rd_cost(mx, my, &sc, src, qp) < self.rd_cost(mx, my, &ci, src, qp) { + return sc; + } + } + } + ci + } + + // ------------------------------------------------------------------ B slices + + /// Spatial direct prediction (§8.4.1.2.2) with direct_8x8_inference. Returns per-8x8 refs and mvs. + fn direct_spatial(&self, mx: usize, my: usize) -> ([[i8; 4]; 2], [[Mv; 4]; 2]) { + let cur = MbInfo::default(); + let mut refs = [-1i8; 2]; + let mut mvp = [Mv::ZERO; 2]; + for l in 0..2 { + let a = self.nb_part(&cur, 0, mx, my, -1, 0, l).map_or(-1, |v| v.0); + let b = self.nb_part(&cur, 0, mx, my, 0, -1, l).map_or(-1, |v| v.0); + let c = self.nb_part(&cur, 0, mx, my, 4, -1, l).or_else(|| self.nb_part(&cur, 0, mx, my, -1, -1, l)).map_or(-1, |v| v.0); + let minp = |x: i8, y: i8| if x >= 0 && y >= 0 { x.min(y) } else { x.max(y) }; + refs[l] = minp(a, minp(b, c)); + } + let zero = refs[0] < 0 && refs[1] < 0; + if zero { + refs = [0, 0]; + } else { + for l in 0..2 { + if refs[l] >= 0 { + mvp[l] = self.mvp(&cur, 0, mx, my, 0, 0, 4, l, refs[l], Shape::Normal); + } + } + } + let col = &self.f.l1.unwrap().mbs[my * self.f.p.mbw + mx]; + let mut out_ref = [[-1i8; 4]; 2]; + let mut out_mv = [[Mv::ZERO; 4]; 2]; + for q in 0..4 { + let corner = [0usize, 3, 12, 15][q]; + let col_zero = if col.kind.is_intra() { + false + } else { + let (rc, mvc) = if col.ref_idx[0][q] >= 0 { (col.ref_idx[0][q], col.mv[0][corner]) } else { (col.ref_idx[1][q], col.mv[1][corner]) }; + rc == 0 && mvc.x.abs() <= 1 && mvc.y.abs() <= 1 + }; + for l in 0..2 { + out_ref[l][q] = refs[l]; + if refs[l] >= 0 { + out_mv[l][q] = if zero || (refs[l] == 0 && col_zero) { Mv::ZERO } else { mvp[l] }; + } + } + } + (out_ref, out_mv) + } + + /// B 16x8 / 8x16 candidate: search each partition in both lists, choose L0/L1/Bi per partition, then derive + /// the motion vector predictors in decoding order from the chosen directions. + fn b_partition(&self, mx: usize, my: usize, part: Part, src: &Src, qp: u8, lam: u32, seeds: [Mv; 2]) -> Option> { + let f = self.f; + let rp = [f.l0?, f.l1?]; + let (px, py) = (mx * 16, my * 16); + let mut mvs = [[Mv::ZERO; 4]; 2]; + for l in 0..2 { + mvs[l] = self.search_part(mx, my, rp[l], part, l, src, lam, seeds[l]).0; + } + let mut preds = [([0u8; 256], [0u8; 64], [0u8; 64]), ([0u8; 256], [0u8; 64], [0u8; 64])]; + for l in 0..2 { + let (a, b, c) = &mut preds[l]; + for pi in 0..2 { + self.pred_block(rp[l], px, py, part.rect(pi), mvs[l][pi], a, b, c); + } + } + let mut bi = preds[0]; + avg_into(&mut bi.0, &preds[1].0); + avg_into(&mut bi.1, &preds[1].1); + avg_into(&mut bi.2, &preds[1].2); + let srcs = [&preds[0], &preds[1], &bi]; + let (mut py_, mut pu, mut pv) = ([0u8; 256], [0u8; 64], [0u8; 64]); + let mut dirs = [0u8; 4]; + let mut cost = ((lam * 8) >> 8) as u64; + for pi in 0..2 { + let (x, y, w, h) = part.rect(pi); + let mut bestd = (u64::MAX, 0usize); + for (d, s) in srcs.iter().enumerate() { + let bits = if d == 2 { 14 } else { 7 }; + let c = satd(&src.y[y * 16 + x..], 16, &s.0[y * 16 + x..], 16, w, h) as u64 + ((lam * bits) >> 8) as u64; + if c < bestd.0 { + bestd = (c, d); + } + } + cost += bestd.0; + dirs[pi] = bestd.1 as u8; + let s = srcs[bestd.1]; + for r in y..y + h { + py_[r * 16 + x..r * 16 + x + w].copy_from_slice(&s.0[r * 16 + x..r * 16 + x + w]); + } + for r in y / 2..(y + h) / 2 { + let o = r * 8 + x / 2; + pu[o..o + w / 2].copy_from_slice(&s.1[o..o + w / 2]); + pv[o..o + w / 2].copy_from_slice(&s.2[o..o + w / 2]); + } + } + // predictors in decoding order with the actual per-list references + let mut info = MbInfo::default(); + let mut refs = [[-1i8; 4]; 2]; + let mut fmvs = [[Mv::ZERO; 4]; 2]; + let mut mvds = [[Mv::ZERO; 4]; 2]; + for pi in 0..2 { + let (x, y, w, h) = part.rect(pi); + let mask = Self::mask_of(part, pi); + for l in 0..2 { + if dirs[pi] == 2 || dirs[pi] as usize == l { + let mvp = self.mvp(&info, mask, mx, my, (x / 4) as i32, (y / 4) as i32, (w / 4) as i32, l, 0, Self::shape_of(part, pi)); + let mv = mvs[l][pi]; + mvds[l][pi] = Mv::new(mv.x as i32 - mvp.x as i32, mv.y as i32 - mvp.y as i32); + refs[l][pi] = 0; + fmvs[l][pi] = mv; + } + } + for l in 0..2 { + for by in y / 4..(y + h) / 4 { + for bx in x / 4..(x + w) / 4 { + info.mv[l][by * 4 + bx] = fmvs[l][pi]; + } + } + for by8 in y / 8..(y + h).div_ceil(8) { + for bx8 in x / 8..(x + w).div_ceil(8) { + info.ref_idx[l][by8 * 2 + bx8] = refs[l][pi]; + } + } + } + } + let mut c = Cand::new(); + c.code.kind = MbKind::BInter; + c.code.part = part; + c.code.bdir = dirs; + c.code.mvd = mvds; + Self::cand_inter_info(&mut c, part, refs, fmvs, mvds); + self.code_inter(src, &py_, &pu, &pv, qp, &mut c); + c.cost = cost; + Some(c) + } + + fn analyse_b(&self, mx: usize, my: usize, qp: u8, av: Avail, e: &MbEdges, src: &Src) -> Box { + let f = self.f; + let (r0, r1) = (f.l0.unwrap(), f.l1.unwrap()); + let lam = f.lam.sad[qp as usize]; + let (px, py) = (mx * 16, my * 16); + // --- direct + let (dref, dmv) = self.direct_spatial(mx, my); + let mut direct_ok = true; + for q in 0..4 { + let rect = Part::P8x8.rect(q); + for l in 0..2 { + if dref[l][q] >= 0 && !self.mv_ok(px + rect.0, py + rect.1, 8, 8, dmv[l][q]) { + direct_ok = false; + } + } + } + let mut best: Option> = None; + if direct_ok { + let mut y = [0u8; 256]; + let mut u = [0u8; 64]; + let mut v = [0u8; 64]; + let mut ty = [0u8; 256]; + let mut tu = [0u8; 64]; + let mut tv = [0u8; 64]; + for q in 0..4 { + let rect = Part::P8x8.rect(q); + let (u0, u1) = (dref[0][q] >= 0, dref[1][q] >= 0); + if u0 { + self.pred_block(r0, px, py, rect, dmv[0][q], &mut y, &mut u, &mut v); + } + if u1 { + if u0 { + self.pred_block(r1, px, py, rect, dmv[1][q], &mut ty, &mut tu, &mut tv); + let (x, yy, w, h) = rect; + for rr in 0..h { + avg_into(&mut y[(yy + rr) * 16 + x..(yy + rr) * 16 + x + w], &ty[(yy + rr) * 16 + x..(yy + rr) * 16 + x + w]); + } + for rr in 0..h / 2 { + let o = (yy / 2 + rr) * 8 + x / 2; + avg_into(&mut u[o..o + w / 2], &tu[o..o + w / 2]); + avg_into(&mut v[o..o + w / 2], &tv[o..o + w / 2]); + } + } else { + self.pred_block(r1, px, py, rect, dmv[1][q], &mut y, &mut u, &mut v); + } + } + } + let mut c = Cand::new(); + c.code.kind = MbKind::BDirect; + c.code.part = Part::P8x8; + // store per-8x8 refs/mvs + let refs = [[dref[0][0], dref[0][1], dref[0][2], dref[0][3]], [dref[1][0], dref[1][1], dref[1][2], dref[1][3]]]; + Self::cand_inter_info(&mut c, Part::P8x8, refs, dmv, [[Mv::ZERO; 4]; 2]); + c.code.part = Part::P16x16; + c.cost = satd(&src.y, 16, &y, 16, 16, 16) as u64 + ((lam * 2) >> 8) as u64; + self.code_inter(src, &y, &u, &v, qp, &mut c); + if c.code.cbp == 0 { + c.code.kind = MbKind::BSkip; + c.info.t8x8 = false; + // a skip is cheap; accept immediately when the prediction is good + let thresh = (f.qt.v4[(qp % 6) as usize][0] << (qp / 6)) as u64 * 24; + if sad16(&src.y, 16, &y, 16, 16) as u64 <= thresh { + return c; + } + } + best = Some(c); + } + // --- L0 / L1 16x16 + let mut res = [(Mv::ZERO, Mv::ZERO, u32::MAX); 2]; + for (l, rp) in [(0usize, r0), (1usize, r1)] { + let (m, p, c) = self.search_part(mx, my, rp, Part::P16x16, l, src, lam, if l == 0 { dmv[0][0] } else { dmv[1][0] }); + res[l] = (m[0], p[0], c); + } + let mut preds = [([0u8; 256], [0u8; 64], [0u8; 64]), ([0u8; 256], [0u8; 64], [0u8; 64])]; + for (l, rp) in [(0usize, r0), (1usize, r1)] { + let (a, b, c) = &mut preds[l]; + self.pred_block(rp, px, py, (0, 0, 16, 16), res[l].0, a, b, c); + } + let mut bi = preds[0]; + avg_into(&mut bi.0, &preds[1].0); + avg_into(&mut bi.1, &preds[1].1); + avg_into(&mut bi.2, &preds[1].2); + let bits = |l: usize| mv_bits(res[l].0, res[l].1); + let costs = [ + res[0].2 as u64 + ((lam * 3) >> 8) as u64, + res[1].2 as u64 + ((lam * 3) >> 8) as u64, + satd(&src.y, 16, &bi.0, 16, 16, 16) as u64 + ((lam * (bits(0) + bits(1) + 5)) >> 8) as u64, + ]; + let dir = (0..3).min_by_key(|&d| costs[d]).unwrap(); + let mut ci = Cand::new(); + ci.code.kind = MbKind::BInter; + ci.code.part = Part::P16x16; + ci.code.bdir[0] = dir as u8; + let mut refs = [[-1i8; 4]; 2]; + let mut mvs = [[Mv::ZERO; 4]; 2]; + let mut mvds = [[Mv::ZERO; 4]; 2]; + for l in 0..2 { + if dir == 2 || dir == l { + refs[l] = [0; 4]; + mvs[l] = [res[l].0; 4]; + mvds[l][0] = Mv::new(res[l].0.x as i32 - res[l].1.x as i32, res[l].0.y as i32 - res[l].1.y as i32); + } + } + ci.code.mvd = mvds; + Self::cand_inter_info(&mut ci, Part::P16x16, refs, mvs, mvds); + let (py_, pu, pv) = if dir == 2 { bi } else { preds[dir] }; + ci.cost = costs[dir]; + let better_than_direct = best.as_ref().is_none_or(|b| ci.cost < b.cost); + if better_than_direct || f.p.rd { + self.code_inter(src, &py_, &pu, &pv, qp, &mut ci); + best = Some(match best { + None => ci, + Some(b) => { + if f.p.rd { + if self.rd_cost(mx, my, &ci, src, qp) < self.rd_cost(mx, my, &b, src, qp) { ci } else { b } + } else { + ci + } + } + }); + } + let mut best = best.unwrap(); + // --- 16x8 / 8x16 partitions with a per-partition prediction direction + if f.p.partitions && best.cost > 256 * 4 { + for part in [Part::P16x8, Part::P8x16] { + if let Some(c) = self.b_partition(mx, my, part, src, qp, lam, [res[0].0, res[1].0]) { + let better = if f.p.rd { self.rd_cost(mx, my, &c, src, qp) < self.rd_cost(mx, my, &best, src, qp) } else { c.cost < best.cost }; + if better { + best = c; + } + } + } + } + let need_intra = f.p.always_intra || best.cost > (256 * 3 + 16 * lam as u64 / 256) * 3; + if need_intra { + let ic = self.analyse_intra(mx, my, qp, av, e, src, best.cost); + if f.p.rd { + if self.rd_cost(mx, my, &ic, src, qp) < self.rd_cost(mx, my, &best, src, qp) { + return ic; + } + } else if ic.cost < best.cost { + return ic; + } + } + best + } +} + +struct Src { + y: [u8; 256], + u: [u8; 64], + v: [u8; 64], +} diff --git a/crates/h264enc/src/syntax.rs b/crates/h264enc/src/syntax.rs new file mode 100644 index 00000000..655b1492 --- /dev/null +++ b/crates/h264enc/src/syntax.rs @@ -0,0 +1,120 @@ +//! Macroblock-layer syntax values produced by mode decision and consumed by the entropy coders. + +use crate::mbinfo::{MbInfo, MbKind, Mv, Part}; +use crate::nal::SliceType; + +#[derive(Clone)] +pub struct MbCode { + pub kind: MbKind, + pub part: Part, + /// Prediction direction per partition: 0 = L0, 1 = L1, 2 = Bi. + pub bdir: [u8; 4], + pub i16_mode: u8, + pub chroma_mode: u8, + /// Per 4x4 (or 8x8, first four entries) block in decoding order: -1 = use predicted mode, else rem_intra_pred_mode. + pub ipred: [i8; 16], + pub t8x8: bool, + /// mvd per list per partition. + pub mvd: [[Mv; 4]; 2], + pub cbp: u8, + pub qp_delta: i32, + /// Intra16x16 DC levels in zig-zag order. + pub luma_dc: [i16; 16], + /// Luma 4x4 levels by luma4x4BlkIdx, zig-zag order (index 0 unused for Intra16x16 AC). + pub luma4: [[i16; 16]; 16], + /// Luma 8x8 levels by 8x8 block, zig-zag order. + pub luma8: [[i16; 64]; 4], + /// Chroma DC levels (Cb, Cr) in raster (c0..c3) order. + pub cdc: [[i16; 4]; 2], + /// Chroma AC levels, Cb blocks 0..3 then Cr 0..3, zig-zag order (index 0 unused). + pub cac: [[i16; 16]; 8], +} + +impl Default for MbCode { + fn default() -> Self { + MbCode { + kind: MbKind::None, + part: Part::P16x16, + bdir: [0; 4], + i16_mode: 0, + chroma_mode: 0, + ipred: [-1; 16], + t8x8: false, + mvd: [[Mv::ZERO; 4]; 2], + cbp: 0, + qp_delta: 0, + luma_dc: [0; 16], + luma4: [[0; 16]; 16], + luma8: [[0; 64]; 4], + cdc: [[0; 4]; 2], + cac: [[0; 16]; 8], + } + } +} + +impl MbCode { + /// B macroblock type number (Table 7-14) for inter B macroblocks. + pub fn b_type_number(&self) -> u32 { + match self.kind { + MbKind::BDirect => 0, + _ => match self.part { + Part::P16x16 => 1 + self.bdir[0] as u32, + Part::P16x8 | Part::P8x16 => { + const PAIRS: [(u8, u8); 9] = [(0, 0), (1, 1), (0, 1), (1, 0), (0, 2), (1, 2), (2, 0), (2, 1), (2, 2)]; + let k = PAIRS.iter().position(|&p| p == (self.bdir[0], self.bdir[1])).unwrap() as u32; + 4 + 2 * k + (self.part == Part::P8x16) as u32 + } + Part::P8x8 => 22, + }, + } + } + pub fn uses_list(&self, part: usize, list: usize) -> bool { + match self.kind { + MbKind::PInter => list == 0, + MbKind::BInter => self.bdir[part] == 2 || self.bdir[part] as usize == list, + _ => false, + } + } +} + +/// Neighbour context for entropy coding one macroblock. +pub struct NbCtx<'a> { + pub left: Option<&'a MbInfo>, + pub top: Option<&'a MbInfo>, + pub cur: &'a MbInfo, + pub slice_type: SliceType, + pub prev_dqp_nz: bool, + pub t8x8_mode: bool, +} + +impl NbCtx<'_> { + /// coded_block_flag-style condition for an unavailable neighbour. + #[inline] + pub fn unavail_cbf(&self) -> u32 { + self.cur.kind.is_intra() as u32 + } + /// Luma 4x4 neighbour non-zero count (A = left when `left`, else B = above). Returns None if unavailable. + #[inline] + pub fn luma_nnz(&self, raster: usize, left: bool) -> Option { + let (x, y) = (raster & 3, raster >> 2); + if left { + if x > 0 { Some(self.cur.nnz[raster - 1]) } else { self.left.map(|m| m.nnz[y * 4 + 3]) } + } else if y > 0 { + Some(self.cur.nnz[raster - 4]) + } else { + self.top.map(|m| m.nnz[12 + x]) + } + } + /// Chroma AC neighbour count for component `c`, block raster `b` (2x2). + #[inline] + pub fn chroma_nnz(&self, c: usize, b: usize, left: bool) -> Option { + let (x, y) = (b & 1, b >> 1); + if left { + if x > 0 { Some(self.cur.nnz_c[c * 4 + b - 1]) } else { self.left.map(|m| m.nnz_c[c * 4 + y * 2 + 1]) } + } else if y > 0 { + Some(self.cur.nnz_c[c * 4 + b - 2]) + } else { + self.top.map(|m| m.nnz_c[c * 4 + 2 + x]) + } + } +} diff --git a/crates/h264enc/src/tables.rs b/crates/h264enc/src/tables.rs new file mode 100644 index 00000000..95ee964c --- /dev/null +++ b/crates/h264enc/src/tables.rs @@ -0,0 +1,1095 @@ +//! Normative constant tables from ITU-T H.264 (scans, CABAC state machine and initialisation values, +//! CAVLC code tables, quantisation/deblocking tables). + +/// 4x4 frame zig-zag scan: scan index -> raster index (y*4+x). +pub const ZIGZAG4: [usize; 16] = [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15]; + +/// 8x8 frame zig-zag scan: scan index -> raster index (y*8+x). +pub const ZIGZAG8: [usize; 64] = [ + 0, 1, 8, 16, 9, 2, 3, 10, 17, 24, 32, 25, 18, 11, 4, 5, 12, 19, 26, 33, 40, 48, 41, 34, 27, 20, 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, + 22, 15, 23, 30, 37, 44, 51, 58, 59, 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63, +]; + +/// Luma 4x4 block index (decoding order) -> (x, y) in 4x4 units within the macroblock. +pub const BLK4_XY: [(usize, usize); 16] = + [(0, 0), (1, 0), (0, 1), (1, 1), (2, 0), (3, 0), (2, 1), (3, 1), (0, 2), (1, 2), (0, 3), (1, 3), (2, 2), (3, 2), (2, 3), (3, 3)]; + +/// Raster 4x4 position (y*4+x) -> luma 4x4 block index. +pub const RASTER_TO_BLK4: [usize; 16] = [0, 1, 4, 5, 2, 3, 6, 7, 8, 9, 12, 13, 10, 11, 14, 15]; + +/// rangeTabLPS[pStateIdx][qCodIRangeIdx]. +pub const RANGE_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], +]; + +/// transIdxLPS. +pub const TRANS_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, +]; + +/// Number of CABAC contexts we keep (0..=459 covers every frame-coded 4:2:0 syntax element). +pub const NUM_CTX: usize = 460; + +/// Context initialisation (m, n) for I slices. Field-coding contexts are left as zero (never used). +pub fn cabac_init_i() -> [(i8, i8); NUM_CTX] { + let mut t = [(0i8, 0i8); NUM_CTX]; + set(&mut t, 0, &COMMON_0_10); + set(&mut t, 60, &COMMON_60_69); + set(&mut t, 70, &I_70_104); + set(&mut t, 105, &I_105_165); + set(&mut t, 166, &I_166_226); + set(&mut t, 227, &I_227_275); + set(&mut t, 399, &I_399_401); + set(&mut t, 402, &I_402_435); + t +} + +/// Context initialisation (m, n) for P/B slices with cabac_init_idc = 0. +pub fn cabac_init_pb0() -> [(i8, i8); NUM_CTX] { + let mut t = [(0i8, 0i8); NUM_CTX]; + set(&mut t, 0, &COMMON_0_10); + set(&mut t, 11, &PB0_11_59); + set(&mut t, 60, &COMMON_60_69); + set(&mut t, 70, &PB0_70_104); + set(&mut t, 105, &PB0_105_165); + set(&mut t, 166, &PB0_166_226); + set(&mut t, 227, &PB0_227_275); + set(&mut t, 399, &PB0_399_401); + set(&mut t, 402, &PB0_402_435); + t +} + +fn set(t: &mut [(i8, i8); NUM_CTX], start: usize, v: &[(i8, i8)]) { + t[start..start + v.len()].copy_from_slice(v); +} + +const COMMON_0_10: [(i8, i8); 11] = [(20, -15), (2, 54), (3, 74), (20, -15), (2, 54), (3, 74), (-28, 127), (-23, 104), (-6, 53), (-1, 54), (7, 51)]; + +const COMMON_60_69: [(i8, i8); 10] = [(0, 41), (0, 63), (0, 63), (0, 63), (-9, 83), (4, 86), (0, 97), (-7, 72), (13, 41), (3, 62)]; + +const PB0_11_59: [(i8, i8); 49] = [ + // 11..23 P mb_skip / mb_type / sub_mb_type + (23, 33), + (23, 2), + (21, 0), + (1, 9), + (0, 49), + (-37, 118), + (5, 57), + (-13, 78), + (-11, 65), + (1, 62), + (12, 49), + (-4, 73), + (17, 50), + // 24..39 B mb_skip / mb_type / sub_mb_type + (18, 64), + (9, 43), + (29, 0), + (26, 67), + (16, 90), + (9, 104), + (-46, 127), + (-20, 104), + (1, 67), + (-13, 78), + (-11, 65), + (1, 62), + (-6, 86), + (-17, 95), + (-6, 61), + (9, 45), + // 40..53 mvd + (-3, 69), + (-6, 81), + (-11, 96), + (6, 55), + (7, 67), + (-5, 86), + (2, 88), + (0, 58), + (-3, 76), + (-10, 94), + (5, 54), + (4, 69), + (-3, 81), + (0, 88), + // 54..59 ref_idx + (-7, 67), + (-5, 74), + (-4, 74), + (-5, 80), + (-7, 72), + (1, 58), +]; + +const I_70_104: [(i8, i8); 35] = [ + (0, 11), + (1, 55), + (0, 69), + (-17, 127), + (-13, 102), + (0, 82), + (-7, 74), + (-21, 107), + (-27, 127), + (-31, 127), + (-24, 127), + (-18, 95), + (-27, 127), + (-21, 114), + (-30, 127), + (-17, 123), + (-12, 115), + (-16, 122), + (-11, 115), + (-12, 63), + (-2, 68), + (-15, 84), + (-13, 104), + (-3, 70), + (-8, 93), + (-10, 90), + (-30, 127), + (-1, 74), + (-6, 97), + (-7, 91), + (-20, 127), + (-4, 56), + (-5, 82), + (-7, 76), + (-22, 125), +]; + +const PB0_70_104: [(i8, i8); 35] = [ + (0, 45), + (-4, 78), + (-3, 96), + (-27, 126), + (-28, 98), + (-25, 101), + (-23, 67), + (-28, 82), + (-20, 94), + (-16, 83), + (-22, 110), + (-21, 91), + (-18, 102), + (-13, 93), + (-29, 127), + (-7, 92), + (-5, 89), + (-7, 96), + (-13, 108), + (-3, 46), + (-1, 65), + (-1, 57), + (-9, 93), + (-3, 74), + (-9, 92), + (-8, 87), + (-23, 126), + (5, 54), + (6, 60), + (6, 59), + (6, 69), + (-1, 48), + (0, 68), + (-4, 69), + (-8, 88), +]; + +const I_105_165: [(i8, i8); 61] = [ + (-7, 93), + (-11, 87), + (-3, 77), + (-5, 71), + (-4, 63), + (-4, 68), + (-12, 84), + (-7, 62), + (-7, 65), + (8, 61), + (5, 56), + (-2, 66), + (1, 64), + (0, 61), + (-2, 78), + (1, 50), + (7, 52), + (10, 35), + (0, 44), + (11, 38), + (1, 45), + (0, 46), + (5, 44), + (31, 17), + (1, 51), + (7, 50), + (28, 19), + (16, 33), + (14, 62), + (-13, 108), + (-15, 100), + (-13, 101), + (-13, 91), + (-12, 94), + (-10, 88), + (-16, 84), + (-10, 86), + (-7, 83), + (-13, 87), + (-19, 94), + (1, 70), + (0, 72), + (-5, 74), + (18, 59), + (-8, 102), + (-15, 100), + (0, 95), + (-4, 75), + (2, 72), + (-11, 75), + (-3, 71), + (15, 46), + (-13, 69), + (0, 62), + (0, 65), + (21, 37), + (-15, 72), + (9, 57), + (16, 54), + (0, 62), + (12, 72), +]; + +const PB0_105_165: [(i8, i8); 61] = [ + (-2, 85), + (-6, 78), + (-1, 75), + (-7, 77), + (2, 54), + (5, 50), + (-3, 68), + (1, 50), + (6, 42), + (-4, 81), + (1, 63), + (-4, 70), + (0, 67), + (2, 57), + (-2, 76), + (11, 35), + (4, 64), + (1, 61), + (11, 35), + (18, 25), + (12, 24), + (13, 29), + (13, 36), + (-10, 93), + (-7, 73), + (-2, 73), + (13, 46), + (9, 49), + (-7, 100), + (9, 53), + (2, 53), + (5, 53), + (-2, 61), + (0, 56), + (0, 56), + (-13, 63), + (-5, 60), + (-1, 62), + (4, 57), + (-6, 69), + (4, 57), + (14, 39), + (4, 51), + (13, 68), + (3, 64), + (1, 61), + (9, 63), + (7, 50), + (16, 39), + (5, 44), + (4, 52), + (11, 48), + (-5, 60), + (-1, 59), + (0, 59), + (22, 33), + (5, 44), + (14, 43), + (-1, 78), + (0, 60), + (9, 69), +]; + +const I_166_226: [(i8, i8); 61] = [ + (24, 0), + (15, 9), + (8, 25), + (13, 18), + (15, 9), + (13, 19), + (10, 37), + (12, 18), + (6, 29), + (20, 33), + (15, 30), + (4, 45), + (1, 58), + (0, 62), + (7, 61), + (12, 38), + (11, 45), + (15, 39), + (11, 42), + (13, 44), + (16, 45), + (12, 41), + (10, 49), + (30, 34), + (18, 42), + (10, 55), + (17, 51), + (17, 46), + (0, 89), + (26, -19), + (22, -17), + (26, -17), + (30, -25), + (28, -20), + (33, -23), + (37, -27), + (33, -23), + (40, -28), + (38, -17), + (33, -11), + (40, -15), + (41, -6), + (38, 1), + (41, 17), + (30, -6), + (27, 3), + (26, 22), + (37, -16), + (35, -4), + (38, -8), + (38, -3), + (37, 3), + (38, 5), + (42, 0), + (35, 16), + (39, 22), + (14, 48), + (27, 37), + (21, 60), + (12, 68), + (2, 97), +]; + +const PB0_166_226: [(i8, i8); 61] = [ + (11, 28), + (2, 40), + (3, 44), + (0, 49), + (0, 46), + (2, 44), + (2, 51), + (0, 47), + (4, 39), + (2, 62), + (6, 46), + (0, 54), + (3, 54), + (2, 58), + (4, 63), + (6, 51), + (6, 57), + (7, 53), + (6, 52), + (6, 55), + (11, 45), + (14, 36), + (8, 53), + (-1, 82), + (7, 55), + (-3, 78), + (15, 46), + (22, 31), + (-1, 84), + (25, 7), + (30, -7), + (28, 3), + (28, 4), + (32, 0), + (34, -1), + (30, 6), + (30, 6), + (32, 9), + (31, 19), + (26, 27), + (26, 30), + (37, 20), + (28, 34), + (17, 70), + (1, 67), + (5, 59), + (9, 67), + (16, 30), + (18, 32), + (18, 35), + (22, 29), + (24, 31), + (23, 38), + (18, 43), + (20, 41), + (11, 63), + (9, 59), + (9, 64), + (-1, 94), + (-2, 89), + (-9, 108), +]; + +const I_227_275: [(i8, i8); 49] = [ + (-3, 71), + (-6, 42), + (-5, 50), + (-3, 54), + (-2, 62), + (0, 58), + (1, 63), + (-2, 72), + (-1, 74), + (-9, 91), + (-5, 67), + (-5, 27), + (-3, 39), + (-2, 44), + (0, 46), + (-16, 64), + (-8, 68), + (-10, 78), + (-6, 77), + (-10, 86), + (-12, 92), + (-15, 55), + (-10, 60), + (-6, 62), + (-4, 65), + (-12, 73), + (-8, 76), + (-7, 80), + (-9, 88), + (-17, 110), + (-11, 97), + (-20, 84), + (-11, 79), + (-6, 73), + (-4, 74), + (-13, 86), + (-13, 96), + (-11, 97), + (-19, 117), + (-8, 78), + (-5, 33), + (-4, 48), + (-2, 53), + (-3, 62), + (-13, 71), + (-10, 79), + (-12, 86), + (-13, 90), + (-14, 97), +]; + +const PB0_227_275: [(i8, i8); 49] = [ + (-6, 76), + (-2, 44), + (0, 45), + (0, 52), + (-3, 64), + (-2, 59), + (-4, 70), + (-4, 75), + (-8, 82), + (-17, 102), + (-9, 77), + (3, 24), + (0, 42), + (0, 48), + (0, 55), + (-6, 59), + (-7, 71), + (-12, 83), + (-11, 87), + (-30, 119), + (1, 58), + (-3, 29), + (-1, 36), + (1, 38), + (2, 43), + (-6, 55), + (0, 58), + (0, 64), + (-3, 74), + (-10, 90), + (0, 70), + (-4, 29), + (5, 31), + (7, 42), + (1, 59), + (-2, 58), + (-3, 72), + (-3, 81), + (-11, 97), + (0, 58), + (8, 5), + (10, 14), + (14, 18), + (13, 27), + (2, 40), + (0, 58), + (-3, 70), + (-6, 79), + (-8, 85), +]; + +const I_399_401: [(i8, i8); 3] = [(31, 21), (31, 31), (25, 50)]; +const PB0_399_401: [(i8, i8); 3] = [(12, 40), (11, 51), (14, 59)]; + +const I_402_435: [(i8, i8); 34] = [ + (-17, 120), + (-20, 112), + (-18, 114), + (-11, 85), + (-15, 92), + (-14, 89), + (-26, 71), + (-15, 81), + (-14, 80), + (0, 68), + (-14, 70), + (-24, 56), + (-23, 68), + (-24, 50), + (-11, 74), + (23, -13), + (26, -13), + (40, -15), + (49, -14), + (44, 3), + (45, 6), + (44, 34), + (33, 54), + (19, 82), + (-3, 75), + (-1, 23), + (1, 34), + (1, 43), + (0, 54), + (-2, 55), + (0, 61), + (1, 64), + (0, 68), + (-9, 92), +]; + +const PB0_402_435: [(i8, i8); 34] = [ + (-4, 79), + (-7, 71), + (-5, 69), + (-9, 70), + (-8, 66), + (-10, 68), + (-19, 73), + (-12, 69), + (-16, 70), + (-15, 67), + (-20, 62), + (-19, 70), + (-16, 66), + (-22, 65), + (-20, 63), + (9, -2), + (26, -9), + (33, -9), + (39, -7), + (41, -2), + (45, 3), + (49, 9), + (45, 27), + (36, 59), + (-6, 66), + (-7, 35), + (-7, 42), + (-8, 45), + (-5, 48), + (-12, 56), + (-6, 60), + (-5, 62), + (-8, 66), + (-8, 76), +]; + +/// significant_coeff_flag ctxIdxInc for 8x8 frame-coded blocks (Table 9-43). +pub const SIG8X8_FRAME: [u8; 63] = [ + 0, 1, 2, 3, 4, 5, 5, 4, 4, 3, 3, 4, 4, 4, 5, 5, 4, 4, 4, 4, 3, 3, 6, 7, 7, 7, 8, 9, 10, 9, 8, 7, 7, 6, 11, 12, 13, 11, 6, 7, 8, 9, 14, 10, 9, 8, 6, 11, 12, + 13, 11, 6, 9, 14, 10, 9, 11, 12, 13, 11, 14, 10, 12, +]; + +/// last_significant_coeff_flag ctxIdxInc for 8x8 blocks (Table 9-43). +pub const LAST8X8: [u8; 63] = [ + 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, + 6, 6, 6, 6, 7, 7, 7, 7, 8, 8, 8, +]; + +// ---------------------------------------------------------------- CAVLC + +/// coeff_token (length, code) indexed [table][total_coeff][trailing_ones] for nC ranges 0..2, 2..4, 4..8. +pub const COEFF_TOKEN_LEN: [[[u8; 4]; 17]; 3] = [ + [ + [1, 0, 0, 0], + [6, 2, 0, 0], + [8, 6, 3, 0], + [9, 8, 7, 5], + [10, 9, 8, 6], + [11, 10, 9, 7], + [13, 11, 10, 8], + [13, 13, 11, 9], + [13, 13, 13, 10], + [14, 14, 13, 11], + [14, 14, 14, 13], + [15, 15, 14, 14], + [15, 15, 15, 14], + [16, 15, 15, 15], + [16, 16, 16, 15], + [16, 16, 16, 16], + [16, 16, 16, 16], + ], + [ + [2, 0, 0, 0], + [6, 2, 0, 0], + [6, 5, 3, 0], + [7, 6, 6, 4], + [8, 6, 6, 4], + [8, 7, 7, 5], + [9, 8, 8, 6], + [11, 9, 9, 6], + [11, 11, 11, 7], + [12, 11, 11, 9], + [12, 12, 12, 11], + [12, 12, 12, 11], + [13, 13, 13, 12], + [13, 13, 13, 13], + [13, 14, 13, 13], + [14, 14, 14, 13], + [14, 14, 14, 14], + ], + [ + [4, 0, 0, 0], + [6, 4, 0, 0], + [6, 5, 4, 0], + [6, 5, 5, 4], + [7, 5, 5, 4], + [7, 5, 5, 4], + [7, 6, 6, 4], + [7, 6, 6, 4], + [8, 7, 7, 5], + [8, 8, 7, 6], + [9, 8, 8, 7], + [9, 9, 8, 8], + [9, 9, 9, 8], + [10, 9, 9, 9], + [10, 10, 10, 10], + [10, 10, 10, 10], + [10, 10, 10, 10], + ], +]; + +pub const COEFF_TOKEN_CODE: [[[u8; 4]; 17]; 3] = [ + [ + [1, 0, 0, 0], + [5, 1, 0, 0], + [7, 4, 1, 0], + [7, 6, 5, 3], + [7, 6, 5, 3], + [7, 6, 5, 4], + [15, 6, 5, 4], + [11, 14, 5, 4], + [8, 10, 13, 4], + [15, 14, 9, 4], + [11, 10, 13, 12], + [15, 14, 9, 12], + [11, 10, 13, 8], + [15, 1, 9, 12], + [11, 14, 13, 8], + [7, 10, 9, 12], + [4, 6, 5, 8], + ], + [ + [3, 0, 0, 0], + [11, 2, 0, 0], + [7, 7, 3, 0], + [7, 10, 9, 5], + [7, 6, 5, 4], + [4, 6, 5, 6], + [7, 6, 5, 8], + [15, 6, 5, 4], + [11, 14, 13, 4], + [15, 10, 9, 4], + [11, 14, 13, 12], + [8, 10, 9, 8], + [15, 14, 13, 12], + [11, 10, 9, 12], + [7, 11, 6, 8], + [9, 8, 10, 1], + [7, 6, 5, 4], + ], + [ + [15, 0, 0, 0], + [15, 14, 0, 0], + [11, 15, 13, 0], + [8, 12, 14, 12], + [15, 10, 11, 11], + [11, 8, 9, 10], + [9, 14, 13, 9], + [8, 10, 9, 8], + [15, 14, 13, 13], + [11, 14, 10, 12], + [15, 10, 13, 12], + [11, 14, 9, 12], + [8, 10, 13, 8], + [13, 7, 9, 12], + [9, 12, 11, 10], + [5, 8, 7, 6], + [1, 4, 3, 2], + ], +]; + +/// Chroma DC (nC = -1) coeff_token [total_coeff][trailing_ones]. +pub const COEFF_TOKEN_CDC_LEN: [[u8; 4]; 5] = [[2, 0, 0, 0], [6, 1, 0, 0], [6, 6, 3, 0], [6, 7, 7, 6], [6, 8, 8, 7]]; +pub const COEFF_TOKEN_CDC_CODE: [[u8; 4]; 5] = [[1, 0, 0, 0], [7, 1, 0, 0], [4, 6, 1, 0], [3, 3, 2, 5], [2, 3, 2, 0]]; + +/// total_zeros for 4x4 blocks: [total_coeff - 1][total_zeros] = (len, code). +pub const TOTAL_ZEROS: [[(u8, u8); 16]; 15] = [ + [(1, 1), (3, 3), (3, 2), (4, 3), (4, 2), (5, 3), (5, 2), (6, 3), (6, 2), (7, 3), (7, 2), (8, 3), (8, 2), (9, 3), (9, 2), (9, 1)], + [(3, 7), (3, 6), (3, 5), (3, 4), (3, 3), (4, 5), (4, 4), (4, 3), (4, 2), (5, 3), (5, 2), (6, 3), (6, 2), (6, 1), (6, 0), (0, 0)], + [(4, 5), (3, 7), (3, 6), (3, 5), (4, 4), (4, 3), (3, 4), (3, 3), (4, 2), (5, 3), (5, 2), (6, 1), (5, 1), (6, 0), (0, 0), (0, 0)], + [(5, 3), (3, 7), (4, 5), (4, 4), (3, 6), (3, 5), (3, 4), (4, 3), (3, 3), (4, 2), (5, 2), (5, 1), (5, 0), (0, 0), (0, 0), (0, 0)], + [(4, 5), (4, 4), (4, 3), (3, 7), (3, 6), (3, 5), (3, 4), (3, 3), (4, 2), (5, 1), (4, 1), (5, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(6, 1), (5, 1), (3, 7), (3, 6), (3, 5), (3, 4), (3, 3), (3, 2), (4, 1), (3, 1), (6, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(6, 1), (5, 1), (3, 5), (3, 4), (3, 3), (2, 3), (3, 2), (4, 1), (3, 1), (6, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(6, 1), (4, 1), (5, 1), (3, 3), (2, 3), (2, 2), (3, 2), (3, 1), (6, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(6, 1), (6, 0), (4, 1), (2, 3), (2, 2), (3, 1), (2, 1), (5, 1), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(5, 1), (5, 0), (3, 1), (2, 3), (2, 2), (2, 1), (4, 1), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(4, 0), (4, 1), (3, 1), (3, 2), (1, 1), (3, 3), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(4, 0), (4, 1), (2, 1), (1, 1), (3, 1), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(3, 0), (3, 1), (1, 1), (2, 1), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(2, 0), (2, 1), (1, 1), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(1, 0), (1, 1), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], +]; + +/// total_zeros for 2x2 chroma DC: [total_coeff - 1][total_zeros]. +pub const TOTAL_ZEROS_CDC: [[(u8, u8); 4]; 3] = [[(1, 1), (2, 1), (3, 1), (3, 0)], [(1, 1), (2, 1), (2, 0), (0, 0)], [(1, 1), (1, 0), (0, 0), (0, 0)]]; + +/// run_before: [min(zeros_left, 7) - 1][run_before] = (len, code). +pub const RUN_BEFORE: [[(u8, u8); 15]; 7] = [ + [(1, 1), (1, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(1, 1), (2, 1), (2, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(2, 3), (2, 2), (2, 1), (2, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(2, 3), (2, 2), (2, 1), (3, 1), (3, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(2, 3), (2, 2), (3, 3), (3, 2), (3, 1), (3, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(2, 3), (3, 0), (3, 1), (3, 3), (3, 2), (3, 5), (3, 4), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0), (0, 0)], + [(3, 7), (3, 6), (3, 5), (3, 4), (3, 3), (3, 2), (3, 1), (4, 1), (5, 1), (6, 1), (7, 1), (8, 1), (9, 1), (10, 1), (11, 1)], +]; + +/// coded_block_pattern -> codeNum for Intra_4x4/8x8 macroblocks (ChromaArrayType 1/2). +pub const CBP_INTRA_TO_CODE: [u8; 48] = invert_cbp(&CODE_TO_CBP_INTRA); +/// coded_block_pattern -> codeNum for inter macroblocks. +pub const CBP_INTER_TO_CODE: [u8; 48] = invert_cbp(&CODE_TO_CBP_INTER); + +const CODE_TO_CBP_INTRA: [u8; 48] = [ + 47, 31, 15, 0, 23, 27, 29, 30, 7, 11, 13, 14, 39, 43, 45, 46, 16, 3, 5, 10, 12, 19, 21, 26, 28, 35, 37, 42, 44, 1, 2, 4, 8, 17, 18, 20, 24, 6, 9, 22, 25, + 32, 33, 34, 36, 40, 38, 41, +]; +const CODE_TO_CBP_INTER: [u8; 48] = [ + 0, 16, 1, 2, 4, 8, 32, 3, 5, 10, 12, 15, 47, 7, 11, 13, 14, 6, 9, 31, 35, 37, 42, 44, 33, 34, 36, 40, 39, 43, 45, 46, 17, 18, 20, 24, 19, 21, 26, 28, 23, + 27, 29, 30, 22, 25, 38, 41, +]; + +const fn invert_cbp(t: &[u8; 48]) -> [u8; 48] { + let mut out = [0u8; 48]; + let mut i = 0; + while i < 48 { + out[t[i] as usize] = i as u8; + i += 1; + } + out +} + +// ---------------------------------------------------------------- quantisation + +/// Forward quantisation multipliers for 4x4 (position classes: even/even, odd/odd, other). +pub const QUANT4: [[i32; 3]; 6] = [[13107, 5243, 8066], [11916, 4660, 7490], [10082, 4194, 6554], [9362, 3647, 5825], [8192, 3355, 5243], [7282, 2893, 4559]]; + +/// normAdjust4x4 (v) values. +pub const DEQUANT4: [[i32; 3]; 6] = [[10, 16, 13], [11, 18, 14], [13, 20, 16], [14, 23, 18], [16, 25, 20], [18, 29, 23]]; + +/// normAdjust8x8 (v) values. +pub const DEQUANT8: [[i32; 6]; 6] = [ + [20, 18, 32, 19, 25, 24], + [22, 19, 35, 21, 28, 26], + [26, 23, 42, 24, 33, 31], + [28, 25, 45, 26, 35, 33], + [32, 28, 51, 30, 40, 38], + [36, 32, 58, 34, 46, 43], +]; + +#[inline] +pub fn pos_class4(i: usize) -> usize { + let (x, y) = (i & 3, i >> 2); + if x % 2 == 0 && y % 2 == 0 { + 0 + } else if x % 2 == 1 && y % 2 == 1 { + 1 + } else { + 2 + } +} + +#[inline] +pub fn pos_class8(i: usize) -> usize { + let (x, y) = (i & 7, i >> 3); + if x % 4 == 0 && y % 4 == 0 { + 0 + } else if x % 2 == 1 && y % 2 == 1 { + 1 + } else if x % 4 == 2 && y % 4 == 2 { + 2 + } else if (x % 4 == 0 && y % 2 == 1) || (x % 2 == 1 && y % 4 == 0) { + 3 + } else if (x % 4 == 0 && y % 4 == 2) || (x % 4 == 2 && y % 4 == 0) { + 4 + } else { + 5 + } +} + +/// QPc as a function of qPI (Table 8-15). +pub fn chroma_qp(qpi: i32) -> u8 { + const T: [u8; 22] = [29, 30, 31, 32, 32, 33, 34, 34, 35, 35, 36, 36, 37, 37, 37, 38, 38, 38, 39, 39, 39, 39]; + let q = qpi.clamp(0, 51); + if q < 30 { q as u8 } else { T[(q - 30) as usize] } +} + +// ---------------------------------------------------------------- deblocking + +pub const ALPHA: [u8; 52] = [ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 4, 4, 5, 6, 7, 8, 9, 10, 12, 13, 15, 17, 20, 22, 25, 28, 32, 36, 40, 45, 50, 56, 63, 71, 80, 90, 101, 113, + 127, 144, 162, 182, 203, 226, 255, 255, +]; +pub const BETA: [u8; 52] = [ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13, 14, 14, 15, 15, 16, + 16, 17, 17, 18, 18, +]; +pub const TC0: [[u8; 3]; 52] = [ + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 0], + [0, 0, 1], + [0, 0, 1], + [0, 0, 1], + [0, 0, 1], + [0, 1, 1], + [0, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 1], + [1, 1, 2], + [1, 1, 2], + [1, 1, 2], + [1, 1, 2], + [1, 2, 3], + [1, 2, 3], + [2, 2, 3], + [2, 2, 4], + [2, 3, 4], + [2, 3, 4], + [3, 3, 5], + [3, 4, 6], + [3, 4, 6], + [4, 5, 7], + [4, 5, 8], + [4, 6, 9], + [5, 7, 10], + [6, 8, 11], + [6, 8, 13], + [7, 10, 14], + [8, 11, 16], + [9, 12, 18], + [10, 13, 20], + [11, 15, 23], + [13, 17, 25], +]; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cbp_tables_are_permutations() { + let mut seen = [false; 48]; + for &c in &CBP_INTRA_TO_CODE { + seen[c as usize] = true; + } + assert!(seen.iter().all(|&s| s)); + let mut seen = [false; 48]; + for &c in &CBP_INTER_TO_CODE { + seen[c as usize] = true; + } + assert!(seen.iter().all(|&s| s)); + } + + #[test] + fn scans_are_permutations() { + let mut s = ZIGZAG8; + s.sort(); + assert!(s.iter().enumerate().all(|(i, &v)| i == v)); + let mut s = ZIGZAG4; + s.sort(); + assert!(s.iter().enumerate().all(|(i, &v)| i == v)); + for (b, &(x, y)) in BLK4_XY.iter().enumerate() { + assert_eq!(RASTER_TO_BLK4[y * 4 + x], b); + } + } + + /// Every VLC table must be prefix-free (a cheap consistency check on transcription). + #[test] + fn vlc_tables_prefix_free() { + fn check(codes: &[(u8, u32)]) { + for (i, &(la, ca)) in codes.iter().enumerate() { + for (j, &(lb, cb)) in codes.iter().enumerate() { + if i == j || la == 0 || lb == 0 || la > lb { + continue; + } + assert!(cb >> (lb - la) != ca, "prefix clash {la}:{ca:b} vs {lb}:{cb:b}"); + } + } + } + for t in 0..3 { + let mut v = vec![]; + for tc in 0..17 { + for t1 in 0..4 { + if t1 <= tc && COEFF_TOKEN_LEN[t][tc][t1] > 0 { + v.push((COEFF_TOKEN_LEN[t][tc][t1], COEFF_TOKEN_CODE[t][tc][t1] as u32)); + } + } + } + assert_eq!(v.len(), 62); + check(&v); + } + let mut v = vec![]; + for tc in 0..5 { + for t1 in 0..4 { + if t1 <= tc && COEFF_TOKEN_CDC_LEN[tc][t1] > 0 { + v.push((COEFF_TOKEN_CDC_LEN[tc][t1], COEFF_TOKEN_CDC_CODE[tc][t1] as u32)); + } + } + } + assert_eq!(v.len(), 14); + check(&v); + for (tc, row) in TOTAL_ZEROS.iter().enumerate() { + let v: Vec<_> = row.iter().take(16 - tc).map(|&(l, c)| (l, c as u32)).collect(); + check(&v); + } + for (i, row) in RUN_BEFORE.iter().enumerate() { + let n = if i < 6 { i + 2 } else { 15 }; + let v: Vec<_> = row.iter().take(n).map(|&(l, c)| (l, c as u32)).collect(); + check(&v); + } + } +} diff --git a/crates/h264enc/src/transform.rs b/crates/h264enc/src/transform.rs new file mode 100644 index 00000000..526712cc --- /dev/null +++ b/crates/h264enc/src/transform.rs @@ -0,0 +1,424 @@ +//! Integer transforms, quantisation and the normative inverse transforms/scaling (§8.5.12). + +use crate::tables::{DEQUANT4, DEQUANT8, QUANT4, pos_class4, pos_class8}; + +/// Forward 4x4 core transform of a residual block (raster order). +#[inline] +pub fn fdct4(d: &[i32; 16]) -> [i32; 16] { + let mut t = [0i32; 16]; + for i in 0..4 { + let (a, b, c, e) = (d[i * 4], d[i * 4 + 1], d[i * 4 + 2], d[i * 4 + 3]); + let s03 = a + e; + let d03 = a - e; + let s12 = b + c; + let d12 = b - c; + t[i * 4] = s03 + s12; + t[i * 4 + 1] = 2 * d03 + d12; + t[i * 4 + 2] = s03 - s12; + t[i * 4 + 3] = d03 - 2 * d12; + } + let mut o = [0i32; 16]; + for i in 0..4 { + let (a, b, c, e) = (t[i], t[4 + i], t[8 + i], t[12 + i]); + let s03 = a + e; + let d03 = a - e; + let s12 = b + c; + let d12 = b - c; + o[i] = s03 + s12; + o[4 + i] = 2 * d03 + d12; + o[8 + i] = s03 - s12; + o[12 + i] = d03 - 2 * d12; + } + o +} + +/// Normative inverse 4x4 transform of scaled coefficients `d` (raster), returns residual r (after the >>6). +#[inline] +pub fn idct4(d: &[i32; 16]) -> [i32; 16] { + let mut f = [0i32; 16]; + for i in 0..4 { + let (d0, d1, d2, d3) = (d[i * 4], d[i * 4 + 1], d[i * 4 + 2], d[i * 4 + 3]); + let e = d0 + d2; + let ff = d0 - d2; + let g = (d1 >> 1) - d3; + let h = d1 + (d3 >> 1); + f[i * 4] = e + h; + f[i * 4 + 1] = ff + g; + f[i * 4 + 2] = ff - g; + f[i * 4 + 3] = e - h; + } + let mut r = [0i32; 16]; + for j in 0..4 { + let (f0, f1, f2, f3) = (f[j], f[4 + j], f[8 + j], f[12 + j]); + let e = f0 + f2; + let ff = f0 - f2; + let g = (f1 >> 1) - f3; + let h = f1 + (f3 >> 1); + r[j] = (e + h + 32) >> 6; + r[4 + j] = (ff + g + 32) >> 6; + r[8 + j] = (ff - g + 32) >> 6; + r[12 + j] = (e - h + 32) >> 6; + } + r +} + +#[inline(always)] +fn fdct8_1d(s: [i32; 8]) -> [i32; 8] { + let a0 = s[0] + s[7]; + let a1 = s[1] + s[6]; + let a2 = s[2] + s[5]; + let a3 = s[3] + s[4]; + let b0 = s[0] - s[7]; + let b1 = s[1] - s[6]; + let b2 = s[2] - s[5]; + let b3 = s[3] - s[4]; + let e0 = a0 + a3; + let e1 = a1 + a2; + let e2 = a0 - a3; + let e3 = a1 - a2; + [ + 8 * (e0 + e1), + 12 * b0 + 10 * b1 + 6 * b2 + 3 * b3, + 8 * e2 + 4 * e3, + 10 * b0 - 3 * b1 - 12 * b2 - 6 * b3, + 8 * (e0 - e1), + 6 * b0 - 12 * b1 + 3 * b2 + 10 * b3, + 4 * e2 - 8 * e3, + 3 * b0 - 6 * b1 + 10 * b2 - 12 * b3, + ] +} + +/// Forward 8x8 transform X = A x A^T (unnormalised; rows of A are the H.264 8x8 basis scaled by 8). +pub fn fdct8(d: &[i32; 64]) -> [i32; 64] { + let mut t = [0i32; 64]; + for y in 0..8 { + let o = fdct8_1d(d[y * 8..y * 8 + 8].try_into().unwrap()); + t[y * 8..y * 8 + 8].copy_from_slice(&o); + } + let mut out = [0i32; 64]; + for x in 0..8 { + let col = [t[x], t[8 + x], t[16 + x], t[24 + x], t[32 + x], t[40 + x], t[48 + x], t[56 + x]]; + let o = fdct8_1d(col); + for k in 0..8 { + out[k * 8 + x] = o[k]; + } + } + out +} + +#[cfg(test)] +const A8: [[i32; 8]; 8] = [ + [8, 8, 8, 8, 8, 8, 8, 8], + [12, 10, 6, 3, -3, -6, -10, -12], + [8, 4, -4, -8, -8, -4, 4, 8], + [10, -3, -12, -6, 6, 12, 3, -10], + [8, -8, -8, 8, 8, -8, -8, 8], + [6, -12, 3, 10, -10, -3, 12, -6], + [4, -8, 8, -4, -4, 8, -8, 4], + [3, -6, 10, -12, 12, -10, 6, -3], +]; + +/// Normative inverse 8x8 transform (§8.5.13), returns residual after (+32)>>6. +pub fn idct8(d: &[i32; 64]) -> [i32; 64] { + fn one(s: [i32; 8]) -> [i32; 8] { + let a0 = s[0] + s[4]; + let a4 = s[0] - s[4]; + let a2 = (s[2] >> 1) - s[6]; + let a6 = s[2] + (s[6] >> 1); + let b0 = a0 + a6; + let b2 = a4 + a2; + let b4 = a4 - a2; + let b6 = a0 - a6; + let a1 = -s[3] + s[5] - s[7] - (s[7] >> 1); + let a3 = s[1] + s[7] - s[3] - (s[3] >> 1); + let a5 = -s[1] + s[7] + s[5] + (s[5] >> 1); + let a7 = s[3] + s[5] + s[1] + (s[1] >> 1); + let b1 = a1 + (a7 >> 2); + let b7 = a7 - (a1 >> 2); + let b3 = a3 + (a5 >> 2); + let b5 = (a3 >> 2) - a5; + [b0 + b7, b2 + b5, b4 + b3, b6 + b1, b6 - b1, b4 - b3, b2 - b5, b0 - b7] + } + let mut g = [0i32; 64]; + for i in 0..8 { + let row: [i32; 8] = d[i * 8..i * 8 + 8].try_into().unwrap(); + let o = one(row); + g[i * 8..i * 8 + 8].copy_from_slice(&o); + } + let mut r = [0i32; 64]; + for j in 0..8 { + let col = [g[j], g[8 + j], g[16 + j], g[24 + j], g[32 + j], g[40 + j], g[48 + j], g[56 + j]]; + let o = one(col); + for i in 0..8 { + r[i * 8 + j] = (o[i] + 32) >> 6; + } + } + r +} + +/// Per-QP quantiser tables. +pub struct QuantTables { + /// mf4[qp%6][pos] + pub mf4: [[i32; 16]; 6], + pub v4: [[i32; 16]; 6], + pub mf8: [[i64; 64]; 6], + pub v8: [[i32; 64]; 6], +} + +impl QuantTables { + pub fn new() -> Self { + let mut mf4 = [[0; 16]; 6]; + let mut v4 = [[0; 16]; 6]; + let mut mf8 = [[0; 64]; 6]; + let mut v8 = [[0; 64]; 6]; + let norm8 = [512.0, 578.0, 320.0, 578.0, 512.0, 578.0, 320.0, 578.0f64]; + for q in 0..6 { + for i in 0..16 { + mf4[q][i] = QUANT4[q][pos_class4(i)]; + v4[q][i] = DEQUANT4[q][pos_class4(i)]; + } + for i in 0..64 { + let v = DEQUANT8[q][pos_class8(i)]; + v8[q][i] = v; + let (x, y) = (i & 7, i >> 3); + mf8[q][i] = ((1u64 << 36) as f64 / (norm8[x] * norm8[y] * v as f64)).round() as i64; + } + } + QuantTables { mf4, v4, mf8, v8 } + } +} + +impl Default for QuantTables { + fn default() -> Self { + Self::new() + } +} + +/// Deadzone rounding fractions (of 2^qbits), as numerator over 6*... we express as shift-based fractions. +#[derive(Clone, Copy)] +pub struct Deadzone { + /// rounding offset as a fraction of one quantisation step, in 1/256 units. + pub f256: i64, +} + +impl Deadzone { + pub const INTRA: Deadzone = Deadzone { f256: 85 }; // ~1/3 + pub const INTER: Deadzone = Deadzone { f256: 43 }; // ~1/6 +} + +/// Quantise a 4x4 block of transform coefficients (raster) in place into levels; returns number of non-zeros. +/// `skip_dc` leaves coefficient 0 untouched (for Intra16x16/chroma AC where DC is coded separately). +#[inline] +pub fn quant4(c: &mut [i32; 16], qt: &QuantTables, qp: u8, dz: Deadzone, skip_dc: bool) -> u32 { + let qbits = 15 + (qp / 6) as u32; + let f = ((1i64 << qbits) * dz.f256) >> 8; + let mf = &qt.mf4[(qp % 6) as usize]; + let mut nz = 0; + let start = if skip_dc { 1 } else { 0 }; + for i in start..16 { + let v = c[i]; + let a = ((v.unsigned_abs() as i64 * mf[i] as i64 + f) >> qbits) as i32; + c[i] = if v < 0 { -a } else { a }; + nz += (a != 0) as u32; + } + nz +} + +/// Dequantise 4x4 levels (raster) for flat scaling matrices (§8.5.12.1). `skip_dc` leaves c[0] unchanged. +#[inline] +pub fn dequant4(c: &mut [i32; 16], qt: &QuantTables, qp: u8, skip_dc: bool) { + let v = &qt.v4[(qp % 6) as usize]; + let sh = (qp / 6) as u32; + let start = if skip_dc { 1 } else { 0 }; + for i in start..16 { + // LevelScale = 16*v; qP>=24: (c*LS) << (qP/6-4); else (c*LS + 2^(3-qP/6)) >> (4-qP/6) == c*v << qP/6. + c[i] = (c[i] * v[i]) << sh; + } +} + +/// Quantise an 8x8 block (raster), returns non-zero count. +pub fn quant8(c: &mut [i32; 64], qt: &QuantTables, qp: u8, dz: Deadzone) -> u32 { + let qbits = 22 + (qp / 6) as u32; + let f = ((1i64 << qbits) * dz.f256) >> 8; + let mf = &qt.mf8[(qp % 6) as usize]; + let mut nz = 0; + for i in 0..64 { + let v = c[i]; + let a = ((v.unsigned_abs() as i64 * mf[i] + f) >> qbits) as i32; + c[i] = if v < 0 { -a } else { a }; + nz += (a != 0) as u32; + } + nz +} + +/// Dequantise 8x8 levels (raster), flat scaling. +pub fn dequant8(c: &mut [i32; 64], qt: &QuantTables, qp: u8) { + let v = &qt.v8[(qp % 6) as usize]; + let q6 = (qp / 6) as i32; + for i in 0..64 { + let ls = 16 * v[i]; + c[i] = if q6 >= 6 { (c[i] * ls) << (q6 - 6) } else { (c[i] * ls + (1 << (5 - q6))) >> (6 - q6) }; + } +} + +/// 4x4 Hadamard used for the Intra16x16 luma DC (unnormalised). +#[inline] +pub fn hadamard4(d: &[i32; 16]) -> [i32; 16] { + let mut t = [0i32; 16]; + for i in 0..4 { + let (a, b, c, e) = (d[i * 4], d[i * 4 + 1], d[i * 4 + 2], d[i * 4 + 3]); + t[i * 4] = a + b + c + e; + t[i * 4 + 1] = a + b - c - e; + t[i * 4 + 2] = a - b - c + e; + t[i * 4 + 3] = a - b + c - e; + } + let mut o = [0i32; 16]; + for i in 0..4 { + let (a, b, c, e) = (t[i], t[4 + i], t[8 + i], t[12 + i]); + o[i] = a + b + c + e; + o[4 + i] = a + b - c - e; + o[8 + i] = a - b - c + e; + o[12 + i] = a - b + c - e; + } + o +} + +/// Forward quantisation of the Intra16x16 DC (input: raster 4x4 of block DC coefficients). +pub fn quant_luma_dc(dc: &[i32; 16], qt: &QuantTables, qp: u8, dz: Deadzone) -> ([i32; 16], u32) { + let h = hadamard4(dc); + let qbits = 16 + (qp / 6) as u32; + let f = ((1i64 << qbits) * dz.f256) >> 8; + let mf = qt.mf4[(qp % 6) as usize][0] as i64; + let mut out = [0i32; 16]; + let mut nz = 0; + for i in 0..16 { + let v = h[i] >> 1; + let a = ((v.unsigned_abs() as i64 * mf + f) >> qbits) as i32; + out[i] = if v < 0 { -a } else { a }; + nz += (a != 0) as u32; + } + (out, nz) +} + +/// Normative Intra16x16 DC inverse (§8.5.10): levels (raster 4x4) -> dcY (raster 4x4). +pub fn dequant_luma_dc(c: &[i32; 16], qt: &QuantTables, qp: u8) -> [i32; 16] { + let f = hadamard4(c); + let ls = 16 * qt.v4[(qp % 6) as usize][0]; + let q6 = (qp / 6) as i32; + let mut o = [0i32; 16]; + for i in 0..16 { + o[i] = if q6 >= 6 { (f[i] * ls) << (q6 - 6) } else { (f[i] * ls + (1 << (5 - q6))) >> (6 - q6) }; + } + o +} + +/// Forward quantisation of chroma DC (2x2 raster of block DC coefficients). +pub fn quant_chroma_dc(dc: &[i32; 4], qt: &QuantTables, qp: u8, dz: Deadzone) -> ([i32; 4], u32) { + let h = [dc[0] + dc[1] + dc[2] + dc[3], dc[0] - dc[1] + dc[2] - dc[3], dc[0] + dc[1] - dc[2] - dc[3], dc[0] - dc[1] - dc[2] + dc[3]]; + let qbits = 16 + (qp / 6) as u32; + let f = ((1i64 << qbits) * dz.f256) >> 8; + let mf = qt.mf4[(qp % 6) as usize][0] as i64; + let mut out = [0i32; 4]; + let mut nz = 0; + for i in 0..4 { + let v = h[i]; + let a = ((v.unsigned_abs() as i64 * mf + f) >> qbits) as i32; + out[i] = if v < 0 { -a } else { a }; + nz += (a != 0) as u32; + } + (out, nz) +} + +/// Normative chroma DC inverse (§8.5.11.2) for 4:2:0: levels (raster 2x2) -> dcC. +pub fn dequant_chroma_dc(c: &[i32; 4], qt: &QuantTables, qpc: u8) -> [i32; 4] { + let f = [c[0] + c[1] + c[2] + c[3], c[0] - c[1] + c[2] - c[3], c[0] + c[1] - c[2] - c[3], c[0] - c[1] - c[2] + c[3]]; + let ls = 16 * qt.v4[(qpc % 6) as usize][0]; + let q6 = (qpc / 6) as u32; + let mut o = [0i32; 4]; + for i in 0..4 { + o[i] = ((f[i] * ls) << q6) >> 5; + } + o +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn roundtrip_4x4_low_qp() { + let qt = QuantTables::new(); + let mut d = [0i32; 16]; + for (i, v) in d.iter_mut().enumerate() { + *v = ((i * 37) % 51) as i32 - 25; + } + let mut c = fdct4(&d); + quant4(&mut c, &qt, 4, Deadzone { f256: 128 }, false); + dequant4(&mut c, &qt, 4, false); + let r = idct4(&c); + for i in 0..16 { + assert!((r[i] - d[i]).abs() <= 1, "{i}: {} vs {}", r[i], d[i]); + } + } + + #[test] + fn roundtrip_8x8_low_qp() { + let qt = QuantTables::new(); + let mut d = [0i32; 64]; + for (i, v) in d.iter_mut().enumerate() { + *v = ((i * 53) % 71) as i32 - 35; + } + let mut c = fdct8(&d); + quant8(&mut c, &qt, 4, Deadzone { f256: 128 }); + dequant8(&mut c, &qt, 4); + let r = idct8(&c); + for i in 0..64 { + assert!((r[i] - d[i]).abs() <= 1, "{i}: {} vs {}", r[i], d[i]); + } + } + + #[test] + fn fdct8_matches_matrix() { + let mut d = [0i32; 64]; + for (i, v) in d.iter_mut().enumerate() { + *v = ((i * 97) % 211) as i32 - 105; + } + let f = fdct8(&d); + for k in 0..8 { + for l in 0..8 { + let mut s = 0; + for y in 0..8 { + for x in 0..8 { + s += A8[k][y] * d[y * 8 + x] * A8[l][x]; + } + } + assert_eq!(f[k * 8 + l], s); + } + } + } + + #[test] + fn luma_dc_roundtrip() { + let qt = QuantTables::new(); + // DC of a flat block of value 10 is 160 per 4x4 block. + let dc = [160i32; 16]; + let (lv, _) = quant_luma_dc(&dc, &qt, 10, Deadzone { f256: 128 }); + let o = dequant_luma_dc(&lv, &qt, 10); + let mut blk = [0i32; 16]; + blk[0] = o[0]; + let r = idct4(&blk); + assert!(r.iter().all(|&v| (v - 10).abs() <= 1), "{r:?}"); + } + + #[test] + fn chroma_dc_roundtrip() { + let qt = QuantTables::new(); + let dc = [16 * 7, 16 * 7, 16 * 7, 16 * 7]; + let (lv, _) = quant_chroma_dc(&dc, &qt, 10, Deadzone { f256: 128 }); + let o = dequant_chroma_dc(&lv, &qt, 10); + let mut blk = [0i32; 16]; + blk[0] = o[0]; + let r = idct4(&blk); + assert!(r.iter().all(|&v| (v - 7).abs() <= 1), "{r:?}"); + } +} diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/avcc.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/avcc.h264 new file mode 100644 index 00000000..f62d3920 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/avcc.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/baseline_lowqp_slices.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/baseline_lowqp_slices.h264 new file mode 100644 index 00000000..2a30ff3f Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/baseline_lowqp_slices.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/baseline_speed.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/baseline_speed.h264 new file mode 100644 index 00000000..c26860c0 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/baseline_speed.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/bframes.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/bframes.h264 new file mode 100644 index 00000000..ef283762 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/bframes.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/cbr.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/cbr.h264 new file mode 100644 index 00000000..40963871 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/cbr.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_crf_b2.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_crf_b2.h264 new file mode 100644 index 00000000..07aa0171 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_crf_b2.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_quality_b3_slices.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_quality_b3_slices.h264 new file mode 100644 index 00000000..a7a7ca72 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_quality_b3_slices.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_speed_odd_size.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_speed_odd_size.h264 new file mode 100644 index 00000000..813d09c2 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/high_speed_odd_size.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/main_b2.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/main_b2.h264 new file mode 100644 index 00000000..2cacf7b8 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/main_b2.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/mandelbrot.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/mandelbrot.h264 new file mode 100644 index 00000000..cc9201a0 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/mandelbrot.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/psnr1080.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/psnr1080.h264 new file mode 100644 index 00000000..ce651c4a Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/psnr1080.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/testsrc2.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/testsrc2.h264 new file mode 100644 index 00000000..8fecc11a Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/testsrc2.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/twopass.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/twopass.h264 new file mode 100644 index 00000000..90018b6b Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/twopass.h264 differ diff --git a/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/vbr1pass.h264 b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/vbr1pass.h264 new file mode 100644 index 00000000..f40f0775 Binary files /dev/null and b/crates/h264enc/target/agent-h264enc/fixtures/h264enc/oracle/vbr1pass.h264 differ diff --git a/crates/h264enc/tests/common/mod.rs b/crates/h264enc/tests/common/mod.rs new file mode 100644 index 00000000..18ddc360 --- /dev/null +++ b/crates/h264enc/tests/common/mod.rs @@ -0,0 +1,154 @@ +//! Shared helpers for the encoder tests: a synthetic content generator and ffmpeg/ffprobe oracle wrappers. +#![allow(dead_code)] + +use std::path::{Path, PathBuf}; +use std::process::Command; + +pub struct Yuv { + pub w: usize, + pub h: usize, + pub y: Vec, + pub u: Vec, + pub v: Vec, +} + +impl Yuv { + pub fn frame(&self) -> filmcraft_h264enc::YuvFrame<'_> { + filmcraft_h264enc::YuvFrame { y: &self.y, u: &self.u, v: &self.v, y_stride: self.w, uv_stride: self.w / 2 } + } + pub fn bytes(&self) -> Vec { + let mut v = self.y.clone(); + v.extend_from_slice(&self.u); + v.extend_from_slice(&self.v); + v + } +} + +fn hash(mut x: u32) -> u32 { + x ^= x >> 16; + x = x.wrapping_mul(0x7feb352d); + x ^= x >> 15; + x = x.wrapping_mul(0x846ca68b); + x ^= x >> 16; + x +} + +/// Synthetic test picture `t`: four regions with a moving gradient, film-grain-like noise, +/// text-like hard edges and a sub-pixel panning texture. +pub fn synth(w: usize, h: usize, t: usize) -> Yuv { + let mut y = vec![0u8; w * h]; + let mut u = vec![0u8; w * h / 4]; + let mut v = vec![0u8; w * h / 4]; + let tf = t as f32; + for yy in 0..h { + for xx in 0..w { + let region = (xx * 2 / w) + 2 * (yy * 2 / h); + let val: f32 = match region { + 0 => { + // moving smooth gradient + let fx = xx as f32 / w as f32; + let fy = yy as f32 / h as f32; + 128.0 + 90.0 * ((fx * 6.0 + tf * 0.07).sin() * (fy * 4.0 - tf * 0.05).cos()) + } + 1 => { + // mild noise over a slow ramp (like film grain) + let n = (hash((xx as u32).wrapping_mul(7919) ^ (yy as u32).wrapping_mul(104729) ^ (t as u32).wrapping_mul(1_000_003)) & 31) as f32; + 60.0 + (xx as f32 * 0.1 + yy as f32 * 0.05) % 120.0 + n + } + 2 => { + // text-like strokes: glyph grid scrolling slowly upwards + let sy = yy + t; + let gx = xx / 12; + let gy = sy / 16; + let lx = xx % 12; + let ly = sy % 16; + let bits = hash((gx as u32) * 31 + gy as u32 * 977); + let inr = |v: usize, a: usize, b: usize| (a..=b).contains(&v); + let bar = inr(lx, 2, 9); + let stroke = (inr(lx, 2, 3) && bits & 1 != 0) + || (inr(lx, 8, 9) && bits & 2 != 0) + || (inr(ly, 3, 4) && bar && bits & 4 != 0) + || (inr(ly, 8, 9) && bar && bits & 8 != 0) + || (inr(ly, 12, 13) && bar && bits & 16 != 0); + if ly < 14 && lx < 11 && stroke { 235.0 } else { 20.0 } + } + _ => { + // panning texture with sub-pixel motion + let px = xx as f32 + tf * 1.75; + let py = yy as f32 + tf * 0.5; + let c = ((px / 9.0).sin() + (py / 7.0).cos() + ((px + py) / 23.0).sin()) * 40.0; + let checker = if ((px / 32.0).floor() as i64 + (py / 32.0).floor() as i64) % 2 == 0 { 30.0 } else { -30.0 }; + 128.0 + c + checker + } + }; + y[yy * w + xx] = val.clamp(0.0, 255.0) as u8; + } + } + for yy in 0..h / 2 { + for xx in 0..w / 2 { + let fx = xx as f32 / (w / 2) as f32; + let fy = yy as f32 / (h / 2) as f32; + u[yy * w / 2 + xx] = (128.0 + 60.0 * (fx * 3.0 + tf * 0.03).sin()) as u8; + v[yy * w / 2 + xx] = (128.0 + 50.0 * (fy * 5.0 - tf * 0.04).cos() * fx) as u8; + } + } + Yuv { w, h, y, u, v } +} + +pub fn ffmpeg() -> Option { + for p in ["/opt/homebrew/bin/ffmpeg", "/usr/local/bin/ffmpeg", "/usr/bin/ffmpeg"] { + if Path::new(p).exists() { + return Some(p.into()); + } + } + None +} + +pub fn ffprobe() -> Option { + for p in ["/opt/homebrew/bin/ffprobe", "/usr/local/bin/ffprobe", "/usr/bin/ffprobe"] { + if Path::new(p).exists() { + return Some(p.into()); + } + } + None +} + +pub fn out_dir(name: &str) -> PathBuf { + let base = std::env::var("CARGO_TARGET_DIR").map(PathBuf::from).unwrap_or_else(|_| PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../target")); + let d = base.join("fixtures").join("h264enc").join(name); + std::fs::create_dir_all(&d).unwrap(); + d +} + +/// Decode an Annex-B file with ffmpeg; returns (stderr, raw yuv420p bytes). +pub fn ffmpeg_decode(file: &Path) -> (String, Vec) { + let out = Command::new(ffmpeg().unwrap()) + .args(["-v", "error", "-ec", "0", "-i"]) + .arg(file) + .args(["-f", "rawvideo", "-pix_fmt", "yuv420p", "-"]) + .output() + .expect("run ffmpeg"); + (String::from_utf8_lossy(&out.stderr).into_owned(), out.stdout) +} + +/// Raw frames produced by ffmpeg from a lavfi source. +pub fn ffmpeg_source(src: &str, w: usize, h: usize, frames: usize) -> Vec { + let out = Command::new(ffmpeg().unwrap()) + .args(["-v", "error", "-f", "lavfi", "-i"]) + .arg(format!("{src}=size={w}x{h}:rate=30")) + .args(["-frames:v", &frames.to_string(), "-f", "rawvideo", "-pix_fmt", "yuv420p", "-"]) + .output() + .expect("run ffmpeg"); + assert!(out.status.success(), "{}", String::from_utf8_lossy(&out.stderr)); + let fs = w * h * 3 / 2; + out.stdout.chunks_exact(fs).map(|c| Yuv { w, h, y: c[..w * h].to_vec(), u: c[w * h..w * h * 5 / 4].to_vec(), v: c[w * h * 5 / 4..].to_vec() }).collect() +} + +pub fn psnr(a: &[u8], b: &[u8]) -> f64 { + assert_eq!(a.len(), b.len()); + let se: f64 = a.iter().zip(b).map(|(&x, &y)| (x as f64 - y as f64).powi(2)).sum(); + if se == 0.0 { + return 99.0; + } + 10.0 * (255.0f64 * 255.0 * a.len() as f64 / se).log10() +} diff --git a/crates/h264enc/tests/ffmpeg_oracle.rs b/crates/h264enc/tests/ffmpeg_oracle.rs new file mode 100644 index 00000000..a754442f --- /dev/null +++ b/crates/h264enc/tests/ffmpeg_oracle.rs @@ -0,0 +1,280 @@ +//! End-to-end tests using ffmpeg/ffprobe strictly as an external decoder oracle (skipped when absent). +//! +//! Every stream must decode with zero errors, with the exact frame count, and the decoder's output must be +//! bit-identical to the encoder's own reconstruction (catches prediction/transform/deblocking mismatches). + +mod common; + +use common::{Yuv, ffmpeg, ffmpeg_decode, ffmpeg_source, ffprobe, out_dir, psnr, synth}; +use filmcraft_h264enc::*; +use std::process::Command; + +fn cfg(w: usize, h: usize, profile: Profile, preset: Preset, bframes: u8, rate: RateControl, slices: usize) -> EncoderConfig { + let mut c = EncoderConfig::new(w as u32, h as u32, 30, 1); + c.profile = profile; + c.preset = preset; + c.bframes = bframes; + c.rate = rate; + c.slices = slices; + c.threads = 4; + c +} + +struct Run { + stream: Vec, + packets: Vec, + recon: Vec, +} + +fn encode(c: EncoderConfig, frames: &[Yuv]) -> Run { + let mut enc = Encoder::new(c).unwrap(); + enc.set_recon_capture(true); + let mut packets = Vec::new(); + for (i, f) in frames.iter().enumerate() { + packets.extend(enc.encode(&f.frame(), i as i64)); + } + packets.extend(enc.flush()); + let mut recon = enc.take_recon(); + recon.sort_by_key(|r| r.pts); + let stream = packets.iter().flat_map(|p| p.data.iter().copied()).collect(); + Run { stream, packets, recon } +} + +/// Decode with ffmpeg, check zero errors, exact frame count and bit-exact match with the reconstruction. +/// Returns the decoded frames (display order). +fn verify(name: &str, run: &Run, w: usize, h: usize, n: usize) -> Vec { + let path = out_dir("oracle").join(format!("{name}.h264")); + std::fs::write(&path, &run.stream).unwrap(); + let (err, dec) = ffmpeg_decode(&path); + assert!(err.trim().is_empty(), "{name}: ffmpeg reported errors: {err}"); + let fs = w * h * 3 / 2; + assert_eq!(dec.len(), n * fs, "{name}: decoded frame count"); + assert_eq!(run.recon.len(), n); + for (i, r) in run.recon.iter().enumerate() { + let d = &dec[i * fs..(i + 1) * fs]; + assert!(d[..w * h] == r.y[..], "{name}: frame {i} luma differs from reconstruction"); + assert!(d[w * h..w * h * 5 / 4] == r.u[..], "{name}: frame {i} Cb differs"); + assert!(d[w * h * 5 / 4..] == r.v[..], "{name}: frame {i} Cr differs"); + } + dec +} + +fn check_timestamps(run: &Run, n: usize) { + let mut pts: Vec = run.packets.iter().map(|p| p.pts).collect(); + for w in run.packets.windows(2) { + assert!(w[1].dts > w[0].dts, "dts must increase"); + } + for p in &run.packets { + assert!(p.dts <= p.pts, "dts {} > pts {}", p.dts, p.pts); + } + pts.sort(); + assert_eq!(pts, (0..n as i64).collect::>()); + assert!(run.packets[0].keyframe); +} + +#[test] +fn bit_exact_profiles_presets() { + if ffmpeg().is_none() { + eprintln!("ffmpeg not found; skipping"); + return; + } + let cases: [(&str, usize, usize, Profile, Preset, u8, RateControl, usize); 6] = [ + ("baseline_speed", 176, 144, Profile::Baseline, Preset::Speed, 0, RateControl::Qp(26), 1), + ("baseline_lowqp_slices", 176, 144, Profile::Baseline, Preset::Balanced, 0, RateControl::Qp(8), 3), + ("main_b2", 320, 240, Profile::Main, Preset::Balanced, 2, RateControl::Qp(30), 2), + ("high_crf_b2", 320, 240, Profile::High, Preset::Balanced, 2, RateControl::Crf(23.0), 1), + ("high_quality_b3_slices", 320, 240, Profile::High, Preset::Quality, 3, RateControl::Qp(22), 3), + ("high_speed_odd_size", 330, 250, Profile::High, Preset::Speed, 1, RateControl::Qp(34), 4), + ]; + for (name, w, h, profile, preset, b, rate, slices) in cases { + let n = 12; + let frames: Vec = (0..n).map(|t| synth(w, h, t)).collect(); + let run = encode(cfg(w, h, profile, preset, b, rate, slices), &frames); + let dec = verify(name, &run, w, h, n); + check_timestamps(&run, n); + // decoded display order must follow the source + let fs = w * h * 3 / 2; + for (i, f) in frames.iter().enumerate() { + let p = psnr(&dec[i * fs..i * fs + w * h], &f.y); + assert!(p > 28.0, "{name}: frame {i} psnr {p:.2} (reordering?)"); + } + } +} + +#[test] +fn bit_exact_lavfi_sources() { + if ffmpeg().is_none() { + return; + } + for src in ["testsrc2", "mandelbrot"] { + let (w, h, n) = (352, 288, 20); + let frames = ffmpeg_source(src, w, h, n); + let run = encode(cfg(w, h, Profile::High, Preset::Balanced, 3, RateControl::Crf(22.0), 2), &frames); + let dec = verify(src, &run, w, h, n); + let fs = w * h * 3 / 2; + let mean: f64 = (0..n).map(|i| psnr(&dec[i * fs..i * fs + w * h], &frames[i].y)).sum::() / n as f64; + assert!(mean > 38.0, "{src}: mean psnr {mean:.2}"); + } +} + +#[test] +fn psnr_1080p_crf20() { + if ffmpeg().is_none() { + return; + } + let (w, h, n) = (1920, 1080, 6); + let frames: Vec = (0..n).map(|t| synth(w, h, t)).collect(); + let mut c = cfg(w, h, Profile::High, Preset::Speed, 2, RateControl::Crf(20.0), 0); + c.threads = 0; + let run = encode(c, &frames); + let dec = verify("psnr1080", &run, w, h, n); + let fs = w * h * 3 / 2; + let ps: Vec = (0..n).map(|i| psnr(&dec[i * fs..i * fs + w * h], &frames[i].y)).collect(); + let mean = ps.iter().sum::() / n as f64; + eprintln!("1080p CRF 20 luma PSNR per frame: {ps:.2?} mean {mean:.2}"); + assert!(mean >= 40.0, "mean luma PSNR {mean:.2} < 40 dB"); + assert!(ps.iter().all(|&p| p >= 38.0), "{ps:?}"); +} + +fn bitrate(run: &Run, n: usize) -> f64 { + run.stream.len() as f64 * 8.0 / (n as f64 / 30.0) / 1000.0 +} + +#[test] +fn abr_hits_target() { + if ffmpeg().is_none() { + return; + } + let (w, h, n) = (640, 360, 90); + let frames: Vec = (0..n).map(|t| synth(w, h, t)).collect(); + for (name, rate, target) in [("cbr", RateControl::Cbr { kbps: 1500 }, 1500.0), ("vbr1pass", RateControl::Vbr { target_kbps: 1000, max_kbps: 2000 }, 1000.0)] + { + let run = encode(cfg(w, h, Profile::High, Preset::Speed, 2, rate, 0), &frames); + verify(name, &run, w, h, n); + let k = bitrate(&run, n); + eprintln!("{name}: {k:.1} kbps (target {target})"); + assert!((k - target).abs() / target <= 0.10, "{name}: {k:.1} kbps vs target {target}"); + } +} + +#[test] +fn two_pass_vbr_hits_target() { + if ffmpeg().is_none() { + return; + } + let (w, h, n) = (640, 360, 90); + let frames: Vec = (0..n).map(|t| synth(w, h, t)).collect(); + let rate = RateControl::Vbr { target_kbps: 800, max_kbps: 2400 }; + let mut c1 = cfg(w, h, Profile::High, Preset::Speed, 2, rate, 0); + c1.pass = Pass::First; + let mut enc = Encoder::new(c1.clone()).unwrap(); + for (i, f) in frames.iter().enumerate() { + enc.encode(&f.frame(), i as i64); + } + enc.flush(); + let stats = enc.pass_stats().unwrap(); + assert_eq!(stats.frames.len(), n); + let p1: u64 = stats.frames.iter().map(|f| f.bits).sum(); + eprintln!("pass 1: {:.1} kbps; {:?}", p1 as f64 / (n as f64 / 30.0) / 1000.0, &stats.frames[..6]); + let stats = PassStats::from_text(&stats.to_text()).unwrap(); + let mut c2 = c1; + c2.pass = Pass::Second(stats); + let run = encode(c2, &frames); + verify("twopass", &run, w, h, n); + let k = bitrate(&run, n); + eprintln!("2-pass: {k:.1} kbps (target 800)"); + assert!((k - 800.0).abs() / 800.0 <= 0.10, "2-pass: {k:.1} kbps vs 800"); +} + +#[test] +fn bframes_order_via_ffprobe() { + if ffmpeg().is_none() || ffprobe().is_none() { + return; + } + let (w, h, n) = (176, 144, 13); + let frames: Vec = (0..n).map(|t| synth(w, h, t)).collect(); + let run = encode(cfg(w, h, Profile::High, Preset::Balanced, 3, RateControl::Qp(28), 1), &frames); + verify("bframes", &run, w, h, n); + check_timestamps(&run, n); + let path = out_dir("oracle").join("bframes.h264"); + let out = Command::new(ffprobe().unwrap()) + .args(["-v", "error", "-select_streams", "v:0", "-show_entries", "frame=pict_type,coded_picture_number", "-of", "csv=p=0"]) + .arg(&path) + .output() + .unwrap(); + let text = String::from_utf8_lossy(&out.stdout); + let types: Vec = text.lines().map(|l| l.split(',').next().unwrap_or("").trim().to_string()).filter(|s| !s.is_empty()).collect(); + assert_eq!(types.len(), n, "{text}"); + // Display order as reported by the decoder must match what we encoded: I B B B P B B B P B B B P + let ours: Vec = { + let mut v: Vec<(i64, FrameType)> = run.packets.iter().map(|p| (p.pts, p.frame_type)).collect(); + v.sort_by_key(|x| x.0); + v.iter() + .map(|x| match x.1 { + FrameType::Idr | FrameType::I => 'I', + FrameType::P => 'P', + FrameType::B => 'B', + }) + .collect() + }; + let theirs: Vec = types.iter().map(|t| t.chars().next().unwrap()).collect(); + assert_eq!(ours, theirs); + assert!(theirs.contains(&'B')); +} + +#[test] +fn length_prefixed_with_avcc() { + if ffmpeg().is_none() { + return; + } + let (w, h, n) = (176, 144, 8); + let frames: Vec = (0..n).map(|t| synth(w, h, t)).collect(); + let mut c = cfg(w, h, Profile::High, Preset::Speed, 1, RateControl::Qp(26), 1); + c.format = PacketFormat::LengthPrefixed; + let mut enc = Encoder::new(c).unwrap(); + enc.set_recon_capture(true); + let avcc = enc.avcc(); + assert_eq!(avcc[0], 1); + assert_eq!(avcc[1], 100); + assert_eq!(avcc[4], 0xFF); + let sps_len = u16::from_be_bytes([avcc[6], avcc[7]]) as usize; + let sps = &avcc[8..8 + sps_len]; + let pps_len = u16::from_be_bytes([avcc[9 + sps_len], avcc[10 + sps_len]]) as usize; + let pps = &avcc[11 + sps_len..11 + sps_len + pps_len]; + let mut packets = Vec::new(); + for (i, f) in frames.iter().enumerate() { + packets.extend(enc.encode(&f.frame(), i as i64)); + } + packets.extend(enc.flush()); + // convert to Annex-B using the avcC parameter sets + let mut stream = Vec::new(); + for nal in [sps, pps] { + stream.extend_from_slice(&[0, 0, 0, 1]); + stream.extend_from_slice(nal); + } + for p in &packets { + let mut i = 0; + while i < p.data.len() { + let len = u32::from_be_bytes(p.data[i..i + 4].try_into().unwrap()) as usize; + stream.extend_from_slice(&[0, 0, 0, 1]); + stream.extend_from_slice(&p.data[i + 4..i + 4 + len]); + i += 4 + len; + } + } + let mut recon = enc.take_recon(); + recon.sort_by_key(|r| r.pts); + verify("avcc", &Run { stream, packets, recon }, w, h, n); +} + +#[test] +fn rejects_bad_config() { + assert!(Encoder::new(EncoderConfig::new(0, 16, 30, 1)).is_err()); + assert!(Encoder::new(EncoderConfig::new(15, 16, 30, 1)).is_err()); + assert!(Encoder::new(EncoderConfig::new(16, 16, 0, 1)).is_err()); + let mut c = EncoderConfig::new(64, 64, 30, 1); + c.bframes = 4; + assert!(Encoder::new(c).is_err()); + let mut enc = Encoder::new(EncoderConfig::new(64, 64, 30, 1)).unwrap(); + let y = vec![0u8; 10]; + assert!(enc.try_encode(&YuvFrame { y: &y, u: &y, v: &y, y_stride: 64, uv_stride: 32 }, 0).is_err()); +}