rusty_h264_encoder/lib.rs
1//! Pure-Rust H.264 encoder — raw I420 frames in, a conformant Annex-B stream out.
2//!
3//! Every frame it emits decodes **bit-exactly under ffmpeg across QP 0–51**,
4//! intra and inter. The crate is `#![forbid(unsafe_code)]`; the optional SIMD
5//! kernels behind the `asm` feature keep their `unsafe` quarantined in
6//! `rusty_h264-accel`, so that guarantee holds either way.
7//!
8//! Coding tools, default-on: `I_16x16`/`I_4x4`/`I_PCM` intra with λ-based
9//! RD/SATD mode decision; P-frames (`P_Skip`, 16×16/16×8/8×16) with quarter-pel
10//! motion compensation, rate-aware ME and a multi-reference DPB; **CABAC**
11//! entropy coding (Main profile — set `RUSTY_H264_LEGACY_CAVLC=1` to restore
12//! the Constrained Baseline + CAVLC bitstream byte-for-byte); **adaptive
13//! quantization**; the per-GOP I-frame QP cascade; in-loop deblocking; and
14//! average-bitrate rate control. Opt-in via [`EncoderConfig`]: B-frames (fixed
15//! or content-adaptive), the 8×8 transform, mb-tree temporal AQ, sub-8×8
16//! partitions and RD `P_Skip`.
17//!
18//! [`Preset`] picks the speed/quality trade-off — `Fast` (SAD, integer-pel),
19//! `Balanced` (adds sub-pel refinement; the default) or `Quality` (full RD
20//! trial-encode). The bitstream is valid either way; only the effort differs.
21//!
22//! ```
23//! use rusty_h264_encoder::{Encoder, EncoderConfig};
24//! use rusty_h264_common::YuvFrame;
25//!
26//! let cfg = EncoderConfig::new(16, 16);
27//! let mut enc = Encoder::new(cfg).unwrap();
28//! let frame = YuvFrame::black(16, 16);
29//! let bitstream = enc.encode(&frame); // Annex-B bytes for one access unit
30//! assert!(!bitstream.is_empty());
31//! ```
32
33pub mod bitacct;
34mod cabac;
35mod config;
36mod lookahead;
37mod mb16;
38mod mbtree;
39
40/// Lookahead candidate evaluations so far (mb-tree cost instrument, H-36) — a
41/// deterministic stand-in for wall time, which this box cannot measure at the
42/// precision the content effect needs. `reset` before an encode, read after.
43pub fn mbtree_satd_calls() -> u64 {
44 mbtree::SATD_CALLS.load(std::sync::atomic::Ordering::Relaxed)
45}
46/// Zeroes [`mbtree_satd_calls`].
47pub fn mbtree_satd_reset() {
48 mbtree::SATD_CALLS.store(0, std::sync::atomic::Ordering::Relaxed)
49}
50mod mvd_cost_tab;
51mod params;
52mod rc;
53mod slice;
54
55pub use crate::mb16::{EXT_MV, ME_PROBE, MVCMP, MVCMP_FRAME};
56pub use config::{EncoderConfig, LookaheadMode, Preset};
57pub use params::{Pps, Sps};
58pub use rc::RateControl;
59
60use rusty_h264_common::{BitWriter, ChromaFormat, NalUnit, NalUnitType, Profile, YuvFrame};
61
62/// Errors that can arise constructing or driving the encoder.
63#[derive(Debug, Clone, PartialEq, Eq)]
64pub enum EncodeError {
65 /// A feature outside the implemented Constrained Baseline subset was asked for.
66 Unsupported(&'static str),
67 /// The supplied frame's dimensions or plane sizes don't match the config.
68 FrameMismatch,
69}
70
71impl core::fmt::Display for EncodeError {
72 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
73 match self {
74 EncodeError::Unsupported(s) => write!(f, "unsupported: {s}"),
75 EncodeError::FrameMismatch => write!(f, "frame dimensions do not match encoder config"),
76 }
77 }
78}
79
80impl std::error::Error for EncodeError {}
81
82/// A Constrained Baseline H.264 encoder.
83#[derive(Debug)]
84pub struct Encoder {
85 cfg: EncoderConfig,
86 sps: Sps,
87 pps: Pps,
88 /// Count of frames fed so far; drives IDR placement via `gop_size`.
89 frame_index: u32,
90 /// `frame_num` of the next picture (resets to 0 at each IDR).
91 next_frame_num: u32,
92 /// Index of the current picture within its GOP (0 at IDR), for POC.
93 gop_index: u32,
94 /// Decoded-picture buffer: recent **deblocked** reconstructions (coded size),
95 /// most-recent first, used as inter references (`ref_idx` 0 = front).
96 refs: Vec<RefFrame>,
97 /// Average-bitrate controller; `None` for constant-QP encoding.
98 rc: Option<RateControl>,
99 /// Per-MB QP offset for the NEXT `encode()` (mb-tree temporal AQ). Set by the
100 /// batch path before each frame; consumed (and cleared) by `try_encode`. Empty /
101 /// `None` → no offset (byte-identical).
102 pending_qpo: Option<Vec<i32>>,
103 /// Frames held by the streaming lookahead (mb-tree needs a whole GOP before it
104 /// can assign any of its QPs). Drained a GOP at a time by `try_encode`, and at
105 /// end of stream by `flush`.
106 la_queue: Vec<YuvFrame>,
107}
108
109impl Drop for Encoder {
110 fn drop(&mut self) {
111 // Dropping with frames still buffered means the caller never flushed and has
112 // silently lost the tail of its stream. Loud in debug, free in release.
113 debug_assert!(
114 self.la_queue.is_empty() || std::thread::panicking(),
115 "Encoder dropped with {} frame(s) still in the lookahead queue — call flush()",
116 self.la_queue.len()
117 );
118 }
119}
120
121/// A reference picture: deblocked reconstruction at coded (MB-grid) resolution.
122/// Stored now (4a); read by motion compensation in 4b.
123#[derive(Clone, Debug)]
124#[allow(dead_code)]
125pub(crate) struct RefFrame {
126 // 16-byte aligned (moved from the encoder's aligned rec planes) so the openh264
127 // MC asm can load aligned reference row chunks.
128 pub y: rusty_h264_common::aligned::AlignedBytes,
129 pub u: rusty_h264_common::aligned::AlignedBytes,
130 pub v: rusty_h264_common::aligned::AlignedBytes,
131 /// Picture Order Count — the DISPLAY position. B ref-lists order L0/L1 by POC
132 /// relative to the current picture; P ignores it.
133 pub poc: i32,
134 /// The picture's `frame_num` (reference frames only advance it).
135 pub frame_num: u32,
136 /// Per-4×4-block List-0 motion (raster, `mb_w*4` wide). Populated for anchors;
137 /// read as the co-located picture (`RefPicList1[0]`) when deriving a B-frame's
138 /// spatial-direct `colZeroFlag`. `ref_idx == -1` marks intra/uncoded blocks.
139 pub mv: Vec<(i32, i32)>,
140 pub ref_idx: Vec<i32>,
141 /// Blocks-wide (`mb_w*4`), so the co-located index is `by*w4 + bx`.
142 pub w4: usize,
143 /// Cached half-pel luma planes, built on first sub-pel motion-search use.
144 ///
145 /// ENCODER-SIDE ONLY, and lazily: the motion search makes ~300 `mc_luma` calls
146 /// per macroblock while final reconstruction makes ~1, so this pays enormously
147 /// in the search and would be pure tax anywhere else. `Arc` so cloning a
148 /// `RefFrame` (the DPB does) does not copy three frame-sized planes.
149 pub hpel: std::sync::OnceLock<std::sync::Arc<rusty_h264_common::inter::HpelPlanes>>,
150}
151
152impl RefFrame {
153 /// The half-pel planes for this picture, filtering them once on first use.
154 pub(crate) fn hpel(&self, cw: usize, ch: usize) -> &rusty_h264_common::inter::HpelPlanes {
155 self.hpel.get_or_init(|| {
156 std::sync::Arc::new(rusty_h264_common::inter::build_hpel_planes(&self.y, cw, ch))
157 })
158 }
159}
160
161/// Sets the sub-pel refinement pattern (U1) for subsequent encodes in this process.
162/// 0 = 8-point ring + iterate, 1 = 4-point diamond + iterate, 2 = 8-point single
163/// pass, 3 = 4-point single pass. Exposed so the pattern can be A/B'd inside ONE
164/// binary, which is the only comparison this machine can resolve.
165/// Enables/disables the U1 online sub-pel dispatcher for subsequent encodes.
166/// Sets the λ-normalised partition-split search threshold (U2). 0 = off.
167/// Enables the U5-struct deferred sub-pel refinement (search all partition shapes at
168/// full-pel, refine only the winner). Bitstream-changing → BD-gated.
169/// Descent B: ME cost-path census [interior-fullpel, edge-fullpel, sub-pel].
170#[cfg(feature = "profile")]
171pub fn satdpath_snapshot() -> Vec<u64> { crate::mb16::satdpath::snapshot() }
172#[cfg(not(feature = "profile"))]
173pub fn satdpath_snapshot() -> Vec<u64> { Vec::new() }
174#[cfg(feature = "profile")]
175pub fn satdpath_reset() { crate::mb16::satdpath::reset() }
176#[cfg(not(feature = "profile"))]
177pub fn satdpath_reset() {}
178
179/// Descent D-2: sub-pel evaluations that re-price an already-priced MV.
180#[cfg(feature = "profile")]
181pub fn spstats_redundant() -> u64 { crate::mb16::spstats::redundant_count() }
182#[cfg(not(feature = "profile"))]
183pub fn spstats_redundant() -> u64 { 0 }
184
185/// Descent D: sub-pel ring census (profile builds only).
186#[cfg(feature = "profile")]
187pub fn spstats_snapshot() -> (Vec<u64>, Vec<u64>) { crate::mb16::spstats::snapshot() }
188#[cfg(not(feature = "profile"))]
189pub fn spstats_snapshot() -> (Vec<u64>, Vec<u64>) { (Vec::new(), Vec::new()) }
190#[cfg(feature = "profile")]
191pub fn spstats_reset() { crate::mb16::spstats::reset() }
192#[cfg(not(feature = "profile"))]
193pub fn spstats_reset() {}
194
195/// Default diamond rung mask (`[16,8,4]`).
196pub const DIA_DEFAULT_MASK: u32 = crate::mb16::DIA_DEFAULT;
197
198/// Descent A: select which rungs of the [64,32,16,8,4] diamond ladder to walk.
199pub fn set_dia_mask(m: u32) { crate::mb16::set_dia_mask(m) }
200/// Track-B B2: SAD-domain full-pel search phase (SATD from sub-pel on) — x264's
201/// cost split. Bitstream-changing; BD-gated; off = byte-identical to pre-B2.
202pub fn set_me_sadfp(on: bool) { crate::mb16::set_me_sadfp(on) }
203/// B2 mode: 0 off, 1 dispatched per frame by the `b2_mgain` probe, 2 force-on.
204pub fn set_me_sadfp_mode(m: u32) { crate::mb16::set_me_sadfp_mode(m) }
205/// Fixed-centre batched diamond passes (both cost domains). Off = cascade.
206pub fn set_me_fc(on: bool) { crate::mb16::set_me_fc(on) }
207/// H-13 split-dispatch threshold in milli-units of the mgain probe (0 = always
208/// search splits, byte-identical to pre-gate). Default 30 (= 0.03).
209pub fn set_split_mg(milli: u32) { crate::mb16::set_split_mg(milli) }
210/// H-23: smooth (x264-shape) mvd cost model in ME. Off = Exp-Golomb step fn.
211pub fn set_mv_smooth(on: bool) { crate::mb16::set_mv_smooth(on) }
212/// H-24 mv-cost mode: 0 off, 1 dispatched per frame by mgain, 2 force-on.
213pub fn set_mv_smooth_mode(m: u32) { crate::mb16::set_mv_smooth_mode(m) }
214/// Fixed-centre batched HALF-PEL sub-pel ring (satd_x4p). Off = cascade.
215pub fn set_sp_fc(on: bool) { crate::mb16::set_sp_fc(on) }
216
217/// The x264-style SUB-PEL EFFORT LADDER (H-10): one level selects a priced
218/// (ring pattern × iteration budget) rung — closing the ~24-vs-9 eval-count gap
219/// vs x264 as a BUDGET choice instead of a blanket cut.
220///
221/// 5 = ring8, iterate to convergence (the quality preset's default — max effort);
222/// 4 = ring8, ≤3 iterations/step; 3 = ring8, ≤2 iterations/step;
223/// 2 = ring8, single pass (= today's balanced preset); 1 = ring4, single pass.
224/// Levels ≥5 restore the defaults. Equivalent env knobs: `RFF_SUBPEL_PAT` +
225/// `RFF_SP_MAXIT`.
226pub fn set_subme(level: u32) {
227 let (pat, cap) = match level {
228 1 => (3, 0),
229 2 => (2, 0),
230 3 => (0, 2),
231 4 => (0, 3),
232 _ => (0, 0),
233 };
234 set_subpel_pattern(pat);
235 crate::mb16::set_sp_maxit(cap);
236}
237
238/// The SUPERFAST-CLASS rung (H-11/H-12): the Quality preset at x264 superfast's
239/// partition SHAPE — P16×16-only (splits gated off), everything else (sub-pel
240/// ladder, B2 dispatch) at defaults. Measured fair-run on foreman: **1.81× faster
241/// than default quality and STILL −0.9% BD vs x264 superfast itself.** The
242/// further effort cuts (subme 2 + SAD-fp force) were measured and REJECTED from
243/// this rung: no speed on top of shape-only (0.27× vs 0.28×) while costing BD
244/// (+1.9% foreman / +8.4% bus) — compose them manually via `set_subme` /
245/// `set_me_sadfp_mode` if wanted. Split-heavy content (bus-class) pays more at
246/// this rung; the per-frame split DISPATCH (H-11 next-brick b) is the eventual
247/// no-tax answer. Env twin: `RFF_SPLIT_T=10000000`.
248pub fn set_turbo(on: bool) {
249 set_split_t(if on { 10_000_000 } else { 0 });
250}
251/// Track-B B3: sub-pel iteration budget (0 = unlimited = byte-identical) — the
252/// bounded walk x264's subme levels have; pairs with B2. BD-gated.
253pub fn set_sp_maxit(n: u32) { crate::mb16::set_sp_maxit(n) }
254
255/// Descent A: diamond per-step evaluation census (profile builds only).
256#[cfg(feature = "profile")]
257pub fn diastats_snapshot() -> Vec<(u64, u64)> { crate::mb16::diastats::snapshot() }
258#[cfg(not(feature = "profile"))]
259pub fn diastats_snapshot() -> Vec<(u64, u64)> { Vec::new() }
260#[cfg(feature = "profile")]
261pub fn diastats_reset() { crate::mb16::diastats::reset() }
262#[cfg(not(feature = "profile"))]
263pub fn diastats_reset() {}
264
265pub fn set_defer_subpel(on: bool) {
266 crate::mb16::DEFER_SUBPEL.store(if on { 1 } else { 0 }, std::sync::atomic::Ordering::Relaxed);
267}
268
269pub fn set_split_t(t: u32) {
270 crate::mb16::SPLIT_T.store(t, std::sync::atomic::Ordering::Relaxed);
271}
272
273pub fn set_subpel_dispatch(on: bool) {
274 crate::mb16::SP_DISPATCH.store(if on { 1 } else { 0 }, std::sync::atomic::Ordering::Relaxed);
275}
276
277pub fn set_subpel_pattern(p: u32) {
278 crate::mb16::SUBPEL_PAT.store(p, std::sync::atomic::Ordering::Relaxed);
279}
280
281impl Encoder {
282 /// Creates an encoder, validating that the configuration is within the
283 /// implemented subset.
284 pub fn new(cfg: EncoderConfig) -> Result<Self, EncodeError> {
285 if !matches!(
286 cfg.profile,
287 Profile::ConstrainedBaseline | Profile::Baseline | Profile::Main | Profile::High
288 ) {
289 return Err(EncodeError::Unsupported("unsupported profile"));
290 }
291 // The 8×8 transform is a High-profile CAVLC feature (our decoder has no CABAC 8×8).
292 if cfg.transform_8x8 && (!matches!(cfg.profile, Profile::High) || cfg.cabac) {
293 return Err(EncodeError::Unsupported("8x8 transform requires High profile + CAVLC"));
294 }
295 // B-frames are illegal in Baseline (the decoder enforces this too): Main only.
296 if cfg.bframes > 0 && !matches!(cfg.profile, Profile::Main) {
297 return Err(EncodeError::Unsupported("B-frames require Main profile"));
298 }
299 if cfg.chroma != ChromaFormat::Yuv420 {
300 return Err(EncodeError::Unsupported("only 4:2:0 chroma"));
301 }
302 if cfg.width == 0 || cfg.height == 0 || cfg.width % 2 != 0 || cfg.height % 2 != 0 {
303 return Err(EncodeError::Unsupported("dimensions must be positive and even"));
304 }
305 let sps = Sps::from_config(&cfg);
306 let pps = Pps::from_config(&cfg);
307 let rc = (cfg.bitrate > 0).then(|| RateControl::new(cfg.bitrate, cfg.framerate, cfg.qp));
308 Ok(Self {
309 cfg,
310 sps,
311 pps,
312 frame_index: 0,
313 next_frame_num: 0,
314 gop_index: 0,
315 refs: Vec::new(),
316 rc,
317 pending_qpo: None,
318 la_queue: Vec::new(),
319 })
320 }
321
322 /// Sets the per-MB QP offset applied to the NEXT [`encode`](Self::encode) call
323 /// (mb-tree temporal AQ). One entry per macroblock (raster). Consumed once.
324 pub(crate) fn set_pending_qpo(&mut self, qpo: Vec<i32>) {
325 self.pending_qpo = Some(qpo);
326 }
327
328 /// The active configuration.
329 pub fn config(&self) -> &EncoderConfig {
330 &self.cfg
331 }
332
333 /// Encodes one frame, returning the Annex-B access unit. Every `gop_size`
334 /// frames (and always the first) is coded as an IDR, prefixed with SPS/PPS.
335 ///
336 /// Generation 1 codes *every* picture as an IDR (all-intra); inter frames
337 /// arrive with motion compensation later.
338 pub fn encode(&mut self, frame: &YuvFrame) -> Vec<u8> {
339 self.try_encode(frame).expect("frame matched config")
340 }
341
342 /// Fallible [`encode`](Self::encode).
343 ///
344 /// With a lookahead feature active (currently mb-tree, on by default in the
345 /// constant-QP path) this BUFFERS: mb-tree needs a whole GOP of future frames
346 /// before it can assign any of their QPs, so the returned `Vec` is empty while
347 /// the GOP fills and then carries that entire GOP's access units at once. The
348 /// concatenation of every return value plus [`flush`](Self::flush) is exactly
349 /// what [`encode_all`](Self::encode_all) produces — byte for byte.
350 ///
351 /// **You must call [`flush`](Self::flush) at end of stream** or the final
352 /// partial GOP is never emitted. (A debug build asserts if the encoder is
353 /// dropped with frames still buffered.) For zero added latency set
354 /// `cfg.mbtree = false`, which restores one-AU-per-call behaviour.
355 pub fn try_encode(&mut self, frame: &YuvFrame) -> Result<Vec<u8>, EncodeError> {
356 if !self.lookahead_active() {
357 return self.encode_direct(frame);
358 }
359 if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
360 return Err(EncodeError::FrameMismatch);
361 }
362 self.la_queue.push(frame.clone());
363 // The GOP is mb-tree's natural window, so buffering exactly one GOP makes
364 // the streaming result identical to the batch path's by construction.
365 if self.la_queue.len() >= self.cfg.gop_size.max(1) as usize {
366 self.emit_lookahead_gop()
367 } else {
368 Ok(Vec::new())
369 }
370 }
371
372 /// True when a feature needs future frames, so [`try_encode`] must buffer.
373 /// B-frames already refuse the streaming API, and rate control drives its own
374 /// sequential path, so mb-tree in constant-QP mode is the only case.
375 fn lookahead_active(&self) -> bool {
376 self.cfg.mbtree && self.cfg.bframes == 0 && self.cfg.bitrate == 0
377 }
378
379 /// Codes every buffered frame with mb-tree's per-GOP QP offsets and returns
380 /// their access units concatenated. Identical to what an `encode_all` worker
381 /// does for the same GOP.
382 fn emit_lookahead_gop(&mut self) -> Result<Vec<u8>, EncodeError> {
383 let frames = std::mem::take(&mut self.la_queue);
384 let offs = mbtree::gop_qp_offsets(&self.cfg, &frames, self.cfg.mbtree_strength);
385 let mut out = Vec::new();
386 for (i, f) in frames.iter().enumerate() {
387 if let Some(o) = offs.get(i) {
388 self.pending_qpo = Some(o.clone());
389 }
390 out.extend_from_slice(&self.encode_direct(f)?);
391 }
392 Ok(out)
393 }
394
395 /// Emits any frames still held by the lookahead queue (end of stream).
396 ///
397 /// Returns the trailing access units, or empty when nothing is buffered — so it
398 /// is always safe to call, including when no lookahead feature is active.
399 pub fn flush(&mut self) -> Vec<u8> {
400 self.try_flush().expect("buffered frames matched the config when accepted")
401 }
402
403 /// Fallible [`flush`](Self::flush).
404 pub fn try_flush(&mut self) -> Result<Vec<u8>, EncodeError> {
405 if self.la_queue.is_empty() {
406 return Ok(Vec::new());
407 }
408 self.emit_lookahead_gop()
409 }
410
411 /// The unbuffered single-frame path: codes `frame` immediately. This is what the
412 /// batch path's workers call, since they compute the lookahead themselves.
413 fn encode_direct(&mut self, frame: &YuvFrame) -> Result<Vec<u8>, EncodeError> {
414 let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Total);
415 if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
416 return Err(EncodeError::FrameMismatch);
417 }
418
419 // B-frames need lookahead (a future anchor coded before the B), which the
420 // one-frame-in streaming API can't provide — use `encode_all` for B.
421 if self.cfg.bframes > 0 {
422 return Err(EncodeError::Unsupported("B-frames need encode_all (lookahead)"));
423 }
424 // GOP placement: an IDR at each `gop_size` boundary, P-frames between.
425 let is_idr = self.cfg.gop_size <= 1 || self.frame_index % self.cfg.gop_size == 0;
426 if is_idr {
427 self.gop_index = 0;
428 self.next_frame_num = 0;
429 self.refs.clear();
430 }
431 let frame_num = self.next_frame_num;
432 let poc_lsb = (2 * self.gop_index) % 16;
433 // mb-tree per-MB QP offset for this frame (empty = none / byte-identical).
434 let qpo = self.pending_qpo.take().unwrap_or_default();
435
436 // Rate control (if enabled) chooses this frame's QP from a cheap
437 // look-ahead complexity estimate; otherwise the QP is fixed.
438 let complexity = if self.rc.is_some() {
439 lookahead::complexity(&self.cfg, frame, if is_idr { None } else { self.refs.first() })
440 } else {
441 0.0
442 };
443 let qp = match &self.rc {
444 Some(rc) => rc.pick_qp(is_idr, complexity),
445 // Constant-QP: apply the per-GOP I-frame cascade offset (0 by default →
446 // byte-identical). Keeps the P-only path consistent with `code_picture`.
447 None if is_idr => (self.cfg.qp as i32 + self.cfg.i_qp_offset).clamp(0, 51) as u8,
448 None => self.cfg.qp,
449 };
450
451 let mut out = Vec::new();
452 // Pre-size the slice writer to a generous fraction of the raw frame so the
453 // CAVLC hot loop never reallocs mid-frame (byte-identical; just capacity).
454 let mut w = BitWriter::with_capacity(self.cfg.width * self.cfg.height / 2 + 4096);
455 let (nal_type, mut reference) = if is_idr {
456 // SPS/PPS precede every IDR so the stream is independently decodable.
457 self.sps.to_nal().write_annex_b(&mut out);
458 self.pps.to_nal().write_annex_b(&mut out);
459 slice::write_idr_slice_header(&mut w, &self.cfg, qp);
460 let r = if self.cfg.cabac {
461 mb16::encode_slice_data_cabac_intra(&mut w, &self.cfg, frame, qp, &qpo)
462 } else {
463 mb16::encode_slice_data(&mut w, &self.cfg, frame, qp, false, &[], &qpo)
464 };
465 (NalUnitType::IdrSlice, r)
466 } else {
467 slice::write_p_slice_header(&mut w, &self.cfg, qp, frame_num, poc_lsb, self.refs.len());
468 let r = if self.cfg.cabac {
469 mb16::encode_slice_data_cabac_p(&mut w, &self.cfg, frame, qp, &self.refs, &qpo)
470 } else {
471 mb16::encode_slice_data(&mut w, &self.cfg, frame, qp, true, &self.refs, &qpo)
472 };
473 (NalUnitType::NonIdrSlice, r)
474 };
475 // POC/frame_num carried on the reference so B-frame ref-lists (when enabled)
476 // can order L0/L1 by display position. Unused on the P-only path.
477 reference.poc = 2 * self.gop_index as i32;
478 reference.frame_num = frame_num;
479 let slice_bytes = w.into_bytes();
480 // Feed the coded slice size (the picture's own bits) back to the controller.
481 if let Some(rc) = &mut self.rc {
482 rc.update(is_idr, slice_bytes.len() * 8, qp, complexity);
483 }
484 {
485 let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::EncNal);
486 NalUnit::new(3, nal_type, slice_bytes).write_annex_b(&mut out);
487 }
488
489 // The deblocked reconstruction enters the DPB (most-recent first), which
490 // is kept to `max_num_ref_frames` by a sliding window.
491 self.refs.insert(0, reference);
492 self.refs.truncate(self.cfg.num_ref_frames.max(1) as usize);
493 self.frame_index += 1;
494 self.gop_index += 1;
495 self.next_frame_num = (self.next_frame_num + 1) % 16;
496 Ok(out)
497 }
498
499 /// Batch-encodes every frame, returning one Annex-B access unit per frame.
500 ///
501 /// At constant QP the GOPs are independent — each begins with an IDR that
502 /// resets the DPB, `frame_num` and POC, and SPS/PPS precede every IDR — so they
503 /// are encoded **in parallel across CPU cores** and the result is
504 /// **byte-identical** to calling [`encode`](Self::encode) frame-by-frame. With
505 /// rate control enabled the per-frame QP depends on history, so this falls back
506 /// to sequential encoding. Within a GOP, P-frames are inherently sequential
507 /// (each predicts from the previous reconstruction); the parallelism is across
508 /// GOPs, so it scales with the number of GOPs in the clip.
509 pub fn encode_all(&self, frames: &[YuvFrame]) -> Result<Vec<Vec<u8>>, EncodeError> {
510 for f in frames {
511 if f.width != self.cfg.width || f.height != self.cfg.height || !f.is_valid() {
512 return Err(EncodeError::FrameMismatch);
513 }
514 }
515 // B-frames need a reorder pipeline (code the future anchor before the B's
516 // that reference it) — a separate sequential path.
517 if self.cfg.bframes > 0 {
518 // Content-adaptive dispatch, PER GOP (codec-content-adaptive-dispatch):
519 // code B-frames only in GOPs whose motion is predictable enough to pay,
520 // so a mixed clip gets B on its smooth segments and P on its busy ones.
521 let gop = self.cfg.gop_size.max(1) as usize;
522 let n_gops = frames.len().div_ceil(gop);
523 let (w, h) = (self.cfg.width, self.cfg.height);
524 // One cheap per-GOP signal drives BOTH content-adaptive knobs: the B/P
525 // structure dispatch AND the I-frame QP-cascade depth.
526 let gop_sig: Vec<f64> = (0..n_gops)
527 .map(|g| gop_bi_residual(&frames[g * gop..((g + 1) * gop).min(frames.len())], w, h, 1))
528 .collect();
529 let gop_fav: Vec<bool> = if self.cfg.bframes_adaptive {
530 gop_sig.iter().map(|&s| bframes_favorable(s)).collect()
531 } else {
532 vec![true; n_gops]
533 };
534 let gop_iqp: Vec<i32> = gop_sig.iter().map(|&s| gop_iqp_offset(s, self.cfg.i_qp_offset)).collect();
535 let gop_bqp: Vec<i32> = gop_sig.iter().map(|&s| gop_bframe_qp_offset(s, self.cfg.bframe_qp_offset)).collect();
536 // Adaptive B-COUNT: how many B's per anchor gap. Fixed `bframes` unless
537 // `auto`, where the 2-gap/1-gap bi-residual RATIO picks it — content that
538 // survives wider anchor spacing (low ratio) carries more cheap B's; simple
539 // translation (high ratio) wants a single equidistant B.
540 let bcount = if self.cfg.bframes_adaptive {
541 adaptive_bcount(frames, w, h, self.cfg.bframes as usize)
542 } else {
543 self.cfg.bframes as usize
544 };
545 if gop_fav.iter().any(|&f| f) {
546 return Ok(self.encode_all_bframes(frames, bcount, &gop_fav, &gop_iqp, &gop_bqp));
547 }
548 // No GOP is B-favorable → pure P-only (byte-identical to bframes=0).
549 let mut pcfg = self.cfg.clone();
550 pcfg.bframes = 0;
551 return Encoder::new(pcfg)?.encode_all(frames);
552 }
553 // Rate control threads state across frames → it must stay sequential. mb-tree
554 // runs in RC mode too: per-GOP lookahead → per-MB offsets (per-GOP centered, so
555 // rate-neutral per GOP), and the controller supplies each frame's base QP.
556 // (MEASURED: routing the cross-frame allocation through the RC's complexity
557 // instead of centering was worse — the centered offsets carry it correctly.)
558 if self.cfg.bitrate > 0 {
559 let mut enc = Encoder::new(self.cfg.clone())?;
560 let offs: Vec<Vec<i32>> = if self.cfg.mbtree {
561 let gop = self.cfg.gop_size.max(1) as usize;
562 frames
563 .chunks(gop)
564 .flat_map(|g| mbtree::gop_qp_offsets(&self.cfg, g, self.cfg.mbtree_strength))
565 .collect()
566 } else {
567 Vec::new()
568 };
569 return frames
570 .iter()
571 .enumerate()
572 .map(|(i, f)| {
573 if let Some(qpo) = offs.get(i) {
574 enc.pending_qpo = Some(qpo.clone());
575 }
576 // Bypass the streaming lookahead buffer: this path supplies the
577 // offsets itself, so buffering here would double-compute them.
578 enc.encode_direct(f)
579 })
580 .collect();
581 }
582 let gop = self.cfg.gop_size.max(1) as usize;
583 let gops: Vec<&[YuvFrame]> = frames.chunks(gop).collect();
584 if gops.is_empty() {
585 return Ok(Vec::new());
586 }
587 let n = std::env::var("RUSTY_THREADS")
588 .ok()
589 .and_then(|v| v.parse().ok())
590 .or_else(|| std::thread::available_parallelism().map(|n| n.get()).ok())
591 .unwrap_or(1)
592 .min(gops.len());
593 // Each GOP is encoded with a fresh encoder (an IDR resets all state), so
594 // GOPs distribute across `n` worker threads with no shared mutable state.
595 let mut out: Vec<Option<Vec<Vec<u8>>>> = (0..gops.len()).map(|_| None).collect();
596 let cfg = &self.cfg;
597 let gops_ref = &gops;
598 std::thread::scope(|s| {
599 let handles: Vec<_> = (0..n)
600 .map(|t| {
601 s.spawn(move || {
602 let mut local = Vec::new();
603 let mut i = t;
604 while i < gops_ref.len() {
605 let mut enc = Encoder::new(cfg.clone()).expect("config");
606 // mb-tree temporal AQ: a per-GOP lookahead over the GOP's
607 // source frames yields per-frame per-MB QP offsets (the GOP
608 // is the natural window — the IDR resets references). Off →
609 // empty → byte-identical.
610 let offs = if cfg.mbtree {
611 mbtree::gop_qp_offsets(cfg, gops_ref[i], cfg.mbtree_strength)
612 } else {
613 Vec::new()
614 };
615 let aus: Vec<Vec<u8>> = gops_ref[i]
616 .iter()
617 .enumerate()
618 .map(|(fi, f)| {
619 if let Some(o) = offs.get(fi) {
620 enc.set_pending_qpo(o.clone());
621 }
622 enc.encode_direct(f).expect("frame matched config")
623 })
624 .collect();
625 local.push((i, aus));
626 i += n;
627 }
628 local
629 })
630 })
631 .collect();
632 for h in handles {
633 for (i, aus) in h.join().expect("encode worker panicked") {
634 out[i] = Some(aus);
635 }
636 }
637 });
638 Ok(out.into_iter().flatten().flatten().collect())
639 }
640
641 /// B-frame reorder pipeline (sequential). Produces access units in **coding
642 /// order** (the decoder reorders to display order by POC). Structure: an IDR
643 /// at each `gop_size` boundary, a P anchor every `bframes+1` frames within a
644 /// GOP, `bframes` non-reference B-frames between consecutive anchors, and the
645 /// last frame forced to an anchor so trailing B's always have a future
646 /// reference. Each anchor is coded before the B's that reference it.
647 /// `gop_favorable[g]` (content-adaptive): GOP `g` codes B-frames only when
648 /// `true`; a `false` GOP is coded all-P (every frame an anchor) so busy segments
649 /// of a mixed clip don't regress. Non-adaptive callers pass all-`true`.
650 fn encode_all_bframes(&self, frames: &[YuvFrame], bcount: usize, gop_favorable: &[bool], gop_iqp: &[i32], gop_bqp: &[i32]) -> Vec<Vec<u8>> {
651 let n = frames.len();
652 if n == 0 {
653 return Vec::new();
654 }
655 let step = bcount.max(1) + 1; // B's per anchor gap + 1 (adaptive in `auto`)
656 let gop = self.cfg.gop_size.max(1) as usize;
657 // A B-capable config: Main profile + ≥2 refs so the DPB holds both anchors.
658 let mut cfg = self.cfg.clone();
659 cfg.num_ref_frames = cfg.num_ref_frames.max(2);
660 let sps = Sps::from_config(&cfg);
661 let pps = Pps::from_config(&cfg);
662
663 // Anchor display-indices: IDR at GOP starts, P anchors every `step`, plus
664 // the frame right before each IDR boundary and the clip's last frame — a
665 // trailing B with no future reference IN ITS OWN GOP would otherwise be
666 // coded after the next GOP's IDR (which clears the DPB), losing its anchors.
667 let mut is_anchor = vec![false; n];
668 for (d, a) in is_anchor.iter_mut().enumerate() {
669 // A non-favorable GOP is coded all-P (every frame an anchor); a favorable
670 // one uses the B structure.
671 *a = if gop_favorable.get(d / gop).copied().unwrap_or(true) {
672 d % gop == 0 || (d % gop) % step == 0 || (d + 1) % gop == 0
673 } else {
674 true
675 };
676 }
677 is_anchor[n - 1] = true;
678
679 // mb-tree temporal AQ over the ANCHOR reference chain: B-frames are
680 // non-reference leaves (mb-tree offsets them at ~0 anyway), so the lookahead
681 // runs over each GOP's anchor sub-sequence — the frames that actually form the
682 // reference chain — and only anchors receive an offset. `mbtree_off[d]` is that
683 // anchor's per-MB offset (empty for B's / when off → byte-identical).
684 let mbtree_off: Vec<Vec<i32>> = if cfg.mbtree {
685 let mut off = vec![Vec::new(); n];
686 let mut g = 0;
687 while g < n {
688 let gop_end = (g + gop).min(n);
689 let anchors: Vec<usize> = (g..gop_end).filter(|&d| is_anchor[d]).collect();
690 let aframes: Vec<YuvFrame> = anchors.iter().map(|&d| frames[d].clone()).collect();
691 let offs = mbtree::gop_qp_offsets(&cfg, &aframes, cfg.mbtree_strength);
692 for (i, &d) in anchors.iter().enumerate() {
693 off[d] = offs[i].clone();
694 }
695 g = gop_end;
696 }
697 off
698 } else {
699 Vec::new()
700 };
701
702 // Coding order: each anchor (display order), then the B's before it.
703 let mut order: Vec<usize> = Vec::with_capacity(n);
704 let mut prev: Option<usize> = None;
705 for d in 0..n {
706 if !is_anchor[d] {
707 continue;
708 }
709 order.push(d);
710 if let Some(p) = prev {
711 order.extend((p + 1)..d);
712 }
713 prev = Some(d);
714 }
715
716 let mut dpb: Vec<RefFrame> = Vec::new();
717 let mut aus: Vec<Vec<u8>> = Vec::with_capacity(n);
718 let mut frame_num: u32 = 0;
719 for &d in &order {
720 let is_idr = d % gop == 0;
721 if is_idr {
722 dpb.clear();
723 frame_num = 0;
724 }
725 let is_b = !is_anchor[d];
726 let gop_start = (d / gop) * gop;
727 let poc = ((d - gop_start) as i32) * 2; // POC = display position within the GOP
728 let iqp = gop_iqp.get(d / gop).copied().unwrap_or(cfg.i_qp_offset);
729 let bqp = gop_bqp.get(d / gop).copied().unwrap_or(cfg.bframe_qp_offset);
730 let qpo: &[i32] = mbtree_off.get(d).map(|v| v.as_slice()).unwrap_or(&[]);
731 let (au, recon) =
732 code_picture(&cfg, &sps, &pps, &frames[d], is_idr, is_b, poc, frame_num, &dpb, iqp, bqp, qpo);
733 aus.push(au);
734 if !is_b {
735 if let Some(r) = recon {
736 dpb.insert(0, r);
737 dpb.truncate(cfg.num_ref_frames as usize);
738 }
739 frame_num = (frame_num + 1) % 16;
740 }
741 }
742 aus
743 }
744}
745
746/// Codes ONE picture (IDR / P anchor / B) with explicit POC + frame_num + DPB.
747/// Returns the access unit and, for reference pictures, the reconstruction to add
748/// to the DPB (B-frames are non-reference → `None`). `dpb` is most-recent-first.
749#[allow(clippy::too_many_arguments)]
750fn code_picture(
751 cfg: &EncoderConfig,
752 sps: &Sps,
753 pps: &Pps,
754 frame: &YuvFrame,
755 is_idr: bool,
756 is_b: bool,
757 poc: i32,
758 frame_num: u32,
759 dpb: &[RefFrame],
760 i_qp_offset: i32,
761 b_qp_offset: i32,
762 qpo: &[i32],
763) -> (Vec<u8>, Option<RefFrame>) {
764 let mut out = Vec::new();
765 let mut w = BitWriter::with_capacity(cfg.width * cfg.height / 2 + 4096);
766 let poc_lsb = (poc as u32) & 0xF; // log2_max_pic_order_cnt_lsb = 4
767 // Per-GOP QP cascade, both offsets content-adaptive: B-frames are non-reference →
768 // quantize HARDER (`b_qp_offset`, deeper on very predictable GOPs); the GOP's
769 // I-frame is the root reference → quantize FINER (`i_qp_offset`, deeper on
770 // predictable GOPs where the I dominates the bits).
771 let qp = if is_b {
772 (cfg.qp as i32 + b_qp_offset).clamp(0, 51) as u8
773 } else if is_idr {
774 (cfg.qp as i32 + i_qp_offset).clamp(0, 51) as u8
775 } else {
776 cfg.qp
777 };
778 let (nal_type, nal_ref_idc, recon) = if is_idr {
779 sps.to_nal().write_annex_b(&mut out);
780 pps.to_nal().write_annex_b(&mut out);
781 slice::write_idr_slice_header(&mut w, cfg, qp);
782 let mut r = if cfg.cabac {
783 mb16::encode_slice_data_cabac_intra(&mut w, cfg, frame, qp, qpo)
784 } else {
785 mb16::encode_slice_data(&mut w, cfg, frame, qp, false, &[], qpo)
786 };
787 r.poc = poc;
788 r.frame_num = frame_num;
789 (NalUnitType::IdrSlice, 3u8, Some(r))
790 } else if is_b {
791 // B is non-reference. We signal one active reference per list: L0[0] =
792 // nearest PAST anchor (highest poc < current), L1[0] = nearest FUTURE anchor
793 // (lowest poc > current) — the heads of the decoder's POC-ordered B lists.
794 let l0 = dpb.iter().filter(|r| r.poc < poc).max_by_key(|r| r.poc);
795 let l1 = dpb.iter().filter(|r| r.poc > poc).min_by_key(|r| r.poc);
796 slice::write_b_slice_header(&mut w, cfg, qp, frame_num, poc_lsb, 1, 1);
797 match (l0, l1) {
798 // B-frames are non-reference leaves — mb-tree offsets them at 0 anyway, so
799 // `qpo` is `&[]` here (the anchor reference chain carries the temporal AQ).
800 (Some(l0), Some(l1)) if cfg.cabac => {
801 mb16::encode_slice_data_cabac_b(&mut w, cfg, frame, qp, poc, l0, l1, &[])
802 }
803 (Some(l0), Some(l1)) => mb16::encode_slice_data_b(&mut w, cfg, frame, qp, poc, l0, l1, &[]),
804 // A B with no bracketing anchor pair can't be List-0/1 coded; fall back
805 // to an all-B_Skip slice (spatial-direct) so the stream stays legal.
806 _ => {
807 let n = cfg.mb_width() * cfg.mb_height();
808 if cfg.cabac {
809 mb16::encode_all_skip_b_cabac(&mut w, cfg, qp, n);
810 } else {
811 w.write_ue(n as u32);
812 w.rbsp_trailing_bits();
813 }
814 }
815 }
816 (NalUnitType::NonIdrSlice, 0u8, None)
817 } else {
818 // P anchor: L0 = the DPB (past anchors), ordered most-recent-first. Both CAVLC
819 // and CABAC now code ref_idx_l0 (cb_ref_idx / parse_ref_idx_cabac), so a P slice
820 // searches + signals the full DPB (`--refs N`) under either entropy coder.
821 let p_dpb: &[RefFrame] = dpb;
822 slice::write_p_slice_header(&mut w, cfg, qp, frame_num, poc_lsb, p_dpb.len());
823 let mut r = if cfg.cabac {
824 mb16::encode_slice_data_cabac_p(&mut w, cfg, frame, qp, p_dpb, qpo)
825 } else {
826 mb16::encode_slice_data(&mut w, cfg, frame, qp, true, dpb, qpo)
827 };
828 r.poc = poc;
829 r.frame_num = frame_num;
830 (NalUnitType::NonIdrSlice, 3u8, Some(r))
831 };
832 let slice_bytes = w.into_bytes();
833 NalUnit::new(nal_ref_idc, nal_type, slice_bytes).write_annex_b(&mut out);
834 (out, recon)
835}
836
837/// The B-favorability threshold on the per-GOP signal (`gop_bi_residual`): below it,
838/// motion is predictable enough that B-frames pay AND the I-frame dominates the GOP's
839/// bits (so it wants a deeper QP cascade); above it the GOP is busy.
840const BI_THRESH: f64 = 4.0;
841
842/// Whether a GOP's temporal residual makes B-frames pay (predictable motion).
843fn bframes_favorable(residual: f64) -> bool {
844 residual < BI_THRESH
845}
846
847/// Content-adaptive per-GOP I-frame QP offset (the ip_ratio cascade, DISPATCHED by
848/// content). `base` is the busy-GOP offset (`cfg.i_qp_offset`, default −3); a
849/// predictable GOP — where the I-frame is a large fraction of the GOP's bits, so
850/// investing in it pays outsized — gets up to 2 QP steps FINER, ramping from `base`
851/// at the threshold to `base−2` at residual 0. Calibrated: busy ≈ −3, compressible
852/// ≈ −5 (−11.6% vs −7.3% at −3). `base == 0` (the opt-out) disables it entirely so
853/// the byte-identical escape hatch survives.
854fn gop_iqp_offset(residual: f64, base: i32) -> i32 {
855 if base == 0 {
856 return 0;
857 }
858 let bonus = (2.0 * ((BI_THRESH - residual) / BI_THRESH).clamp(0.0, 1.0)).round() as i32;
859 base - bonus
860}
861
862/// Content-adaptive per-GOP B-frame QP offset. B-frames are non-reference, so on a
863/// VERY predictable GOP (bi-pred + spatial-direct nail them → tiny residual) they can
864/// be quantized much HARDER for near-free bits. But the optimum is KNIFE-EDGE in the
865/// signal — measured ~+8 at residual 0.10 yet ~+2 by residual 0.29 (and a heavy LOSS
866/// at +12 there) — so unlike the I-cascade this ramp is STEEP and confined to the
867/// near-perfect-motion regime: `base` (default +2) everywhere, boosted up to +4 only
868/// as residual → 0 (decaying to `base` by ~0.3/px). Deliberately conservative — it
869/// helps near-static / clean-pan content and must never touch the common range.
870fn gop_bframe_qp_offset(residual: f64, base: i32) -> i32 {
871 const RAMP: f64 = 0.3; // residual above this gets no boost (steep — see calibration)
872 let boost = (4.0 * ((RAMP - residual) / RAMP).clamp(0.0, 1.0)).round() as i32;
873 base + boost
874}
875
876/// Adaptive B-COUNT (B-frames per anchor gap) for `auto` mode. The RATIO of the
877/// 2-gap to 1-gap bi-prediction residual measures how fast bi-pred degrades as the
878/// anchor spacing widens: LOW ratio (content survives wider gaps) carries MORE cheap
879/// non-reference B's; HIGH ratio (simple translation — degrades fast, so wider anchors
880/// cost more than the extra B's save) wants a single equidistant B. Calibrated on
881/// pans/zoom: ratio ≥ 1.8 → 1, ≥ 1.4 → 2, else 3. Capped at `max_b` (the `auto` cap).
882fn adaptive_bcount(frames: &[YuvFrame], w: usize, h: usize, max_b: usize) -> usize {
883 let cap = max_b.clamp(1, 3);
884 let g1 = gop_bi_residual(frames, w, h, 1);
885 let g2 = gop_bi_residual(frames, w, h, 2);
886 if !g1.is_finite() || !g2.is_finite() {
887 return 1;
888 }
889 let ratio = g2 / g1.max(1e-3);
890 // Calibrated on this encoder's (subsampled global-ME) ratios: a simple
891 // translation degrades to ~1.5 (→ 1 B), predictable-under-wide-gaps content sits
892 // ~1.3 or below (→ 3 B).
893 let c = if ratio >= 1.4 { 1 } else if ratio >= 1.3 { 2 } else { 3 };
894 c.clamp(1, cap)
895}
896
897/// Cheap content signal for the content-adaptive dispatch: the mean per-pixel
898/// residual of a coarse GLOBAL-motion BI-prediction, over a subsample of interior
899/// frames. Low = temporally predictable (bi-pred + spatial-direct cheap → B-frames
900/// WIN, and the I-frame dominates → deeper QP cascade); high = busy motion.
901/// `f64::INFINITY` when the GOP is too short to measure (treated as busy).
902///
903/// Global (not block) ME keeps it O(pixels)-cheap and biases toward "coherent
904/// motion", which is what spatial-direct/skip exploit. Thresholds calibrated on
905/// extremes (pan ~0.03/px, high-motion ~12.3/px); refine on a corpus.
906fn gop_bi_residual(frames: &[YuvFrame], w: usize, h: usize, gap: usize) -> f64 {
907 let n = frames.len();
908 if n < 2 * gap + 1 || w < 48 || h < 48 {
909 return f64::INFINITY;
910 }
911 // Subsampled SAD of `cur` vs `rf` shifted by (dx,dy): interior pixels only
912 // (|shift| ≤ 15 stays in-bounds, no clamping), every 4th pixel for speed.
913 let sad = |cur: &[u8], rf: &[u8], dx: isize, dy: isize| -> u64 {
914 let mut s = 0u64;
915 let mut y = 16;
916 while y < h - 16 {
917 let cbase = (y * w) as isize;
918 let rbase = ((y as isize + dy) * w as isize) + dx;
919 let mut x = 16isize;
920 while x < (w - 16) as isize {
921 let c = cur[(cbase + x) as usize] as i32;
922 let r = rf[(rbase + x) as usize] as i32;
923 s += (c - r).unsigned_abs() as u64;
924 x += 8;
925 }
926 y += 8;
927 }
928 s
929 };
930 // Coarse global ME: ±12 step 4, then refine ±3 step 1.
931 let global_me = |cur: &[u8], rf: &[u8]| -> (isize, isize) {
932 let (mut best, mut bc) = ((0isize, 0isize), u64::MAX);
933 let mut dy = -12;
934 while dy <= 12 {
935 let mut dx = -12;
936 while dx <= 12 {
937 let c = sad(cur, rf, dx, dy);
938 if c < bc {
939 bc = c;
940 best = (dx, dy);
941 }
942 dx += 4;
943 }
944 dy += 4;
945 }
946 for dy in best.1 - 3..=best.1 + 3 {
947 for dx in best.0 - 3..=best.0 + 3 {
948 let c = sad(cur, rf, dx, dy);
949 if c < bc {
950 bc = c;
951 best = (dx, dy);
952 }
953 }
954 }
955 best
956 };
957 let mut n_samp = 0usize;
958 {
959 let mut y = 16;
960 while y < h - 16 {
961 let mut x = 16;
962 while x < w - 16 {
963 n_samp += 1;
964 x += 8;
965 }
966 y += 8;
967 }
968 }
969 let step = (n / 5).max(1);
970 let (mut total, mut cnt) = (0f64, 0usize);
971 // `gap` frames each side (1 = adjacent, for the B/P dispatch; 2 probes how well
972 // bi-prediction survives WIDER anchor spacing, for the adaptive B-count).
973 let mut d = gap;
974 while d < n - gap {
975 let (cur, past, fut) = (&frames[d].y, &frames[d - gap].y, &frames[d + gap].y);
976 let (mpx, mpy) = global_me(cur, past);
977 let (mfx, mfy) = global_me(cur, fut);
978 let mut bi = 0u64;
979 let mut y = 16;
980 while y < h - 16 {
981 let mut x = 16isize;
982 while x < (w - 16) as isize {
983 let c = cur[y * w + x as usize] as i32;
984 let p = past[((y as isize + mpy) * w as isize + x + mpx) as usize] as i32;
985 let f = fut[((y as isize + mfy) * w as isize + x + mfx) as usize] as i32;
986 bi += (c - ((p + f + 1) >> 1)).unsigned_abs() as u64;
987 x += 8;
988 }
989 y += 8;
990 }
991 total += bi as f64 / n_samp as f64;
992 cnt += 1;
993 d += step;
994 }
995 if cnt > 0 {
996 total / cnt as f64
997 } else {
998 f64::INFINITY
999 }
1000}
1001
1002#[cfg(test)]
1003mod tests {
1004 use super::*;
1005
1006 #[test]
1007 fn rejects_unsupported_config() {
1008 // High profile is supported (8x8 transform); a High-profile 8x8 stream must be
1009 // CAVLC (our decoder has no CABAC 8x8) — that combination is rejected.
1010 let mut cfg = EncoderConfig::new(16, 16);
1011 cfg.profile = Profile::High;
1012 cfg.transform_8x8 = true;
1013 cfg.cabac = true;
1014 assert!(matches!(Encoder::new(cfg), Err(EncodeError::Unsupported(_))));
1015 }
1016
1017 #[test]
1018 fn encodes_access_unit_with_sps_pps_idr() {
1019 use rusty_h264_common::nal::split_annex_b;
1020 let cfg = EncoderConfig::new(32, 32);
1021 let mut enc = Encoder::new(cfg).unwrap();
1022 let frame = YuvFrame::black(32, 32);
1023 let au = enc.encode(&frame);
1024
1025 let nals = split_annex_b(&au);
1026 assert_eq!(nals.len(), 3);
1027 assert_eq!(NalUnitType::from_id(nals[0][0]), NalUnitType::Sps);
1028 assert_eq!(NalUnitType::from_id(nals[1][0]), NalUnitType::Pps);
1029 assert_eq!(NalUnitType::from_id(nals[2][0]), NalUnitType::IdrSlice);
1030 }
1031
1032 #[test]
1033 fn encode_all_matches_sequential_cqp() {
1034 // GOP-parallel batch encoding must be byte-identical to frame-by-frame
1035 // sequential encoding at constant QP (GOPs are independent).
1036 let (w, h) = (48usize, 32usize);
1037 let mut cfg = EncoderConfig::new(w, h);
1038 cfg.gop_size = 4; // 10 frames → 3 GOPs (4,4,2)
1039 let frames: Vec<YuvFrame> = (0..10u8)
1040 .map(|t| YuvFrame {
1041 width: w,
1042 height: h,
1043 y: (0..w * h).map(|i| (i as u8).wrapping_add(t.wrapping_mul(7))).collect(),
1044 u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
1045 v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
1046 })
1047 .collect();
1048 let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
1049 let mut seq: Vec<u8> = frames.iter().flat_map(|f| seq_enc.encode(f)).collect();
1050 seq.extend_from_slice(&seq_enc.flush()); // end of stream (lookahead tail)
1051 let par: Vec<u8> = Encoder::new(cfg).unwrap().encode_all(&frames).unwrap().concat();
1052 assert_eq!(seq, par, "GOP-parallel must equal sequential+flush at CQP");
1053 }
1054
1055 #[test]
1056 fn encode_all_matches_sequential_quality_preset() {
1057 // Same invariant on the QUALITY preset, whose per-frame dispatch decisions
1058 // (b2_mgain SAD/mv-cost routing) once lived in a process-global and RACED
1059 // across GOP workers — divergence only appears with >1 GOP in flight, which
1060 // the single-GOP hash harness never exercised. Content varies per frame so
1061 // the per-frame routing decisions actually differ between GOPs.
1062 let (w, h) = (48usize, 32usize);
1063 let mut cfg = EncoderConfig::new(w, h);
1064 cfg.gop_size = 3; // 12 frames → 4 GOPs, several workers in flight
1065 cfg.preset = crate::config::Preset::Quality;
1066 let frames: Vec<YuvFrame> = (0..12u8)
1067 .map(|t| YuvFrame {
1068 width: w,
1069 height: h,
1070 y: (0..w * h)
1071 .map(|i| {
1072 // alternate calm and busy frames so the mgain probe flips
1073 let base = (i as u8).wrapping_add(t.wrapping_mul(3));
1074 if t % 2 == 0 { base } else { base.wrapping_mul(37).wrapping_add(i as u8) }
1075 })
1076 .collect(),
1077 u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
1078 v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
1079 })
1080 .collect();
1081 let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
1082 let mut seq: Vec<u8> = frames.iter().flat_map(|f| seq_enc.encode(f)).collect();
1083 seq.extend_from_slice(&seq_enc.flush());
1084 let par: Vec<u8> = Encoder::new(cfg).unwrap().encode_all(&frames).unwrap().concat();
1085 assert_eq!(seq, par, "quality-preset GOP-parallel must equal sequential+flush");
1086 }
1087
1088 #[test]
1089 fn rejects_mismatched_frame() {
1090 let cfg = EncoderConfig::new(16, 16);
1091 let mut enc = Encoder::new(cfg).unwrap();
1092 let frame = YuvFrame::black(32, 16);
1093 assert_eq!(enc.try_encode(&frame), Err(EncodeError::FrameMismatch));
1094 }
1095}