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