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