Skip to main content

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}