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//! // The default config carries a lookahead (mb-tree), so `encode()` may
30//! // buffer — `flush()` at end of stream is part of the streaming contract.
31//! let mut bitstream = enc.encode(&frame);
32//! bitstream.extend_from_slice(&enc.flush());
33//! assert!(!bitstream.is_empty());
34//! ```
35
36#![cfg_attr(not(feature = "std"), no_std)]
37
38#[allow(unused_imports)]
39use rusty_h264_common::once::OnceLock;
40extern crate alloc;
41#[cfg(test)]
42extern crate std;
43
44#[cfg(all(not(feature = "std"), not(feature = "libm")))]
45compile_error!(
46 "rusty_h264-encoder without `std` needs the `libm` feature for its floating-point math"
47);
48
49// ---------------------------------------------------------------------------
50// `no_std` shims (see rusty_h264-common for the knob shim). Every `std` use
51// left in this crate is a diagnostic or a per-thread convenience; without
52// `std` a knob reads unset, a print is a no-op and a thread-local is built
53// fresh per `with` call. Defined before the modules: textual scope.
54// ---------------------------------------------------------------------------
55#[cfg(not(feature = "std"))]
56#[allow(unused_macros)]
57macro_rules! eprintln {
58 ($($t:tt)*) => {{
59 let _ = ::core::format_args!($($t)*);
60 }};
61}
62#[cfg(not(feature = "std"))]
63#[allow(unused_macros)]
64macro_rules! println {
65 ($($t:tt)*) => {{
66 let _ = ::core::format_args!($($t)*);
67 }};
68}
69/// `thread_local!` without threads: each `NAME.with(|v| ..)` builds the value
70/// fresh. The encoder uses these for recycled per-frame scratch, so without
71/// `std` that scratch is allocated per frame instead of recycled.
72#[cfg(not(feature = "std"))]
73macro_rules! thread_local {
74 () => {};
75 ($(#[$m:meta])* $vis:vis static $name:ident: $ty:ty = const { $init:expr }; $($rest:tt)*) => {
76 $(#[$m])* #[allow(non_camel_case_types)] $vis struct $name;
77 impl $name {
78 #[allow(dead_code)]
79 pub fn with<R>(&self, f: impl FnOnce(&$ty) -> R) -> R {
80 let v: $ty = $init;
81 f(&v)
82 }
83 }
84 thread_local!($($rest)*);
85 };
86 ($(#[$m:meta])* $vis:vis static $name:ident: $ty:ty = $init:expr; $($rest:tt)*) => {
87 $(#[$m])* #[allow(non_camel_case_types)] $vis struct $name;
88 impl $name {
89 #[allow(dead_code)]
90 pub fn with<R>(&self, f: impl FnOnce(&$ty) -> R) -> R {
91 let v: $ty = $init;
92 f(&v)
93 }
94 }
95 thread_local!($($rest)*);
96 };
97}
98
99#[allow(unused_imports)]
100use alloc::string::{String, ToString};
101#[allow(unused_imports)]
102use alloc::vec;
103#[allow(unused_imports)]
104use alloc::vec::Vec;
105#[allow(unused_imports)]
106use rusty_h264_common::fmath::{F32Ext as _, F64Ext as _};
107
108pub mod bitacct;
109mod cabac;
110mod config;
111mod fastmath;
112mod lookahead;
113pub mod mb16;
114mod mbtree;
115
116/// Prometheus telemetry hooks — the CABAC entropy-bin tap for offline
117/// probability-law discovery by the private Prometheus refinery (CASC
118/// campaign). Opt-in behind the `prometheus-telemetry` feature; the
119/// production build is byte-identical without it (and with it — the tap
120/// only observes the emit path, never steers it).
121#[cfg(feature = "prometheus-telemetry")]
122pub mod telemetry;
123#[cfg(feature = "prometheus-telemetry")]
124pub mod prometheus_telemetry {
125 pub use crate::telemetry::{enable, p_zero_q8, take, CabacBin, SliceTap};
126 #[allow(unused_imports)]
127 use alloc::{
128 boxed::Box,
129 format,
130 string::{String, ToString},
131 vec,
132 vec::Vec,
133 };
134 #[allow(unused_imports)]
135 use rusty_h264_common::once::OnceLock;
136}
137
138/// Lookahead candidate evaluations so far (mb-tree cost instrument, H-36) — a
139/// deterministic stand-in for wall time, which this box cannot measure at the
140/// precision the content effect needs. `reset` before an encode, read after.
141pub fn mbtree_satd_calls() -> u64 {
142 mbtree::SATD_CALLS.load(core::sync::atomic::Ordering::Relaxed)
143}
144/// Zeroes [`mbtree_satd_calls`].
145pub fn mbtree_satd_reset() {
146 mbtree::SATD_CALLS.store(0, core::sync::atomic::Ordering::Relaxed)
147}
148/// Gate fire-rate census (Tier 1 of the gate-regression harness): `(fired,
149/// seen)` per tracked gate, in [`gate_census_names`] order. Deterministic —
150/// one run is the verdict. See `signals::census`.
151/// Per-GOP mb-tree gate telemetry (Front-B harvest seam).
152pub use mbtree::gopstats;
153
154pub fn gate_census() -> Vec<(u64, u64)> {
155 signals::census::snapshot().to_vec()
156}
157/// Per-gate `(fired, seen)` split by the macroblock's TRANSFORM SIZE:
158/// `[0]` = macroblocks coded 4x4, `[1]` = coded 8x8, each in [`gate_census_names`]
159/// order. A LABEL on the existing counters, not a new gate — it answers whether a
160/// per-transform-size threshold could ever be worth fitting, before one is.
161pub fn gate_census_by_t8() -> [Vec<(u64, u64)>; 2] {
162 let s = signals::census::snapshot_by_t8();
163 [s[0].to_vec(), s[1].to_vec()]
164}
165/// Per-reference LUMA weight estimation for one P picture — the weightp fade
166/// detector (x264-parity `weightp`). A DC-ratio fit at denom 6 per reference,
167/// kept only when a subsampled zero-MV SAD improves by >1% (the x264-style
168/// keep test); identity `(64, 0)` otherwise, so non-fade content's slices
169/// carry only the table's flag bits. LUMA only — matching the streams x264's
170/// own weightp emits (its chroma stays unweighted there too).
171fn estimate_luma_weights(
172 cfg: &EncoderConfig,
173 frame: &YuvPlanes<'_>,
174 refs: &[RefFrame],
175) -> Vec<(i32, i32)> {
176 let cw = cfg.mb_width() * 16;
177 let (w, h) = (frame.width.min(cw), frame.height);
178 // The CURRENT frame's subsample grid is identical for every reference, yet
179 // it was re-walked per ref — twice per ref when the keep test ran (up to 6
180 // full grid passes at the default refs = 3). One pass now collects the sum
181 // AND the samples; each reference likewise caches its samples on its means
182 // pass so the keep test re-reads a compact buffer instead of re-striding
183 // the plane. Same samples in the same row-major order: BIT-IDENTICAL.
184 let mut cur: Vec<u8> = Vec::new();
185 let mut sc = 0u64;
186 let mut y = 0;
187 while y < h {
188 // Row slice (`w <= frame.width` by construction): the loop guard
189 // `x < w == row.len()` lets LLVM discharge the indexing, where the
190 // multiplied form re-proved bounds per sample.
191 let row = &frame.y[y * frame.width..][..w];
192 let mut x = 0;
193 while x < w {
194 let p = row[x];
195 sc += p as u64;
196 cur.push(p);
197 x += 4;
198 }
199 y += 4;
200 }
201 let n = cur.len() as u64;
202 let mut rbuf: Vec<u8> = Vec::with_capacity(cur.len());
203 refs.iter()
204 .map(|r| {
205 // Subsampled reference mean (every 4th pixel, both axes).
206 rbuf.clear();
207 let mut sr = 0u64;
208 let mut y = 0;
209 while y < h {
210 let row = &r.y[y * cw..][..w]; // w <= cw by construction
211 let mut x = 0;
212 while x < w {
213 let p = row[x];
214 sr += p as u64;
215 rbuf.push(p);
216 x += 4;
217 }
218 y += 4;
219 }
220 if n == 0 || sr == 0 {
221 return (64, 0);
222 }
223 let (mc, mr) = (sc as f64 / n as f64, sr as f64 / n as f64);
224 let lw = (rusty_h264_common::fmath::round(mc * 64.0 / mr) as i32).clamp(1, 127);
225 let lo = (rusty_h264_common::fmath::round(mc - (lw as f64) * mr / 64.0) as i32)
226 .clamp(-128, 127);
227 if (lw, lo) == (64, 0) {
228 return (64, 0);
229 }
230 // Keep test: the weighted reference must actually predict better.
231 let (mut sad_u, mut sad_w) = (0u64, 0u64);
232 for (&c, &rr) in cur.iter().zip(rbuf.iter()) {
233 let (c, rr) = (c as i32, rr as i32);
234 let rw = (((rr * lw + 32) >> 6) + lo).clamp(0, 255);
235 sad_u += c.abs_diff(rr) as u64;
236 sad_w += c.abs_diff(rw) as u64;
237 }
238 if sad_w * 100 < sad_u * 99 {
239 (lw, lo)
240 } else {
241 (64, 0)
242 }
243 })
244 .collect()
245}
246
247/// B-frame gate signal probe (harness surface for the bframes-v2 dispatch
248/// fit): per-GOP `(bi_residual_1gap, gmc_residual, mgain, dcfrac, is_screen,
249/// grain_signature)` — the same estimators the shipping gates consult, on the
250/// GOP's leading frames. Frame dimensions must be MB multiples (probe use).
251pub fn bframes_gate_signals(
252 cfg: &EncoderConfig,
253 frames: &[YuvFrame],
254) -> (f64, f64, f64, f64, bool, bool) {
255 let (w, h) = (cfg.width, cfg.height);
256 let bi = gop_bi_residual(frames, w, h, 1);
257 if frames.len() < 2 || w % 16 != 0 || h % 16 != 0 {
258 return (bi, f64::INFINITY, 0.0, 0.0, false, false);
259 }
260 let sig = signals::FrameSignals::new(&frames[1].y, w, w / 16, h / 16, Some(&frames[0].y));
261 let (mg, dc) = sig.mgain_dc();
262 (
263 bi,
264 sig.gmc_residual(),
265 mg,
266 dc,
267 sig.is_screen(),
268 sig.grain_signature(),
269 )
270}
271
272/// Scene-cut pair ratios for a frame sequence (calibration probe surface —
273/// the same detector `segment_gops` consults; index `i` is the pair
274/// `(frames[i], frames[i+1])`).
275pub fn scene_cut_ratios(cfg: &EncoderConfig, frames: &[YuvFrame]) -> Vec<f64> {
276 lookahead::all_pair_ratios(cfg, frames)
277}
278
279/// Deterministic WORK counts (`best_part`, `mb_plan`, `mb_coded`) — the speed
280/// instrument that needs no pinning. See `signals::census`.
281pub fn gate_work() -> Vec<u64> {
282 signals::census::work_snapshot().to_vec()
283}
284/// Names for [`gate_work`], same order.
285pub fn gate_work_names() -> &'static [&'static str] {
286 &signals::census::WORK_NAMES
287}
288/// Names for [`gate_census`], same order.
289pub fn gate_census_names() -> &'static [&'static str] {
290 &signals::census::NAMES
291}
292/// Zeroes the gate census.
293pub fn gate_census_reset() {
294 signals::census::reset()
295}
296
297/// LIVENESS tap: dump `gate,fired,seen` to `$RFF_CENSUS_CSV`, once, at the end
298/// of an encode. No-op when the env var is unset.
299///
300/// `fired` says a gate routed a unit. `seen` says its decision site was
301/// CONSULTED AT ALL — and that second number is the one no refit harness here
302/// could previously read. Without it, three states are indistinguishable:
303///
304/// * consulted, never routed -> the corpus lacks the content (extend it)
305/// * NEVER CONSULTED -> the path is dead (fix the configuration)
306/// * routed, output unchanged -> the arm is a no-op (delete the gate)
307///
308/// They have opposite fixes, so collapsing them sends you to the wrong work.
309/// The case that motivated this: `sub8_grain` sits behind `num_refs == 1`, so
310/// an audit run at `--refs 3` measured a gate that was switched off and
311/// reported it as neutral. A hand-written comment caught that one; this makes
312/// it mechanical.
313///
314/// Why here and not `gatecheck`: that binary reads the same counters, but it
315/// builds its OWN `EncoderConfig`, so its numbers describe a different encode
316/// than the one a refit run is judging — which is exactly how the `--refs 3`
317/// mismatch survived. This tap fires from the same process, same flags, same
318/// encode the harness is measuring.
319#[cfg(not(feature = "std"))]
320pub fn gate_census_dump_csv() {}
321#[cfg(feature = "std")]
322pub fn gate_census_dump_csv() {
323 use core::fmt::Write as _;
324 use std::io::Write as _;
325 let Ok(path) = std::env::var("RFF_CENSUS_CSV") else {
326 return;
327 };
328 let snap = signals::census::snapshot();
329 let mut s = String::from("gate,fired,seen\n");
330 for (i, name) in signals::census::NAMES.iter().enumerate() {
331 let _ = writeln!(s, "{name},{},{}", snap[i].0, snap[i].1);
332 }
333 if let Ok(mut f) = std::fs::File::create(&path) {
334 let _ = f.write_all(s.as_bytes());
335 }
336}
337
338mod mvd_cost_tab;
339mod params;
340mod rc;
341mod signals;
342mod slice;
343
344#[cfg(feature = "std")]
345pub use crate::mb16::EXT_MV;
346
347pub use crate::mb16::{ME_PROBE, MVCMP, MVCMP_FRAME};
348
349/// Test-only surface for gating the CABAC *encoder* against the decoder's parser.
350#[doc(hidden)]
351pub mod cabac_enc_test {
352 pub use crate::cabac::CabacEncoder;
353 pub use crate::mb16::b_part_mb_type;
354 pub use crate::mb16::cb_cbp;
355 pub use crate::mb16::cb_mb_qp_delta;
356 pub use crate::mb16::cb_mb_type_b;
357 pub use crate::mb16::cb_ref_idx;
358 #[allow(unused_imports)]
359 use alloc::{
360 boxed::Box,
361 format,
362 string::{String, ToString},
363 vec,
364 vec::Vec,
365 };
366 #[allow(unused_imports)]
367 use rusty_h264_common::once::OnceLock;
368}
369pub use config::{EncoderConfig, LookaheadMode, MemoryEstimate, Preset};
370pub use params::{Pps, Sps};
371pub use rc::RateControl;
372
373use rusty_h264_common::{
374 BitWriter, ChromaFormat, NalUnit, NalUnitType, Profile, YuvFrame, YuvPlanes,
375};
376
377/// Errors that can arise constructing or driving the encoder.
378#[derive(Debug, Clone, PartialEq, Eq)]
379pub enum EncodeError {
380 /// A feature outside the implemented Constrained Baseline subset was asked for.
381 Unsupported(&'static str),
382 /// The supplied frame's dimensions or plane sizes don't match the config.
383 FrameMismatch,
384 /// The caller's output buffer cannot hold the access unit; `needed` is the
385 /// size it would have taken. The frame has been coded and the encoder's
386 /// state advanced (references, GOP position), so the caller treats it as a
387 /// dropped picture and continues — the next call still works.
388 BufferTooSmall {
389 /// Bytes the access unit needed.
390 needed: usize,
391 },
392}
393
394impl core::fmt::Display for EncodeError {
395 fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
396 match self {
397 EncodeError::Unsupported(s) => write!(f, "unsupported: {s}"),
398 EncodeError::FrameMismatch => write!(f, "frame dimensions do not match encoder config"),
399 EncodeError::BufferTooSmall { needed } => {
400 write!(
401 f,
402 "output buffer too small: the access unit needs {needed} bytes"
403 )
404 }
405 }
406 }
407}
408
409impl core::error::Error for EncodeError {}
410
411/// Where an access unit goes: a `Vec` the caller receives, or the caller's own
412/// buffer, written in place. Both take the same NAL bytes from the same
413/// emulation-prevention scan; the slice form keeps counting past its end so a
414/// [`EncodeError::BufferTooSmall`] can say exactly how much was needed.
415enum Sink<'a> {
416 Vec(&'a mut Vec<u8>),
417 Slice {
418 buf: &'a mut [u8],
419 len: usize,
420 needed: usize,
421 },
422}
423
424impl Sink<'_> {
425 /// Append one Annex-B NAL unit: start code, header, emulation-prevented payload.
426 fn nal(&mut self, ref_idc: u8, nal_type: NalUnitType, rbsp: &[u8]) {
427 let header = rusty_h264_common::nal::nal_header_byte(ref_idc, nal_type);
428 match self {
429 Sink::Vec(v) => {
430 v.extend_from_slice(&[0, 0, 0, 1, header]);
431 rusty_h264_common::nal::emulation_prevent_into(rbsp, v);
432 }
433 Sink::Slice { buf, len, needed } => {
434 let mut put = |b: u8| {
435 if *len < buf.len() {
436 buf[*len] = b;
437 *len += 1;
438 }
439 *needed += 1;
440 };
441 for b in [0, 0, 0, 1, header] {
442 put(b);
443 }
444 rusty_h264_common::nal::emulation_prevent_with(rbsp, put);
445 }
446 }
447 }
448}
449
450/// Copy an access unit into a caller-owned buffer, or say how big it needed to be.
451fn copy_out(au: &[u8], out: &mut [u8]) -> Result<usize, EncodeError> {
452 if out.len() < au.len() {
453 return Err(EncodeError::BufferTooSmall { needed: au.len() });
454 }
455 out[..au.len()].copy_from_slice(au);
456 Ok(au.len())
457}
458
459/// A Constrained Baseline H.264 encoder.
460#[derive(Debug)]
461pub struct Encoder {
462 cfg: EncoderConfig,
463 sps: Sps,
464 pps: Pps,
465 /// Count of frames fed so far; drives IDR placement via `gop_size`.
466 frame_index: u32,
467 /// `frame_num` of the next picture (resets to 0 at each IDR).
468 next_frame_num: u32,
469 /// Index of the current picture within its GOP (0 at IDR), for POC.
470 gop_index: u32,
471 /// Decoded-picture buffer: recent **deblocked** reconstructions (coded size),
472 /// most-recent first, used as inter references (`ref_idx` 0 = front).
473 refs: Vec<RefFrame>,
474 /// Average-bitrate controller; `None` for constant-QP encoding.
475 rc: Option<RateControl>,
476 /// Per-MB QP offset for the NEXT `encode()` (mb-tree temporal AQ). Set by the
477 /// batch path before each frame; consumed (and cleared) by `try_encode`. Empty /
478 /// `None` → no offset (byte-identical).
479 pending_qpo: Option<Vec<i32>>,
480 /// AQ grain probe for the NEXT frame IF it is an IDR: the previous
481 /// display-order SOURCE frame (docs/gate-ledger.md aq-grain-veto — an IDR
482 /// has no coding reference, so the veto's temporal signals read
483 /// source-vs-source). Set by the batch paths, consumed per frame; `None`
484 /// (streaming, or the stream's first frame) → the veto fails open.
485 pending_aq_probe: Option<YuvFrame>,
486 /// Frames held by the streaming lookahead (mb-tree needs a whole GOP before it
487 /// can assign any of its QPs). Drained a GOP at a time by `try_encode`, and at
488 /// end of stream by `flush`.
489 la_queue: Vec<YuvFrame>,
490 /// Frames coded since the last IDR (0 = the next frame IS an IDR). The
491 /// scenecut counter that replaced `frame_index % gop_size` — cadence is no
492 /// longer periodic once cuts place IDRs (x264-parity keyint/min-keyint).
493 since_idr: u32,
494 /// One-shot IDR request (scene cut, or a batch segment boundary). Consumed
495 /// by `encode_direct`.
496 force_idr: bool,
497 /// A caller's [`request_keyframe`](Self::request_keyframe), pending until the
498 /// next submitted frame. Kept apart from `force_idr`, which a scene cut
499 /// raises for the FIRST buffered frame: a lookahead drain must spend that
500 /// one and not this one.
501 keyframe_requested: bool,
502 /// Previous display-order SOURCE frame, retained by the streaming path for
503 /// the causal scene-cut pair AND as the cut-IDR's AQ grain probe. `None`
504 /// with scenecut off (no clone cost on the anchor path).
505 last_src: Option<YuvFrame>,
506 /// `last_src`'s detector preparation (coded + half-res planes), carried so
507 /// each pair ratio preps only the NEW frame — the previous frame was
508 /// prepped by the last call (`None` on the first pair; rebuilt on demand).
509 last_prep: Option<mbtree::PairPrep>,
510 /// The two most recent scene-cut pair ratios (the spike-rule baseline),
511 /// most recent first. `1.0` = no history yet (nothing can spike over it).
512 cut_hist: [f64; 2],
513 /// Scratch for a padded (non-tight) [`YuvPlanes`] view: rows are gathered
514 /// here once per call and the frame is reused, never reallocated.
515 gather: Option<YuvFrame>,
516 /// The per-frame slice bit writer, kept across frames so its buffer is
517 /// allocated once (`BitWriter::clear` keeps the capacity).
518 bw: BitWriter,
519 /// SPS and PPS as NAL units, built once; every IDR writes them.
520 sps_nal: NalUnit,
521 pps_nal: NalUnit,
522}
523
524#[cfg(feature = "std")]
525fn panicking() -> bool {
526 std::thread::panicking()
527}
528#[cfg(not(feature = "std"))]
529fn panicking() -> bool {
530 false
531}
532
533impl Drop for Encoder {
534 fn drop(&mut self) {
535 // Dropping with frames still buffered means the caller never flushed and has
536 // silently lost the tail of its stream. Loud in debug, free in release.
537 debug_assert!(
538 self.la_queue.is_empty() || panicking(),
539 "Encoder dropped with {} frame(s) still in the lookahead queue — call flush()",
540 self.la_queue.len()
541 );
542 }
543}
544
545/// A reference picture: deblocked reconstruction at coded (MB-grid) resolution.
546/// Stored now (4a); read by motion compensation in 4b.
547#[derive(Clone, Debug)]
548#[allow(dead_code)]
549pub(crate) struct RefFrame {
550 // 16-byte aligned (moved from the encoder's aligned rec planes) so the openh264
551 // MC asm can load aligned reference row chunks.
552 pub y: rusty_h264_common::aligned::AlignedBytes,
553 pub u: rusty_h264_common::aligned::AlignedBytes,
554 pub v: rusty_h264_common::aligned::AlignedBytes,
555 /// Picture Order Count — the DISPLAY position. B ref-lists order L0/L1 by POC
556 /// relative to the current picture; P ignores it.
557 pub poc: i32,
558 /// The picture's `frame_num` (reference frames only advance it).
559 pub frame_num: u32,
560 /// Per-4×4-block List-0 motion (raster, `mb_w*4` wide). Populated for anchors;
561 /// read as the co-located picture (`RefPicList1[0]`) when deriving a B-frame's
562 /// spatial-direct `colZeroFlag`. `ref_idx == -1` marks intra/uncoded blocks.
563 pub mv: Vec<(i32, i32)>,
564 pub ref_idx: Vec<i32>,
565 /// List-1 motion of the picture (b-pyramid: a REFERENCE B can be the
566 /// co-located picture, and its L1-only blocks read from here — the exact
567 /// List-1 colZeroFlag defect the decoder already root-caused and fixed;
568 /// the encoder's direct derivation must mirror it or pyramid recon
569 /// drifts). EMPTY for P/I references (no List 1 exists there).
570 pub mv1: Vec<(i32, i32)>,
571 pub ref_idx1: Vec<i32>,
572 /// Blocks-wide (`mb_w*4`), so the co-located index is `by*w4 + bx`.
573 pub w4: usize,
574 /// Cached half-pel luma planes, built on first sub-pel motion-search use.
575 ///
576 /// ENCODER-SIDE ONLY, and lazily: the motion search makes ~300 `mc_luma` calls
577 /// per macroblock while final reconstruction makes ~1, so this pays enormously
578 /// in the search and would be pure tax anywhere else. `Arc` so cloning a
579 /// `RefFrame` (the DPB does) does not copy three frame-sized planes.
580 pub hpel: HpelOnce<alloc::sync::Arc<rusty_h264_common::inter::HpelPlanes>>,
581}
582
583/// The once-cell behind a reference frame's half-pel planes: `OnceLock` with
584/// `std` (frames cross threads in the parallel GOP path), `OnceCell` without.
585#[cfg(feature = "std")]
586pub type HpelOnce<T> = rusty_h264_common::once::OnceLock<T>;
587/// See the `std` variant.
588#[cfg(not(feature = "std"))]
589pub type HpelOnce<T> = core::cell::OnceCell<T>;
590
591impl RefFrame {
592 /// The half-pel planes for this picture, filtering them once on first use.
593 pub(crate) fn hpel(&self, cw: usize, ch: usize) -> &rusty_h264_common::inter::HpelPlanes {
594 self.hpel.get_or_init(|| {
595 alloc::sync::Arc::new(rusty_h264_common::inter::build_hpel_planes(&self.y, cw, ch))
596 })
597 }
598}
599
600/// Sets the sub-pel refinement pattern (U1) for subsequent encodes in this process.
601/// 0 = 8-point ring + iterate, 1 = 4-point diamond + iterate, 2 = 8-point single
602/// pass, 3 = 4-point single pass. Exposed so the pattern can be A/B'd inside ONE
603/// binary, which is the only comparison this machine can resolve.
604/// Enables/disables the U1 online sub-pel dispatcher for subsequent encodes.
605/// Sets the λ-normalised partition-split search threshold (U2). 0 = off.
606/// Enables the U5-struct deferred sub-pel refinement (search all partition shapes at
607/// full-pel, refine only the winner). Bitstream-changing → BD-gated.
608/// Descent B: ME cost-path census [interior-fullpel, edge-fullpel, sub-pel].
609#[cfg(feature = "profile")]
610pub fn satdpath_snapshot() -> Vec<u64> {
611 crate::mb16::satdpath::snapshot()
612}
613#[cfg(not(feature = "profile"))]
614pub fn satdpath_snapshot() -> Vec<u64> {
615 Vec::new()
616}
617#[cfg(feature = "profile")]
618pub fn satdpath_reset() {
619 crate::mb16::satdpath::reset()
620}
621#[cfg(not(feature = "profile"))]
622pub fn satdpath_reset() {}
623
624/// Descent D-2: sub-pel evaluations that re-price an already-priced MV.
625#[cfg(feature = "profile")]
626pub fn spstats_redundant() -> u64 {
627 crate::mb16::spstats::redundant_count()
628}
629#[cfg(not(feature = "profile"))]
630pub fn spstats_redundant() -> u64 {
631 0
632}
633
634/// Descent D: sub-pel ring census (profile builds only).
635#[cfg(feature = "profile")]
636pub fn spstats_snapshot() -> (Vec<u64>, Vec<u64>) {
637 crate::mb16::spstats::snapshot()
638}
639#[cfg(not(feature = "profile"))]
640pub fn spstats_snapshot() -> (Vec<u64>, Vec<u64>) {
641 (Vec::new(), Vec::new())
642}
643#[cfg(feature = "profile")]
644pub fn spstats_reset() {
645 crate::mb16::spstats::reset()
646}
647#[cfg(not(feature = "profile"))]
648pub fn spstats_reset() {}
649
650/// Default diamond rung mask (`[16,8,4]`).
651pub const DIA_DEFAULT_MASK: u32 = crate::mb16::DIA_DEFAULT;
652
653/// Descent A: select which rungs of the [64,32,16,8,4] diamond ladder to walk.
654pub fn set_dia_mask(m: u32) {
655 crate::mb16::set_dia_mask(m)
656}
657/// Track-B B2: SAD-domain full-pel search phase (SATD from sub-pel on) — x264's
658/// cost split. Bitstream-changing; BD-gated; off = byte-identical to pre-B2.
659pub fn set_me_sadfp(on: bool) {
660 crate::mb16::set_me_sadfp(on)
661}
662/// B2 mode: 0 off, 1 dispatched per frame by the `b2_mgain` probe, 2 force-on.
663pub fn set_me_sadfp_mode(m: u32) {
664 crate::mb16::set_me_sadfp_mode(m)
665}
666/// Fixed-centre batched diamond passes (both cost domains). Off = cascade.
667pub fn set_me_fc(on: bool) {
668 crate::mb16::set_me_fc(on)
669}
670/// H-13 split-dispatch threshold in milli-units of the mgain probe (0 = always
671/// search splits, byte-identical to pre-gate). Default 30 (= 0.03).
672pub fn set_split_mg(milli: u32) {
673 crate::mb16::set_split_mg(milli)
674}
675/// H-23: smooth (x264-shape) mvd cost model in ME. Off = Exp-Golomb step fn.
676pub fn set_mv_smooth(on: bool) {
677 crate::mb16::set_mv_smooth(on)
678}
679/// H-24 mv-cost mode: 0 off, 1 dispatched per frame by mgain, 2 force-on.
680pub fn set_mv_smooth_mode(m: u32) {
681 crate::mb16::set_mv_smooth_mode(m)
682}
683/// Fixed-centre batched HALF-PEL sub-pel ring (satd_x4p). Off = cascade.
684pub fn set_sp_fc(on: bool) {
685 crate::mb16::set_sp_fc(on)
686}
687
688/// The x264-style SUB-PEL EFFORT LADDER (H-10): one level selects a priced
689/// (ring pattern × iteration budget) rung — closing the ~24-vs-9 eval-count gap
690/// vs x264 as a BUDGET choice instead of a blanket cut.
691///
692/// 5 = ring8, iterate to convergence (the quality preset's default — max effort);
693/// 4 = ring8, ≤3 iterations/step; 3 = ring8, ≤2 iterations/step;
694/// 2 = ring8, single pass (= today's balanced preset); 1 = ring4, single pass.
695/// Levels ≥5 restore the defaults. Equivalent env knobs: `RFF_SUBPEL_PAT` +
696/// `RFF_SP_MAXIT`.
697pub fn set_subme(level: u32) {
698 let (pat, cap) = match level {
699 1 => (3, 0),
700 2 => (2, 0),
701 3 => (0, 2),
702 4 => (0, 3),
703 _ => (0, 0),
704 };
705 set_subpel_pattern(pat);
706 crate::mb16::set_sp_maxit(cap);
707}
708
709/// The SUPERFAST-CLASS rung (H-11/H-12): the Quality preset at x264 superfast's
710/// partition SHAPE — P16×16-only (splits gated off), everything else (sub-pel
711/// ladder, B2 dispatch) at defaults. Measured fair-run on foreman: **1.81× faster
712/// than default quality and STILL −0.9% BD vs x264 superfast itself.** The
713/// further effort cuts (subme 2 + SAD-fp force) were measured and REJECTED from
714/// this rung: no speed on top of shape-only (0.27× vs 0.28×) while costing BD
715/// (+1.9% foreman / +8.4% bus) — compose them manually via `set_subme` /
716/// `set_me_sadfp_mode` if wanted. Split-heavy content (bus-class) pays more at
717/// this rung; the per-frame split DISPATCH (H-11 next-brick b) is the eventual
718/// no-tax answer. Env twin: `RFF_SPLIT_T=10000000`.
719pub fn set_turbo(on: bool) {
720 set_split_t(if on { 10_000_000 } else { 0 });
721}
722/// Track-B B3: sub-pel iteration budget (0 = unlimited = byte-identical) — the
723/// bounded walk x264's subme levels have; pairs with B2. BD-gated.
724pub fn set_sp_maxit(n: u32) {
725 crate::mb16::set_sp_maxit(n)
726}
727
728/// Descent A: diamond per-step evaluation census (profile builds only).
729#[cfg(feature = "profile")]
730pub fn diastats_snapshot() -> Vec<(u64, u64)> {
731 crate::mb16::diastats::snapshot()
732}
733#[cfg(not(feature = "profile"))]
734pub fn diastats_snapshot() -> Vec<(u64, u64)> {
735 Vec::new()
736}
737#[cfg(feature = "profile")]
738pub fn diastats_reset() {
739 crate::mb16::diastats::reset()
740}
741#[cfg(not(feature = "profile"))]
742pub fn diastats_reset() {}
743
744pub fn set_defer_subpel(on: bool) {
745 crate::mb16::DEFER_SUBPEL.store(
746 if on { 1 } else { 0 },
747 core::sync::atomic::Ordering::Relaxed,
748 );
749}
750
751pub fn set_split_t(t: u32) {
752 crate::mb16::SPLIT_T.store(t, core::sync::atomic::Ordering::Relaxed);
753}
754
755pub fn set_subpel_dispatch(on: bool) {
756 crate::mb16::SP_DISPATCH.store(
757 if on { 1 } else { 0 },
758 core::sync::atomic::Ordering::Relaxed,
759 );
760}
761
762pub fn set_subpel_pattern(p: u32) {
763 crate::mb16::SUBPEL_PAT.store(p, core::sync::atomic::Ordering::Relaxed);
764}
765
766impl Encoder {
767 /// Creates an encoder, validating that the configuration is within the
768 /// implemented subset.
769 pub fn new(cfg: EncoderConfig) -> Result<Self, EncodeError> {
770 if !matches!(
771 cfg.profile,
772 Profile::ConstrainedBaseline | Profile::Baseline | Profile::Main | Profile::High
773 ) {
774 return Err(EncodeError::Unsupported("unsupported profile"));
775 }
776 // The 8x8 transform is a High-profile feature, available under BOTH entropy
777 // coders since R6: `transform_size_8x8_flag` at both of its syntax positions
778 // (ctxIdxOffset 399) plus the ctxBlockCat-5 residual writer (sig 402 / last
779 // 417 / levels 426, and NO coded_block_flag -- presence comes from
780 // CodedBlockPatternLuma).
781 //
782 // CLAMPED, not refused. It is default-ON, so a caller who narrows the profile
783 // to Main is asking for Main-compatible output, not asking for an error -- and
784 // refusing here would make `EncoderConfig::new()` plus `profile = Main`, an
785 // entirely reasonable pair, fail outright. A profile is a compatibility
786 // ceiling; High-only tools clamp to it, exactly as `-profile:v` behaves
787 // elsewhere. The CLI promotes to High whenever `--transform-8x8 1` is passed,
788 // so an explicit request is never silently dropped there.
789 let mut cfg = cfg;
790 if cfg.transform_8x8 && !matches!(cfg.profile, Profile::High) {
791 cfg.transform_8x8 = false;
792 }
793 // B-frames are illegal in Baseline / Constrained Baseline (the decoder
794 // enforces this too). HIGH is a superset of Main and permits B slices —
795 // this used to demand Main exactly, which rejected the perfectly legal
796 // High + B-frames combination and blocked the 8x8 + B measurement.
797 // R6-5: 8x8 + B-frames. The B rule for transform_size_8x8_flag is derived
798 // in `plan_inter_mb` (`allow_t8`) and mirrored at the B emit: with
799 // direct_8x8_inference_flag = 0, a B_Direct_16x16 macroblock may not carry
800 // the flag, so the plan must not pick 8x8 for one. Gating at PLAN time (not
801 // emit time) is what keeps our reconstruction and the decoder's in step.
802 if cfg.bframes > 0 && !matches!(cfg.profile, Profile::Main | Profile::High) {
803 return Err(EncodeError::Unsupported(
804 "B-frames require Main or High profile",
805 ));
806 }
807 if cfg.chroma != ChromaFormat::Yuv420 {
808 return Err(EncodeError::Unsupported("only 4:2:0 chroma"));
809 }
810 if cfg.width == 0 || cfg.height == 0 || cfg.width % 2 != 0 || cfg.height % 2 != 0 {
811 return Err(EncodeError::Unsupported(
812 "dimensions must be positive and even",
813 ));
814 }
815 let sps = Sps::from_config(&cfg);
816 let pps = Pps::from_config(&cfg);
817 let rc = (cfg.bitrate > 0).then(|| RateControl::new(cfg.bitrate, cfg.framerate, cfg.qp));
818 let (cfg_w, cfg_h) = (cfg.width, cfg.height);
819 let sps_nal = sps.to_nal();
820 let pps_nal = pps.to_nal();
821 Ok(Self {
822 cfg,
823 sps,
824 pps,
825 frame_index: 0,
826 next_frame_num: 0,
827 gop_index: 0,
828 refs: Vec::new(),
829 rc,
830 pending_qpo: None,
831 pending_aq_probe: None,
832 la_queue: Vec::new(),
833 since_idr: 0,
834 force_idr: false,
835 keyframe_requested: false,
836 last_src: None,
837 last_prep: None,
838 cut_hist: [1.0, 1.0],
839 gather: None,
840 // Pre-size the slice writer to a generous fraction of the raw frame so
841 // the CAVLC hot loop never reallocs mid-frame (byte-identical; just capacity).
842 bw: BitWriter::with_capacity(cfg_w * cfg_h / 2 + 4096),
843 sps_nal,
844 pps_nal,
845 })
846 }
847
848 /// Sets the per-MB QP offset applied to the NEXT [`encode`](Self::encode) call
849 /// (mb-tree temporal AQ). One entry per macroblock (raster). Consumed once.
850 pub(crate) fn set_pending_qpo(&mut self, qpo: Vec<i32>) {
851 self.pending_qpo = Some(qpo);
852 }
853
854 /// Sets the AQ grain probe (the previous display-order SOURCE frame) for the
855 /// NEXT call if it codes an IDR. Consumed once per frame either way.
856 pub(crate) fn set_aq_probe(&mut self, f: YuvFrame) {
857 self.pending_aq_probe = Some(f);
858 }
859
860 /// The active configuration.
861 pub fn config(&self) -> &EncoderConfig {
862 &self.cfg
863 }
864
865 /// Encodes one frame, returning zero or more Annex-B access units. IDR
866 /// placement follows the keyint model (`gop_size` ceiling, `min_keyint`,
867 /// scene cuts), each IDR prefixed with SPS/PPS.
868 ///
869 /// With a lookahead feature active (mb-tree, on by default) frames buffer
870 /// until a window fills, so a call may return EMPTY bytes — call
871 /// [`flush`](Self::flush) at end of stream to drain the tail.
872 pub fn encode(&mut self, frame: &YuvFrame) -> Vec<u8> {
873 self.try_encode(frame).expect("frame matched config")
874 }
875
876 /// [`encode`](Self::encode) over a **borrowed** frame — the camera's DMA
877 /// buffer, not a copy. A tight view feeds the coder directly; a padded one
878 /// is gathered into a reused scratch frame first. Byte-identical to
879 /// `encode(&view.to_frame())`.
880 pub fn encode_planes(&mut self, frame: &YuvPlanes<'_>) -> Result<Vec<u8>, EncodeError> {
881 if frame.width != self.cfg.width || frame.height != self.cfg.height {
882 return Err(EncodeError::FrameMismatch);
883 }
884 if frame.is_valid() {
885 return self.try_encode_planes(frame);
886 }
887 let mut g = self
888 .gather
889 .take()
890 .unwrap_or_else(|| YuvFrame::black(frame.width, frame.height));
891 frame.copy_into(&mut g);
892 let r = self.try_encode_planes(&g.as_planes());
893 self.gather = Some(g);
894 r
895 }
896
897 /// [`encode_planes`](Self::encode_planes) into a caller-owned buffer — the
898 /// packetizer's, on a chip — returning the access unit's length.
899 ///
900 /// An access unit that does not fit is
901 /// [`EncodeError::BufferTooSmall`] with the size it needed; the picture is
902 /// then lost (the encoder has already coded it and moved on) and the next
903 /// call still works, so size the buffer for the worst case: an IDR at low
904 /// QP can approach `width * height * 3 / 2`. Bytes are identical to
905 /// [`encode_planes`].
906 pub fn encode_into(
907 &mut self,
908 frame: &YuvPlanes<'_>,
909 out: &mut [u8],
910 ) -> Result<usize, EncodeError> {
911 if frame.width != self.cfg.width || frame.height != self.cfg.height {
912 return Err(EncodeError::FrameMismatch);
913 }
914 // A configuration that buffers (lookahead) or looks at frame pairs
915 // (scene cut) hands out access units on its own schedule; those go
916 // through the `Vec` path and are copied. The chip configuration —
917 // `baseline()`: no lookahead, no scene cut — is written in place:
918 // no access-unit `Vec`, no second copy, the slice writer's buffer
919 // reused from frame to frame.
920 let buffered = self.lookahead_active() || (self.cfg.scenecut > 0 && self.cfg.gop_size > 1);
921 if buffered {
922 let au = self.encode_planes(frame)?;
923 return copy_out(&au, out);
924 }
925 if core::mem::take(&mut self.keyframe_requested) {
926 self.force_idr = true;
927 }
928 let mut sink = Sink::Slice {
929 buf: out,
930 len: 0,
931 needed: 0,
932 };
933 if frame.is_valid() {
934 self.encode_direct_into(frame, &mut sink)?;
935 } else {
936 let mut g = self
937 .gather
938 .take()
939 .unwrap_or_else(|| YuvFrame::black(frame.width, frame.height));
940 frame.copy_into(&mut g);
941 let r = self.encode_direct_into(&g.as_planes(), &mut sink);
942 self.gather = Some(g);
943 r?;
944 }
945 match sink {
946 Sink::Slice { len, needed, .. } if needed > len => {
947 Err(EncodeError::BufferTooSmall { needed })
948 }
949 Sink::Slice { len, .. } => Ok(len),
950 Sink::Vec(_) => unreachable!("the direct path writes to the slice sink"),
951 }
952 }
953
954 /// [`flush`](Self::flush) into a caller-owned buffer; see
955 /// [`encode_into`](Self::encode_into) for the buffer contract.
956 pub fn flush_into(&mut self, out: &mut [u8]) -> Result<usize, EncodeError> {
957 let tail = self.try_flush()?;
958 copy_out(&tail, out)
959 }
960
961 /// Make the next picture an IDR (a late joiner on the mesh needs one now,
962 /// not at the next GOP boundary). Rate-control state, the frame counter and
963 /// the scene-cut history survive; only the DPB restarts, as at any IDR.
964 /// With a lookahead active the buffered frames are coded first and the IDR
965 /// lands on the next frame *submitted*, in coding order.
966 pub fn request_keyframe(&mut self) {
967 self.keyframe_requested = true;
968 }
969
970 /// Fallible [`encode`](Self::encode).
971 ///
972 /// With a lookahead feature active (currently mb-tree, on by default in the
973 /// constant-QP path) this BUFFERS: mb-tree needs a whole GOP of future frames
974 /// before it can assign any of their QPs, so the returned `Vec` is empty while
975 /// the GOP fills and then carries that entire GOP's access units at once. The
976 /// concatenation of every return value plus [`flush`](Self::flush) is exactly
977 /// what [`encode_all`](Self::encode_all) produces — byte for byte.
978 ///
979 /// **You must call [`flush`](Self::flush) at end of stream** or the final
980 /// partial GOP is never emitted. (A debug build asserts if the encoder is
981 /// dropped with frames still buffered.) For zero added latency set
982 /// `cfg.mbtree = false`, which restores one-AU-per-call behaviour.
983 pub fn try_encode(&mut self, frame: &YuvFrame) -> Result<Vec<u8>, EncodeError> {
984 self.try_encode_planes(&frame.as_planes())
985 }
986
987 /// The streaming entry over a tight view. Every `encode*` lands here.
988 fn try_encode_planes(&mut self, frame: &YuvPlanes<'_>) -> Result<Vec<u8>, EncodeError> {
989 // A keyframe request with frames still buffered: code them first (they
990 // predate the request; a scene-cut IDR pending for the first of them
991 // is spent by that drain, as it should be), then the request becomes
992 // THIS frame's forced IDR.
993 if core::mem::take(&mut self.keyframe_requested) {
994 let mut out = if self.la_queue.is_empty() {
995 Vec::new()
996 } else {
997 self.try_flush()?
998 };
999 self.force_idr = true;
1000 out.extend_from_slice(&self.try_encode_planes(frame)?);
1001 return Ok(out);
1002 }
1003 // CAUSAL scene-cut detection (x264 keyint parity), BEFORE any
1004 // buffering: the pair is (previous source, this source) — exactly the
1005 // pair `segment_gops` reads in the batch path, so streaming == batch
1006 // holds under cuts. On a cut: flush whatever the lookahead holds (an
1007 // mb-tree window must not straddle an IDR), then request the IDR and
1008 // hand the detector's retained frame over as the grain probe.
1009 if self.cfg.scenecut > 0 && self.cfg.gop_size > 1 {
1010 if self.last_src.is_some() {
1011 // History advances on EVERY pair (the spike baseline must see
1012 // the pairs before a legal cut position too — this is what
1013 // keeps streaming == batch, whose `segment_gops` scores the
1014 // same pairs). Two exact reductions, mirroring the batch
1015 // scorer's rolling lazy cursor:
1016 // * prep only the NEW frame — the previous frame's coded +
1017 // half-res planes were built by the last call (`last_prep`);
1018 // * skip the ratio entirely while no decision can read it: a
1019 // decision fires at counter >= min_keyint and consults this
1020 // pair's ratio for at most the two following frames, so a
1021 // pair at counter < min_keyint - 2 is UNREADABLE. Its slot
1022 // rolls a placeholder that is never consulted (at the
1023 // earliest decision, both baseline slots came from counters
1024 // >= min_keyint - 2 — computed).
1025 let counter = self.since_idr + self.la_queue.len() as u32;
1026 let minki = self.cfg.min_keyint.max(1);
1027 let r = if counter + 2 >= minki {
1028 let cur_prep = mbtree::pair_prep(&self.cfg, frame);
1029 let prev_prep = match self.last_prep.take() {
1030 Some(p) => p,
1031 None => mbtree::pair_prep(
1032 &self.cfg,
1033 &self.last_src.as_ref().expect("guarded").as_planes(),
1034 ),
1035 };
1036 let r = mbtree::pair_ratio_prepped(&self.cfg, &cur_prep, &prev_prep);
1037 self.last_prep = Some(cur_prep);
1038 r
1039 } else {
1040 self.last_prep = None;
1041 1.0 // placeholder: unreadable (see above)
1042 };
1043 let (p1, p2) = (self.cut_hist[0], self.cut_hist[1]);
1044 self.cut_hist = [r, p1];
1045 if counter >= minki && lookahead::is_scene_cut(&self.cfg, r, p1, p2) {
1046 let flushed = self.try_flush()?;
1047 self.force_idr = true;
1048 self.pending_aq_probe = self.last_src.take();
1049 self.last_src = Some(frame.to_frame());
1050 let mut out = flushed;
1051 out.extend_from_slice(&self.try_encode_inner(frame)?);
1052 return Ok(out);
1053 }
1054 }
1055 self.last_src = Some(frame.to_frame());
1056 }
1057 self.try_encode_inner(frame)
1058 }
1059
1060 fn try_encode_inner(&mut self, frame: &YuvPlanes<'_>) -> Result<Vec<u8>, EncodeError> {
1061 if !self.lookahead_active() {
1062 return self.encode_direct(frame);
1063 }
1064 if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
1065 return Err(EncodeError::FrameMismatch);
1066 }
1067 self.la_queue.push(frame.to_frame());
1068 // mb-tree's window: bounded by `lookahead` (x264's rc-lookahead — a
1069 // 250-frame keyint must not mean a 250-frame buffer) and by the
1070 // frames REMAINING until the forced IDR, so a window never straddles
1071 // an IDR and ends exactly where the batch path's segment does (the
1072 // streaming == batch identity under the keyint model). `since_idr %
1073 // gop` treats the just-completed GOP (`since_idr == gop`) as a fresh
1074 // one — the next frame IS the IDR.
1075 let gop = self.cfg.gop_size.max(1);
1076 let rem_to_idr = (gop - (self.since_idr % gop)) as usize;
1077 let window = (self.cfg.lookahead.max(1) as usize).min(rem_to_idr);
1078 if self.la_queue.len() >= window {
1079 self.emit_lookahead_gop()
1080 } else {
1081 Ok(Vec::new())
1082 }
1083 }
1084
1085 /// True when a feature needs future frames, so [`try_encode`] must buffer.
1086 /// B-frames already refuse the streaming API, and rate control drives its own
1087 /// sequential path, so mb-tree in constant-QP mode is the only case.
1088 fn lookahead_active(&self) -> bool {
1089 self.cfg.mbtree && self.cfg.bframes == 0 && self.cfg.bitrate == 0
1090 }
1091
1092 /// Codes every buffered frame with mb-tree's per-GOP QP offsets and returns
1093 /// their access units concatenated. Identical to what an `encode_all` worker
1094 /// does for the same GOP.
1095 fn emit_lookahead_gop(&mut self) -> Result<Vec<u8>, EncodeError> {
1096 let frames = core::mem::take(&mut self.la_queue);
1097 let offs = mbtree::gop_qp_offsets(&self.cfg, &frames, self.cfg.mbtree_strength);
1098 let mut out = Vec::new();
1099 for (i, f) in frames.iter().enumerate() {
1100 if let Some(o) = offs.get(i) {
1101 self.pending_qpo = Some(o.clone());
1102 }
1103 out.extend_from_slice(&self.encode_direct(&f.as_planes())?);
1104 }
1105 Ok(out)
1106 }
1107
1108 /// Emits any frames still held by the lookahead queue (end of stream).
1109 ///
1110 /// Returns the trailing access units, or empty when nothing is buffered — so it
1111 /// is always safe to call, including when no lookahead feature is active.
1112 pub fn flush(&mut self) -> Vec<u8> {
1113 self.try_flush()
1114 .expect("buffered frames matched the config when accepted")
1115 }
1116
1117 /// Fallible [`flush`](Self::flush).
1118 pub fn try_flush(&mut self) -> Result<Vec<u8>, EncodeError> {
1119 if self.la_queue.is_empty() {
1120 return Ok(Vec::new());
1121 }
1122 self.emit_lookahead_gop()
1123 }
1124
1125 /// The unbuffered single-frame path: codes `frame` immediately. This is what the
1126 /// batch path's workers call, since they compute the lookahead themselves.
1127 fn encode_direct(&mut self, frame: &YuvPlanes<'_>) -> Result<Vec<u8>, EncodeError> {
1128 let mut out = Vec::new();
1129 self.encode_direct_into(frame, &mut Sink::Vec(&mut out))?;
1130 Ok(out)
1131 }
1132
1133 /// [`encode_direct`](Self::encode_direct) into a [`Sink`]: the one body
1134 /// behind both the `Vec` and the caller-buffer paths.
1135 fn encode_direct_into(
1136 &mut self,
1137 frame: &YuvPlanes<'_>,
1138 out: &mut Sink<'_>,
1139 ) -> Result<(), EncodeError> {
1140 let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Total);
1141 if frame.width != self.cfg.width || frame.height != self.cfg.height || !frame.is_valid() {
1142 return Err(EncodeError::FrameMismatch);
1143 }
1144
1145 // B-frames need lookahead (a future anchor coded before the B), which the
1146 // one-frame-in streaming API can't provide — use `encode_all` for B.
1147 if self.cfg.bframes > 0 {
1148 return Err(EncodeError::Unsupported(
1149 "B-frames need encode_all (lookahead)",
1150 ));
1151 }
1152 // GOP placement (x264 keyint model): an IDR when forced (scene cut or
1153 // batch segment boundary), at stream start / after reset, or when the
1154 // since-IDR counter reaches `gop_size` (the keyint ceiling). The
1155 // counter replaced `frame_index % gop_size` — cadence is not periodic
1156 // once cuts place IDRs; with `scenecut = 0` the counter reproduces the
1157 // modulo exactly (the bisection anchor).
1158 let forced = core::mem::take(&mut self.force_idr);
1159 let is_idr = self.cfg.gop_size <= 1
1160 || forced
1161 || self.since_idr == 0
1162 || self.since_idr >= self.cfg.gop_size;
1163 if is_idr {
1164 self.gop_index = 0;
1165 self.next_frame_num = 0;
1166 self.refs.clear();
1167 }
1168 let frame_num = self.next_frame_num;
1169 let poc_lsb = (2 * self.gop_index) % 256;
1170 // mb-tree per-MB QP offset for this frame (empty = none / byte-identical).
1171 let qpo = self.pending_qpo.take().unwrap_or_default();
1172 // AQ grain probe (IDR only; consumed unconditionally so it cannot go stale).
1173 let aq_probe_owned = self.pending_aq_probe.take();
1174 let aq_probe = aq_probe_owned.as_ref().map(YuvFrame::as_planes);
1175
1176 // Rate control (if enabled) chooses this frame's QP from a cheap
1177 // look-ahead complexity estimate; otherwise the QP is fixed.
1178 let complexity = if self.rc.is_some() {
1179 lookahead::complexity(
1180 &self.cfg,
1181 frame,
1182 if is_idr { None } else { self.refs.first() },
1183 )
1184 } else {
1185 0.0
1186 };
1187 let qp = match &self.rc {
1188 Some(rc) => rc.pick_qp(is_idr, complexity),
1189 // Constant-QP: apply the per-GOP I-frame cascade offset (0 by default →
1190 // byte-identical). Keeps the P-only path consistent with `code_picture`.
1191 None if is_idr => (self.cfg.qp as i32 + self.cfg.i_qp_offset).clamp(0, 51) as u8,
1192 None => self.cfg.qp,
1193 };
1194
1195 // The reused slice writer: its buffer was sized once in `new`.
1196 let mut w = core::mem::take(&mut self.bw);
1197 w.clear();
1198 let (nal_type, mut reference) = if is_idr {
1199 // SPS/PPS precede every IDR so the stream is independently decodable.
1200 out.nal(
1201 self.sps_nal.ref_idc,
1202 self.sps_nal.nal_type,
1203 &self.sps_nal.rbsp,
1204 );
1205 out.nal(
1206 self.pps_nal.ref_idc,
1207 self.pps_nal.nal_type,
1208 &self.pps_nal.rbsp,
1209 );
1210 slice::write_idr_slice_header(&mut w, &self.cfg, qp);
1211 // The batch paths park the previous source frame in
1212 // `pending_aq_probe` (taken above); pure streaming callers have no
1213 // previous frame retained — the grain veto fails open there.
1214 let r = if self.cfg.cabac {
1215 mb16::encode_slice_data_cabac_intra(
1216 &mut w,
1217 &self.cfg,
1218 frame,
1219 qp,
1220 &qpo,
1221 aq_probe.as_ref(),
1222 )
1223 } else {
1224 mb16::encode_slice_data(
1225 &mut w,
1226 &self.cfg,
1227 frame,
1228 qp,
1229 false,
1230 &[],
1231 &qpo,
1232 aq_probe.as_ref(),
1233 &[],
1234 )
1235 };
1236 (NalUnitType::IdrSlice, r)
1237 } else {
1238 // weightp (x264 parity): estimate per-reference luma weights; the
1239 // slice header's table and the coder's post-MC apply must be the
1240 // SAME list or the stream and the recon disagree.
1241 let wp: Vec<(i32, i32)> = if self.cfg.weightp {
1242 estimate_luma_weights(&self.cfg, frame, &self.refs)
1243 } else {
1244 Vec::new()
1245 };
1246 let wp_hdr = if self.cfg.weightp {
1247 Some(wp.as_slice())
1248 } else {
1249 None
1250 };
1251 slice::write_p_slice_header(
1252 &mut w,
1253 &self.cfg,
1254 qp,
1255 frame_num,
1256 poc_lsb,
1257 self.refs.len(),
1258 wp_hdr,
1259 );
1260 let r = if self.cfg.cabac {
1261 mb16::encode_slice_data_cabac_p(&mut w, &self.cfg, frame, qp, &self.refs, &qpo, &wp)
1262 } else {
1263 mb16::encode_slice_data(
1264 &mut w, &self.cfg, frame, qp, true, &self.refs, &qpo, None, &wp,
1265 )
1266 };
1267 (NalUnitType::NonIdrSlice, r)
1268 };
1269 // POC/frame_num carried on the reference so B-frame ref-lists (when enabled)
1270 // can order L0/L1 by display position. Unused on the P-only path.
1271 reference.poc = 2 * self.gop_index as i32;
1272 reference.frame_num = frame_num;
1273 let slice_bytes = w.finish();
1274 // Feed the coded slice size (the picture's own bits) back to the controller.
1275 if let Some(rc) = &mut self.rc {
1276 rc.update(is_idr, slice_bytes.len() * 8, qp, complexity);
1277 }
1278 {
1279 let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::EncNal);
1280 out.nal(3, nal_type, slice_bytes);
1281 }
1282 self.bw = w;
1283
1284 // The deblocked reconstruction enters the DPB (most-recent first), which
1285 // is kept to `max_num_ref_frames` by a sliding window.
1286 self.refs.insert(0, reference);
1287 self.refs.truncate(self.cfg.num_ref_frames.max(1) as usize);
1288 self.frame_index += 1;
1289 self.gop_index += 1;
1290 self.next_frame_num = (self.next_frame_num + 1) % 16;
1291 // SELF-FILL the AQ grain probe: if this encoder's NEXT frame opens a GOP,
1292 // retain THIS source frame as its probe (one clone per GOP, none on other
1293 // frames). Sequential paths (streaming, RC) get IDR coverage for free and
1294 // byte-match the batch workers, which set the same frame externally
1295 // (fresh encoder per GOP — no state to self-fill from).
1296 self.since_idr = if is_idr { 1 } else { self.since_idr + 1 };
1297 if self.cfg.gop_size > 1 && self.since_idr >= self.cfg.gop_size {
1298 self.pending_aq_probe = Some(frame.to_frame());
1299 }
1300 Ok(())
1301 }
1302
1303 /// Batch-encodes every frame, returning one Annex-B access unit per frame.
1304 ///
1305 /// At constant QP the GOPs are independent — each begins with an IDR that
1306 /// resets the DPB, `frame_num` and POC, and SPS/PPS precede every IDR — so they
1307 /// are encoded **in parallel across CPU cores** and the result is
1308 /// **byte-identical** to calling [`encode`](Self::encode) frame-by-frame. With
1309 /// rate control enabled the per-frame QP depends on history, so this falls back
1310 /// to sequential encoding. Within a GOP, P-frames are inherently sequential
1311 /// (each predicts from the previous reconstruction); the parallelism is across
1312 /// GOPs, so it scales with the number of GOPs in the clip.
1313 pub fn encode_all(&self, frames: &[YuvFrame]) -> Result<Vec<Vec<u8>>, EncodeError> {
1314 for f in frames {
1315 if f.width != self.cfg.width || f.height != self.cfg.height || !f.is_valid() {
1316 return Err(EncodeError::FrameMismatch);
1317 }
1318 }
1319 // SUB-PEL GRAIN VETO, decided ONCE for the sequence (grain is a property of the
1320 // SOURCE, not of a frame). Sub-pel interpolates, and on grain it interpolates
1321 // NOISE. `RFF_GRAIN_SUBPEL=0` opts out.
1322 let grain_seq = self.cfg.preset != Preset::Fast
1323 && rusty_h264_common::knob("RFF_GRAIN_SUBPEL")
1324 .map(|v| v != "0")
1325 .unwrap_or(true)
1326 && frames.len() >= 2
1327 && {
1328 // CODED-size planes, not the raw display planes: FrameSignals
1329 // walks the MB grid (`mb_h*16` rows), and a display height
1330 // that is not a multiple of 16 (1080p!) is SHORTER than that —
1331 // this exact line crashed on blue_sky_1080p the first time a
1332 // non-MB-multiple clip went through `encode_all` (pre-existing;
1333 // surfaced by the bframes-v2 holdout run).
1334 let cl0 = mbtree::coded_luma(&self.cfg, &frames[0].as_planes());
1335 let cl1 = mbtree::coded_luma(&self.cfg, &frames[1].as_planes());
1336 crate::signals::FrameSignals::new(
1337 &cl1,
1338 self.cfg.mb_width() * 16,
1339 self.cfg.mb_width(),
1340 self.cfg.mb_height(),
1341 Some(&cl0),
1342 )
1343 .grain_signature()
1344 };
1345 let _guard = mb16::SeqFastPath::set(grain_seq);
1346 // B-frames need a reorder pipeline (code the future anchor before the B's
1347 // that reference it) — a separate sequential path.
1348 if self.cfg.bframes > 0 {
1349 // Content-adaptive dispatch, PER GOP (codec-content-adaptive-dispatch):
1350 // code B-frames only in GOPs whose motion is predictable enough to pay,
1351 // so a mixed clip gets B on its smooth segments and P on its busy ones.
1352 // Scene-cut segmentation first (x264 keyint model) — every per-GOP
1353 // signal below is per-SEGMENT, so the dispatch decides on real
1354 // scene units instead of arbitrary fixed windows.
1355 let seg_starts = lookahead::segment_gops(&self.cfg, frames);
1356 let seg_range = |k: usize| {
1357 (
1358 seg_starts[k],
1359 seg_starts.get(k + 1).copied().unwrap_or(frames.len()),
1360 )
1361 };
1362 let n_gops = seg_starts.len();
1363 let (w, h) = (self.cfg.width, self.cfg.height);
1364 // One cheap per-GOP signal drives BOTH content-adaptive knobs: the B/P
1365 // structure dispatch AND the I-frame QP-cascade depth.
1366 let gop_sig: Vec<f64> = (0..n_gops)
1367 .map(|g| {
1368 let (s, e) = seg_range(g);
1369 gop_bi_residual(&frames[s..e], w, h, 1)
1370 })
1371 .collect();
1372 // bframes-v2 dispatch: the SEGMENT level carries only the vetoes
1373 // (screen content — the +1.19% leak class, whose bi-residual is
1374 // near zero and sails under any threshold); the bi-residual
1375 // decision itself moved to the ANCHOR-GAP level inside
1376 // `encode_all_bframes`, where crew-class episodic content (flash
1377 // gaps P, calm gaps B) is separable in a way no clip-level scalar
1378 // was.
1379 let gop_fav: Vec<bool> = if self.cfg.bframes_adaptive {
1380 (0..n_gops)
1381 .map(|g| {
1382 let (s, e) = seg_range(g);
1383 if e - s < 2 {
1384 return bframes_favorable(gop_sig[g]);
1385 }
1386 let cl0 = mbtree::coded_luma(&self.cfg, &frames[s].as_planes());
1387 let cl1 = mbtree::coded_luma(&self.cfg, &frames[s + 1].as_planes());
1388 let cw = self.cfg.mb_width() * 16;
1389 let sig = signals::FrameSignals::new(
1390 &cl1,
1391 cw,
1392 self.cfg.mb_width(),
1393 self.cfg.mb_height(),
1394 Some(&cl0),
1395 );
1396 !sig.is_screen()
1397 })
1398 .collect()
1399 } else {
1400 vec![true; n_gops]
1401 };
1402 let gop_iqp: Vec<i32> = gop_sig
1403 .iter()
1404 .map(|&s| gop_iqp_offset(s, self.cfg.i_qp_offset))
1405 .collect();
1406 let gop_bqp: Vec<i32> = gop_sig
1407 .iter()
1408 .map(|&s| gop_bframe_qp_offset(s, self.cfg.bframe_qp_offset))
1409 .collect();
1410 // Adaptive B-COUNT: how many B's per anchor gap. Fixed `bframes` unless
1411 // `auto`, where the 2-gap/1-gap bi-residual RATIO picks it — content that
1412 // survives wider anchor spacing (low ratio) carries more cheap B's; simple
1413 // translation (high ratio) wants a single equidistant B.
1414 let bcount = if self.cfg.bframes_adaptive {
1415 adaptive_bcount(frames, w, h, self.cfg.bframes as usize)
1416 } else {
1417 self.cfg.bframes as usize
1418 };
1419 if gop_fav.iter().any(|&f| f) {
1420 return Ok(self.encode_all_bframes(
1421 frames,
1422 bcount,
1423 &seg_starts,
1424 &gop_fav,
1425 &gop_iqp,
1426 &gop_bqp,
1427 ));
1428 }
1429 // No GOP is B-favorable → pure P-only (byte-identical to bframes=0).
1430 let mut pcfg = self.cfg.clone();
1431 pcfg.bframes = 0;
1432 return Encoder::new(pcfg)?.encode_all(frames);
1433 }
1434 // Rate control threads state across frames → it must stay sequential. mb-tree
1435 // runs in RC mode too: per-GOP lookahead → per-MB offsets (per-GOP centered, so
1436 // rate-neutral per GOP), and the controller supplies each frame's base QP.
1437 // (MEASURED: routing the cross-frame allocation through the RC's complexity
1438 // instead of centering was worse — the centered offsets carry it correctly.)
1439 if self.cfg.bitrate > 0 {
1440 let mut enc = Encoder::new(self.cfg.clone())?;
1441 // Same segmentation as the CQP path (scene-cut IDRs under the
1442 // keyint ceiling), mb-tree windowed to `lookahead` within each
1443 // segment; segment starts force the IDR through the controller's
1444 // sequential encoder.
1445 let seg_starts = lookahead::segment_gops(&self.cfg, frames);
1446 let offs: Vec<Vec<i32>> = if self.cfg.mbtree {
1447 seg_starts
1448 .iter()
1449 .enumerate()
1450 .flat_map(|(k, &s)| {
1451 let end = seg_starts.get(k + 1).copied().unwrap_or(frames.len());
1452 frames[s..end].chunks(self.cfg.lookahead.max(1) as usize)
1453 })
1454 .flat_map(|w| mbtree::gop_qp_offsets(&self.cfg, w, self.cfg.mbtree_strength))
1455 .collect()
1456 } else {
1457 Vec::new()
1458 };
1459 let mut next_seg = 1usize; // seg_starts[0] == 0 is the natural first IDR
1460 return frames
1461 .iter()
1462 .enumerate()
1463 .map(|(i, f)| {
1464 if seg_starts.get(next_seg) == Some(&i) {
1465 enc.force_idr = true;
1466 next_seg += 1;
1467 }
1468 if let Some(qpo) = offs.get(i) {
1469 enc.pending_qpo = Some(qpo.clone());
1470 }
1471 // Bypass the streaming lookahead buffer: this path supplies the
1472 // offsets itself, so buffering here would double-compute them.
1473 enc.encode_direct(&f.as_planes())
1474 })
1475 .collect();
1476 }
1477 // Variable GOP segmentation (x264 keyint model): scene-cut driven IDRs
1478 // under the `gop_size` ceiling. With `scenecut = 0` this IS
1479 // `chunks(gop_size)` — byte-identical to the fixed-cadence encoder.
1480 let seg_starts = lookahead::segment_gops(&self.cfg, frames);
1481 let gops: Vec<&[YuvFrame]> = seg_starts
1482 .iter()
1483 .enumerate()
1484 .map(|(k, &s)| &frames[s..seg_starts.get(k + 1).copied().unwrap_or(frames.len())])
1485 .collect();
1486 if gops.is_empty() || frames.is_empty() {
1487 return Ok(Vec::new());
1488 }
1489 // Each GOP is encoded with a fresh encoder (an IDR resets all state), so
1490 // the GOPs are independent: with `std` they distribute across worker
1491 // threads, without it they run in order. Same closure, same bytes.
1492 let mut out: Vec<Option<Vec<Vec<u8>>>> = (0..gops.len()).map(|_| None).collect();
1493 let cfg = &self.cfg;
1494 let gops_ref = &gops;
1495 let encode_gop = |i: usize| -> Vec<Vec<u8>> {
1496 let mut enc = Encoder::new(cfg.clone()).expect("config");
1497 // mb-tree temporal AQ: a per-GOP lookahead over the GOP's
1498 // source frames yields per-frame per-MB QP offsets (the GOP
1499 // is the natural window — the IDR resets references). Off →
1500 // empty → byte-identical.
1501 // mb-tree windows of ≤ `lookahead` frames within
1502 // the segment (a 250-frame scenecut GOP must not
1503 // be one propagation window); aligned to the
1504 // segment start, exactly like the streaming path.
1505 let offs: Vec<Vec<i32>> = if cfg.mbtree {
1506 gops_ref[i]
1507 .chunks(cfg.lookahead.max(1) as usize)
1508 .flat_map(|w| mbtree::gop_qp_offsets(cfg, w, cfg.mbtree_strength))
1509 .collect()
1510 } else {
1511 Vec::new()
1512 };
1513 // The GOP's IDR probes grain against the PREVIOUS GOP's
1514 // last source frame — `gops_ref` is the full shared
1515 // slice, so this is identical under any thread count,
1516 // and it is exactly the frame `encode_direct`'s
1517 // self-fill would have retained in a sequential run
1518 // (the documented streaming==batch invariant). The
1519 // stream's FIRST IDR fails open on every path — pure
1520 // streaming cannot see frame 1 at frame 0.
1521 if i > 0 {
1522 if let Some(pf) = gops_ref[i - 1].last() {
1523 enc.set_aq_probe(pf.clone());
1524 }
1525 }
1526 let aus: Vec<Vec<u8>> = gops_ref[i]
1527 .iter()
1528 .enumerate()
1529 .map(|(fi, f)| {
1530 if let Some(o) = offs.get(fi) {
1531 enc.set_pending_qpo(o.clone());
1532 }
1533 enc.encode_direct(&f.as_planes())
1534 .expect("frame matched config")
1535 })
1536 .collect();
1537 aus
1538 };
1539 #[cfg(feature = "std")]
1540 {
1541 let n = rusty_h264_common::knob("RUSTY_THREADS")
1542 .and_then(|v| v.parse().ok())
1543 .or_else(|| std::thread::available_parallelism().map(|n| n.get()).ok())
1544 .unwrap_or(1)
1545 .min(gops.len());
1546 let encode_gop = &encode_gop;
1547 std::thread::scope(|s| {
1548 let handles: Vec<_> = (0..n)
1549 .map(|t| {
1550 s.spawn(move || {
1551 // The flag is thread-local; each worker carries
1552 // the caller's decision.
1553 let _fast = mb16::SeqFastPath::set(grain_seq);
1554 let mut local = Vec::new();
1555 let mut i = t;
1556 while i < gops_ref.len() {
1557 local.push((i, encode_gop(i)));
1558 i += n;
1559 }
1560 local
1561 })
1562 })
1563 .collect();
1564 for h in handles {
1565 for (i, aus) in h.join().expect("encode worker panicked") {
1566 out[i] = Some(aus);
1567 }
1568 }
1569 });
1570 }
1571 #[cfg(not(feature = "std"))]
1572 {
1573 for i in 0..gops.len() {
1574 out[i] = Some(encode_gop(i));
1575 }
1576 }
1577 Ok(out.into_iter().flatten().flatten().collect())
1578 }
1579
1580 /// B-frame reorder pipeline (sequential). Produces access units in **coding
1581 /// order** (the decoder reorders to display order by POC). Structure: an IDR
1582 /// at each `gop_size` boundary, a P anchor every `bframes+1` frames within a
1583 /// GOP, `bframes` non-reference B-frames between consecutive anchors, and the
1584 /// last frame forced to an anchor so trailing B's always have a future
1585 /// reference. Each anchor is coded before the B's that reference it.
1586 /// `gop_favorable[g]` (content-adaptive): GOP `g` codes B-frames only when
1587 /// `true`; a `false` GOP is coded all-P (every frame an anchor) so busy segments
1588 /// of a mixed clip don't regress. Non-adaptive callers pass all-`true`.
1589 fn encode_all_bframes(
1590 &self,
1591 frames: &[YuvFrame],
1592 bcount: usize,
1593 seg_starts: &[usize],
1594 gop_favorable: &[bool],
1595 gop_iqp: &[i32],
1596 gop_bqp: &[i32],
1597 ) -> Vec<Vec<u8>> {
1598 let n = frames.len();
1599 if n == 0 {
1600 return Vec::new();
1601 }
1602 let step = bcount.max(1) + 1; // B's per anchor gap + 1 (adaptive in `auto`)
1603 // Per-display-index segment map (variable GOPs under scenecut; with
1604 // scenecut=0 the segments are the old fixed `gop_size` chunks and every
1605 // derived quantity below reproduces the old `% gop` arithmetic).
1606 let mut seg_of = vec![0usize; n];
1607 let mut seg_start_of = vec![0usize; n];
1608 {
1609 let mut k = 0usize;
1610 for d in 0..n {
1611 if k + 1 < seg_starts.len() && seg_starts[k + 1] == d {
1612 k += 1;
1613 }
1614 seg_of[d] = k;
1615 seg_start_of[d] = seg_starts[k];
1616 }
1617 }
1618 // A B-capable config: Main profile + ≥2 refs so the DPB holds both anchors.
1619 let mut cfg = self.cfg.clone();
1620 cfg.num_ref_frames = cfg.num_ref_frames.max(if cfg.b_pyramid { 3 } else { 2 }); // pyramid: 2 anchors + 1 B-ref must fit the DPB
1621 let sps = Sps::from_config(&cfg);
1622 let pps = Pps::from_config(&cfg);
1623
1624 // Anchor display-indices: IDR at GOP starts, P anchors every `step`, plus
1625 // the frame right before each IDR boundary and the clip's last frame — a
1626 // trailing B with no future reference IN ITS OWN GOP would otherwise be
1627 // coded after the next GOP's IDR (which clears the DPB), losing its anchors.
1628 let mut is_anchor = vec![false; n];
1629 // Segment-OUTER derivation: the favorability lookup, segment start and
1630 // next-boundary test hoist to once per SEGMENT (each was a bounds-
1631 // checked lookup per frame), and `off % step` — an integer divide per
1632 // frame, `step` being runtime-valued — becomes a rolling phase counter
1633 // (`phase == 0` exactly when `(d - seg_start) % step == 0`). A
1634 // non-favorable GOP is coded all-P (every frame an anchor); a
1635 // favorable one uses the B structure.
1636 for (k, &s) in seg_starts.iter().enumerate() {
1637 let e = seg_starts.get(k + 1).copied().unwrap_or(n).min(n);
1638 if !gop_favorable.get(k).copied().unwrap_or(true) {
1639 is_anchor[s..e].fill(true);
1640 continue;
1641 }
1642 // `d + 1 == next segment start` marked the frame right before each
1643 // IDR boundary; the clip's own end is handled by the line below.
1644 let boundary = k + 1 < seg_starts.len();
1645 let mut phase = 0usize;
1646 for d in s..e {
1647 is_anchor[d] = phase == 0 || (boundary && d + 1 == e);
1648 phase += 1;
1649 if phase == step {
1650 phase = 0;
1651 }
1652 }
1653 }
1654 is_anchor[n - 1] = true;
1655 // bframes-v2: PER-GAP favorability (adaptive mode only). Each anchor
1656 // gap is priced by its OWN bi-prediction residual; an unfavorable gap
1657 // codes all-P while its neighbours keep their B's — episodic content
1658 // (crew's camera flashes, one busy passage of a calm clip) dispatches
1659 // at the scale the phenomenon actually has. Fixed `--bframes N` keeps
1660 // x264's flat structure.
1661 if self.cfg.bframes_adaptive {
1662 let (w, h) = (self.cfg.width, self.cfg.height);
1663 // Frame means, MEMOIZED across pairs and gaps: consecutive pairs
1664 // share a frame and adjacent gaps share their anchor, so the eager
1665 // per-pair closure computed every interior mean twice (and anchor
1666 // means once per adjoining gap). NaN = not yet computed; the
1667 // sample count is `ceil(len/64)` — exactly what the counting loop
1668 // produced. Same sums, same divide: BIT-IDENTICAL, and the
1669 // short-circuiting `any` still computes no mean it never reads.
1670 fn memo_mean(frames: &[YuvFrame], means: &mut [f64], x: usize) -> f64 {
1671 if means[x].is_nan() {
1672 let f = &frames[x];
1673 let mut s = 0u64;
1674 let mut i = 0;
1675 while i < f.y.len() {
1676 s += f.y[i] as u64;
1677 i += 64;
1678 }
1679 let c = f.y.len().div_ceil(64) as u64;
1680 means[x] = s as f64 / c.max(1) as f64;
1681 }
1682 means[x]
1683 }
1684 let mut means = vec![f64::NAN; n];
1685 let mut a = 0usize;
1686 while a + 1 < n {
1687 if !is_anchor[a] {
1688 a += 1;
1689 continue;
1690 }
1691 // The gap = frames (a, next_anchor); slice inclusive of both ends.
1692 let mut b = a + 1;
1693 while b < n && !is_anchor[b] {
1694 b += 1;
1695 }
1696 if b > a + 1 && b < n {
1697 // FLASH VETO, per pair: a camera flash is a global DC jump
1698 // between adjacent frames — B-averaging across it blends
1699 // two exposures (crew's +5.72% at fixed B; the bi-residual
1700 // alone let its calmer flash gaps through at +3.34%). A
1701 // subsampled mean-luma delta is ~free and pair-precise in
1702 // a way no clip-level scalar was. Gradual fades move ~1
1703 // level per frame and stay far under the threshold.
1704 let flash = (a..b).any(|x| {
1705 (memo_mean(frames, &mut means, x) - memo_mean(frames, &mut means, x + 1))
1706 .abs()
1707 > 2.5
1708 });
1709 let res = gop_bi_residual(&frames[a..=b], w, h, 1);
1710 if flash || !bframes_favorable(res) {
1711 for x in a + 1..b {
1712 is_anchor[x] = true;
1713 }
1714 }
1715 }
1716 a = b;
1717 }
1718 }
1719
1720 // mb-tree temporal AQ over the ANCHOR reference chain: B-frames are
1721 // non-reference leaves (mb-tree offsets them at ~0 anyway), so the lookahead
1722 // runs over each GOP's anchor sub-sequence — the frames that actually form the
1723 // reference chain — and only anchors receive an offset. `mbtree_off[d]` is that
1724 // anchor's per-MB offset (empty for B's / when off → byte-identical).
1725 let mbtree_off: Vec<Vec<i32>> = if cfg.mbtree {
1726 let mut off = vec![Vec::new(); n];
1727 for (k, &s) in seg_starts.iter().enumerate() {
1728 let seg_end = seg_starts.get(k + 1).copied().unwrap_or(n).min(n);
1729 let anchors: Vec<usize> = (s..seg_end).filter(|&d| is_anchor[d]).collect();
1730 // Anchor chains windowed to `lookahead` (a 250-frame scenecut
1731 // segment's chain must not be one propagation window).
1732 for aw in anchors.chunks(cfg.lookahead.max(1) as usize) {
1733 // Borrowed window (gop_qp_offsets_refs): no per-anchor frame
1734 // deep-clones. `offs` is owned — MOVE the rows into place
1735 // instead of cloning a per-MB Vec per anchor.
1736 let aframes: Vec<&YuvFrame> = aw.iter().map(|&d| &frames[d]).collect();
1737 let offs = mbtree::gop_qp_offsets_refs(&cfg, &aframes, cfg.mbtree_strength);
1738 for (o, &d) in offs.into_iter().zip(aw.iter()) {
1739 off[d] = o;
1740 }
1741 }
1742 }
1743 off
1744 } else {
1745 Vec::new()
1746 };
1747
1748 // b-pyramid (x264 `normal` parity): in gaps carrying 2+ B's, the
1749 // display-middle B becomes a REFERENCE — coded right after its future
1750 // anchor, entered into the DPB, so the remaining leaves bracket
1751 // against it (the nearest-POC L0/L1 selection finds it naturally) at
1752 // half the prediction distance. CABAC-path v1.
1753 let pyramid = cfg.b_pyramid && cfg.cabac && step >= 3;
1754 let mut is_bref = vec![false; n];
1755 // Coding order: each anchor (display order), then — pyramid — the
1756 // gap's reference B, then the leaf B's.
1757 let mut order: Vec<usize> = Vec::with_capacity(n);
1758 let mut prev: Option<usize> = None;
1759 for d in 0..n {
1760 if !is_anchor[d] {
1761 continue;
1762 }
1763 order.push(d);
1764 if let Some(p) = prev {
1765 if pyramid && d - p > 2 {
1766 let m = (p + d) / 2;
1767 is_bref[m] = true;
1768 order.push(m);
1769 order.extend(((p + 1)..d).filter(|&x| x != m));
1770 } else {
1771 order.extend((p + 1)..d);
1772 }
1773 }
1774 prev = Some(d);
1775 }
1776
1777 let mut dpb: Vec<RefFrame> = Vec::new();
1778 let mut aus: Vec<Vec<u8>> = Vec::with_capacity(n);
1779 let mut frame_num: u32 = 0;
1780 for &d in &order {
1781 let is_idr = d == seg_start_of[d];
1782 if is_idr {
1783 dpb.clear();
1784 frame_num = 0;
1785 }
1786 let is_b = !is_anchor[d];
1787 let bref = is_b && is_bref[d];
1788 let poc = ((d - seg_start_of[d]) as i32) * 2; // POC = display position within the GOP
1789 let iqp = gop_iqp.get(seg_of[d]).copied().unwrap_or(cfg.i_qp_offset);
1790 let bqp_leaf = gop_bqp
1791 .get(seg_of[d])
1792 .copied()
1793 .unwrap_or(cfg.bframe_qp_offset);
1794 // A REFERENCE B must not take the full "quantize harder, nothing
1795 // depends on it" leaf offset — leaves predict FROM it. Half, like
1796 // x264's pyramid B-ref QP sitting between P and leaf-B.
1797 let bqp = if bref { (bqp_leaf + 1) / 2 } else { bqp_leaf };
1798 let qpo: &[i32] = mbtree_off.get(d).map(|v| v.as_slice()).unwrap_or(&[]);
1799 // AQ grain probe for IDRs: an IDR has no coding reference, but the AQ
1800 // grain veto needs a temporal signal, and the PREVIOUS SOURCE frame is
1801 // an even better probe than a reconstruction (no quantization in the
1802 // loop). The first frame of the stream has none — the veto fails open.
1803 let aq_probe = if is_idr && d > 0 {
1804 frames.get(d - 1).map(YuvFrame::as_planes)
1805 } else {
1806 None
1807 };
1808 let (au, recon) = code_picture(
1809 &cfg,
1810 &sps,
1811 &pps,
1812 &frames[d].as_planes(),
1813 is_idr,
1814 is_b,
1815 bref,
1816 poc,
1817 frame_num,
1818 &dpb,
1819 iqp,
1820 bqp,
1821 qpo,
1822 aq_probe.as_ref(),
1823 );
1824 aus.push(au);
1825 if !is_b {
1826 if let Some(r) = recon {
1827 dpb.insert(0, r);
1828 dpb.truncate(cfg.num_ref_frames as usize);
1829 }
1830 frame_num = (frame_num + 1) % 16;
1831 } else if let Some(r) = recon {
1832 // Reference B: into the DPB, frame_num advances (reference
1833 // pictures only) — mirroring the decoder's sliding window.
1834 dpb.insert(0, r);
1835 dpb.truncate(cfg.num_ref_frames as usize);
1836 frame_num = (frame_num + 1) % 16;
1837 }
1838 }
1839 aus
1840 }
1841}
1842
1843/// Codes ONE picture (IDR / P anchor / B) with explicit POC + frame_num + DPB.
1844/// Returns the access unit and, for reference pictures, the reconstruction to add
1845/// to the DPB (B-frames are non-reference → `None`). `dpb` is most-recent-first.
1846#[allow(clippy::too_many_arguments)]
1847fn code_picture(
1848 cfg: &EncoderConfig,
1849 sps: &Sps,
1850 pps: &Pps,
1851 frame: &YuvPlanes<'_>,
1852 is_idr: bool,
1853 is_b: bool,
1854 b_is_ref: bool,
1855 poc: i32,
1856 frame_num: u32,
1857 dpb: &[RefFrame],
1858 i_qp_offset: i32,
1859 b_qp_offset: i32,
1860 qpo: &[i32],
1861 aq_probe: Option<&YuvPlanes<'_>>,
1862) -> (Vec<u8>, Option<RefFrame>) {
1863 let mut out = Vec::new();
1864 let mut w = BitWriter::with_capacity(cfg.width * cfg.height / 2 + 4096);
1865 let poc_lsb = (poc as u32) & 0xFF; // log2_max_pic_order_cnt_lsb = 8
1866 // Per-GOP QP cascade, both offsets content-adaptive: B-frames are non-reference →
1867 // quantize HARDER (`b_qp_offset`, deeper on very predictable GOPs); the GOP's
1868 // I-frame is the root reference → quantize FINER (`i_qp_offset`, deeper on
1869 // predictable GOPs where the I dominates the bits).
1870 let qp = if is_b {
1871 (cfg.qp as i32 + b_qp_offset).clamp(0, 51) as u8
1872 } else if is_idr {
1873 (cfg.qp as i32 + i_qp_offset).clamp(0, 51) as u8
1874 } else {
1875 cfg.qp
1876 };
1877 let (nal_type, nal_ref_idc, recon) = if is_idr {
1878 sps.to_nal().write_annex_b(&mut out);
1879 pps.to_nal().write_annex_b(&mut out);
1880 slice::write_idr_slice_header(&mut w, cfg, qp);
1881 let mut r = if cfg.cabac {
1882 mb16::encode_slice_data_cabac_intra(&mut w, cfg, frame, qp, qpo, aq_probe)
1883 } else {
1884 mb16::encode_slice_data(&mut w, cfg, frame, qp, false, &[], qpo, aq_probe, &[])
1885 };
1886 r.poc = poc;
1887 r.frame_num = frame_num;
1888 (NalUnitType::IdrSlice, 3u8, Some(r))
1889 } else if is_b {
1890 // B is non-reference. We signal one active reference per list: L0[0] =
1891 // nearest PAST anchor (highest poc < current), L1[0] = nearest FUTURE anchor
1892 // (lowest poc > current) — the heads of the decoder's POC-ordered B lists.
1893 let l0 = dpb.iter().filter(|r| r.poc < poc).max_by_key(|r| r.poc);
1894 let l1 = dpb.iter().filter(|r| r.poc > poc).min_by_key(|r| r.poc);
1895 // b-pyramid: a reference B is CABAC-path v1 (the default config); the
1896 // CAVLC B coder stays leaf-only and `b_is_ref` is never set for it.
1897 let as_ref = b_is_ref && cfg.cabac && l0.is_some() && l1.is_some();
1898 slice::write_b_slice_header(&mut w, cfg, qp, frame_num, poc_lsb, 1, 1, as_ref);
1899 let mut b_recon: Option<RefFrame> = None;
1900 match (l0, l1) {
1901 // Leaf B's are non-reference — mb-tree offsets them at 0 anyway, so
1902 // `qpo` is `&[]` here (the anchor reference chain carries the temporal AQ).
1903 (Some(l0), Some(l1)) if cfg.cabac => {
1904 b_recon = mb16::encode_slice_data_cabac_b(
1905 &mut w,
1906 cfg,
1907 frame,
1908 qp,
1909 poc,
1910 l0,
1911 l1,
1912 &[],
1913 as_ref,
1914 );
1915 }
1916 (Some(l0), Some(l1)) => {
1917 mb16::encode_slice_data_b(&mut w, cfg, frame, qp, poc, l0, l1, &[]);
1918 }
1919 // A B with no bracketing anchor pair can't be List-0/1 coded; fall back
1920 // to an all-B_Skip slice (spatial-direct) so the stream stays legal.
1921 _ => {
1922 let n = cfg.mb_width() * cfg.mb_height();
1923 if cfg.cabac {
1924 mb16::encode_all_skip_b_cabac(&mut w, cfg, qp, n);
1925 } else {
1926 w.write_ue(n as u32);
1927 w.rbsp_trailing_bits();
1928 }
1929 }
1930 }
1931 if as_ref {
1932 if let Some(r) = &mut b_recon {
1933 r.poc = poc;
1934 r.frame_num = frame_num;
1935 }
1936 (NalUnitType::NonIdrSlice, 2u8, b_recon)
1937 } else {
1938 (NalUnitType::NonIdrSlice, 0u8, None)
1939 }
1940 } else {
1941 // P anchor: L0 = the DPB (past anchors), ordered most-recent-first. Both CAVLC
1942 // and CABAC now code ref_idx_l0 (cb_ref_idx / parse_ref_idx_cabac), so a P slice
1943 // searches + signals the full DPB (`--refs N`) under either entropy coder.
1944 let p_dpb: &[RefFrame] = dpb;
1945 let wp: Vec<(i32, i32)> = if cfg.weightp {
1946 estimate_luma_weights(cfg, frame, p_dpb)
1947 } else {
1948 Vec::new()
1949 };
1950 let wp_hdr = if cfg.weightp {
1951 Some(wp.as_slice())
1952 } else {
1953 None
1954 };
1955 slice::write_p_slice_header(&mut w, cfg, qp, frame_num, poc_lsb, p_dpb.len(), wp_hdr);
1956 let mut r = if cfg.cabac {
1957 mb16::encode_slice_data_cabac_p(&mut w, cfg, frame, qp, p_dpb, qpo, &wp)
1958 } else {
1959 mb16::encode_slice_data(&mut w, cfg, frame, qp, true, dpb, qpo, None, &wp)
1960 };
1961 r.poc = poc;
1962 r.frame_num = frame_num;
1963 (NalUnitType::NonIdrSlice, 3u8, Some(r))
1964 };
1965 let slice_bytes = w.into_bytes();
1966 NalUnit::new(nal_ref_idc, nal_type, slice_bytes).write_annex_b(&mut out);
1967 (out, recon)
1968}
1969
1970/// The B-favorability threshold on the per-GOP signal (`gop_bi_residual`): below it,
1971/// motion is predictable enough that B-frames pay AND the I-frame dominates the GOP's
1972/// bits (so it wants a deeper QP cascade); above it the GOP is busy.
1973const BI_THRESH: f64 = 4.0;
1974
1975/// Whether a GOP's temporal residual makes B-frames pay (predictable motion).
1976/// B-favorability threshold on the bi-prediction residual, evaluated PER
1977/// ANCHOR GAP (bframes-v2). Refit 2026-08-26 from the 12-clip BD truth table:
1978/// the original 4.0 (shared with the QP-cascade ramps, which keep it) captured
1979/// only the near-static winners and left five pan/texture/noise clips worth
1980/// -11..-17% each on the table (mobile 7.98, shields 8.07, grain 7.61, city
1981/// 5.87, tempete 4.69 all WIN at fixed B); the documented fastmotion losers
1982/// sit above (football 15.1, park_joy 9.1, crowd_run 8.3). 8.2 splits them.
1983/// The per-GAP unit (not per clip) is what handles crew: its flash gaps read
1984/// unfavorable and code P while its calm gaps take the B win — the clip-level
1985/// scalars could not separate crew from tempete (0.006 apart in dcfrac, a
1986/// margin that is fitting noise, refused). Holdout-gated on six unseen clips.
1987const B_GAP_THRESH: f64 = 8.2;
1988
1989fn bframes_favorable(residual: f64) -> bool {
1990 residual < B_GAP_THRESH
1991}
1992
1993/// Content-adaptive per-GOP I-frame QP offset (the ip_ratio cascade, DISPATCHED by
1994/// content). `base` is the busy-GOP offset (`cfg.i_qp_offset`, default −3); a
1995/// predictable GOP — where the I-frame is a large fraction of the GOP's bits, so
1996/// investing in it pays outsized — gets up to 2 QP steps FINER, ramping from `base`
1997/// at the threshold to `base−2` at residual 0. Calibrated: busy ≈ −3, compressible
1998/// ≈ −5 (−11.6% vs −7.3% at −3). `base == 0` (the opt-out) disables it entirely so
1999/// the byte-identical escape hatch survives.
2000fn gop_iqp_offset(residual: f64, base: i32) -> i32 {
2001 if base == 0 {
2002 return 0;
2003 }
2004 let bonus =
2005 rusty_h264_common::fmath::round(2.0 * ((BI_THRESH - residual) / BI_THRESH).clamp(0.0, 1.0))
2006 as i32;
2007 base - bonus
2008}
2009
2010/// Content-adaptive per-GOP B-frame QP offset. B-frames are non-reference, so on a
2011/// VERY predictable GOP (bi-pred + spatial-direct nail them → tiny residual) they can
2012/// be quantized much HARDER for near-free bits. But the optimum is KNIFE-EDGE in the
2013/// signal — measured ~+8 at residual 0.10 yet ~+2 by residual 0.29 (and a heavy LOSS
2014/// at +12 there) — so unlike the I-cascade this ramp is STEEP and confined to the
2015/// near-perfect-motion regime: `base` (default +2) everywhere, boosted up to +4 only
2016/// as residual → 0 (decaying to `base` by ~0.3/px). Deliberately conservative — it
2017/// helps near-static / clean-pan content and must never touch the common range.
2018fn gop_bframe_qp_offset(residual: f64, base: i32) -> i32 {
2019 const RAMP: f64 = 0.3; // residual above this gets no boost (steep — see calibration)
2020 let boost =
2021 rusty_h264_common::fmath::round(4.0 * ((RAMP - residual) / RAMP).clamp(0.0, 1.0)) as i32;
2022 base + boost
2023}
2024
2025/// Adaptive B-COUNT (B-frames per anchor gap) for `auto` mode. The RATIO of the
2026/// 2-gap to 1-gap bi-prediction residual measures how fast bi-pred degrades as the
2027/// anchor spacing widens: LOW ratio (content survives wider gaps) carries MORE cheap
2028/// non-reference B's; HIGH ratio (simple translation — degrades fast, so wider anchors
2029/// cost more than the extra B's save) wants a single equidistant B. Calibrated on
2030/// pans/zoom: ratio ≥ 1.8 → 1, ≥ 1.4 → 2, else 3. Capped at `max_b` (the `auto` cap).
2031/// TEMPORAL PREDICTABILITY probe (Great Gate P3 item 4). Returns the
2032/// `2-gap / 1-gap` motion-compensated residual ratio for a frame window --
2033/// the axis the mb-tree dispatch has been waiting on.
2034///
2035/// Why THIS signal for mb-tree specifically: mb-tree lowers QP on blocks whose
2036/// quality PROPAGATES to the frames that reference them. That model assumes the
2037/// referenced pixels are still there N frames later. On a pan they translate out
2038/// of frame, so propagation decays and the lookahead over-credits the block. The
2039/// ratio measures exactly that decay -- how much worse prediction gets when the
2040/// reference gap doubles -- where `lv_spread`/`flat_run` (the REFUSED candidate)
2041/// are spatial statistics that merely correlate with panning on this corpus.
2042/// `f64::INFINITY` when the window is too short to measure.
2043pub fn temporal_decay_ratio(frames: &[YuvFrame], w: usize, h: usize) -> f64 {
2044 let g1 = gop_bi_residual(frames, w, h, 1);
2045 let g2 = gop_bi_residual(frames, w, h, 2);
2046 if !g1.is_finite() || !g2.is_finite() {
2047 return f64::INFINITY;
2048 }
2049 g2 / g1.max(1e-3)
2050}
2051
2052fn adaptive_bcount(frames: &[YuvFrame], w: usize, h: usize, max_b: usize) -> usize {
2053 let cap = max_b.clamp(1, 3);
2054 let g1 = gop_bi_residual(frames, w, h, 1);
2055 let g2 = gop_bi_residual(frames, w, h, 2);
2056 if !g1.is_finite() || !g2.is_finite() {
2057 return 1;
2058 }
2059 let ratio = g2 / g1.max(1e-3);
2060 // Calibrated on this encoder's (subsampled global-ME) ratios: a simple
2061 // translation degrades to ~1.5 (→ 1 B), predictable-under-wide-gaps content sits
2062 // ~1.3 or below (→ 3 B).
2063 let c = if ratio >= 1.4 {
2064 1
2065 } else if ratio >= 1.3 {
2066 2
2067 } else {
2068 3
2069 };
2070 c.clamp(1, cap)
2071}
2072
2073/// Cheap content signal for the content-adaptive dispatch: the mean per-pixel
2074/// residual of a coarse GLOBAL-motion BI-prediction, over a subsample of interior
2075/// frames. Low = temporally predictable (bi-pred + spatial-direct cheap → B-frames
2076/// WIN, and the I-frame dominates → deeper QP cascade); high = busy motion.
2077/// `f64::INFINITY` when the GOP is too short to measure (treated as busy).
2078///
2079/// Global (not block) ME keeps it O(pixels)-cheap and biases toward "coherent
2080/// motion", which is what spatial-direct/skip exploit. Thresholds calibrated on
2081/// extremes (pan ~0.03/px, high-motion ~12.3/px); refine on a corpus.
2082fn gop_bi_residual(frames: &[YuvFrame], w: usize, h: usize, gap: usize) -> f64 {
2083 let n = frames.len();
2084 if n < 2 * gap + 1 || w < 48 || h < 48 {
2085 return f64::INFINITY;
2086 }
2087 // Subsampled SAD of `cur` vs `rf` shifted by (dx,dy): interior pixels only
2088 // (|shift| ≤ 15 stays in-bounds, no clamping), every 4th pixel for speed.
2089 let sad = |cur: &[u8], rf: &[u8], dx: isize, dy: isize| -> u64 {
2090 let mut s = 0u64;
2091 let mut y = 16;
2092 while y < h - 16 {
2093 let cbase = (y * w) as isize;
2094 let rbase = ((y as isize + dy) * w as isize) + dx;
2095 let mut x = 16isize;
2096 while x < (w - 16) as isize {
2097 let c = cur[(cbase + x) as usize] as i32;
2098 let r = rf[(rbase + x) as usize] as i32;
2099 s += (c - r).unsigned_abs() as u64;
2100 x += 8;
2101 }
2102 y += 8;
2103 }
2104 s
2105 };
2106 // Coarse global ME: ±12 step 4, then refine ±3 step 1.
2107 let global_me = |cur: &[u8], rf: &[u8]| -> (isize, isize) {
2108 let (mut best, mut bc) = ((0isize, 0isize), u64::MAX);
2109 let mut dy = -12;
2110 while dy <= 12 {
2111 let mut dx = -12;
2112 while dx <= 12 {
2113 let c = sad(cur, rf, dx, dy);
2114 if c < bc {
2115 bc = c;
2116 best = (dx, dy);
2117 }
2118 dx += 4;
2119 }
2120 dy += 4;
2121 }
2122 // The refine window re-evaluates its own CENTRE — the coarse best,
2123 // whose cost `bc` already carries. `c < bc` is strict and `bc` only
2124 // decreases, so that re-visit can never win: skipping it is
2125 // decision-identical and saves one full subsampled-SAD call per
2126 // `global_me`. The range expressions stay on `best` (NOT hoisted):
2127 // the inner range re-reads the current `best.0` per row exactly as it
2128 // always did, so the search visits the same points minus the centre.
2129 let centre = best;
2130 for dy in best.1 - 3..=best.1 + 3 {
2131 for dx in best.0 - 3..=best.0 + 3 {
2132 if (dx, dy) == centre {
2133 continue;
2134 }
2135 let c = sad(cur, rf, dx, dy);
2136 if c < bc {
2137 bc = c;
2138 best = (dx, dy);
2139 }
2140 }
2141 }
2142 best
2143 };
2144 let mut n_samp = 0usize;
2145 {
2146 let mut y = 16;
2147 while y < h - 16 {
2148 let mut x = 16;
2149 while x < w - 16 {
2150 n_samp += 1;
2151 x += 8;
2152 }
2153 y += 8;
2154 }
2155 }
2156 let step = (n / 5).max(1);
2157 let (mut total, mut cnt) = (0f64, 0usize);
2158 // `gap` frames each side (1 = adjacent, for the B/P dispatch; 2 probes how well
2159 // bi-prediction survives WIDER anchor spacing, for the adaptive B-count).
2160 let mut d = gap;
2161 while d < n - gap {
2162 let (cur, past, fut) = (&frames[d].y, &frames[d - gap].y, &frames[d + gap].y);
2163 let (mpx, mpy) = global_me(cur, past);
2164 let (mfx, mfy) = global_me(cur, fut);
2165 let mut bi = 0u64;
2166 let mut y = 16;
2167 while y < h - 16 {
2168 let mut x = 16isize;
2169 while x < (w - 16) as isize {
2170 let c = cur[y * w + x as usize] as i32;
2171 let p = past[((y as isize + mpy) * w as isize + x + mpx) as usize] as i32;
2172 let f = fut[((y as isize + mfy) * w as isize + x + mfx) as usize] as i32;
2173 bi += (c - ((p + f + 1) >> 1)).unsigned_abs() as u64;
2174 x += 8;
2175 }
2176 y += 8;
2177 }
2178 total += bi as f64 / n_samp as f64;
2179 cnt += 1;
2180 d += step;
2181 }
2182 if cnt > 0 {
2183 total / cnt as f64
2184 } else {
2185 f64::INFINITY
2186 }
2187}
2188
2189#[cfg(test)]
2190mod tests {
2191 use super::*;
2192 #[allow(unused_imports)]
2193 use alloc::{
2194 boxed::Box,
2195 format,
2196 string::{String, ToString},
2197 vec,
2198 vec::Vec,
2199 };
2200 #[allow(unused_imports)]
2201 use rusty_h264_common::once::OnceLock;
2202
2203 fn textured(w: usize, h: usize, t: usize) -> YuvFrame {
2204 let mut f = YuvFrame::black(w, h);
2205 for y in 0..h {
2206 for x in 0..w {
2207 let v = ((x + 3 * t) * 255 / w) as u32 + ((y * 7 + x * 3 + t) % 23) as u32;
2208 f.y[y * w + x] = v.min(255) as u8;
2209 }
2210 }
2211 f
2212 }
2213
2214 /// The half-pel plane cache is built lazily and only when the search is
2215 /// not `Fast` — a fact, not a comment, because the chip configuration's
2216 /// memory model has no half-pel term. `baseline()` codes P-frames with
2217 /// every reference's cache empty; the same configuration on
2218 /// `Preset::Balanced` fills it.
2219 #[test]
2220 fn fast_preset_never_builds_the_half_pel_cache() {
2221 let (w, h) = (64, 48);
2222 let frames: Vec<YuvFrame> = (0..4).map(|t| textured(w, h, t)).collect();
2223 for refs in [1, 3] {
2224 let mut cfg = EncoderConfig::baseline(w, h);
2225 cfg.gop_size = 8;
2226 cfg.min_keyint = 8;
2227 cfg.num_ref_frames = refs;
2228 assert_eq!(cfg.preset, Preset::Fast);
2229 let mut enc = Encoder::new(cfg).unwrap();
2230 for f in &frames {
2231 let _ = enc.encode(f);
2232 }
2233 assert!(!enc.refs.is_empty());
2234 assert!(
2235 enc.refs.iter().all(|r| r.hpel.get().is_none()),
2236 "Fast built a half-pel cache (refs={refs})"
2237 );
2238 }
2239 let mut cfg = EncoderConfig::baseline(w, h);
2240 cfg.gop_size = 8;
2241 cfg.min_keyint = 8;
2242 cfg.num_ref_frames = 3;
2243 cfg.preset = Preset::Balanced;
2244 let mut enc = Encoder::new(cfg).unwrap();
2245 for f in &frames {
2246 let _ = enc.encode(f);
2247 }
2248 assert!(
2249 enc.refs.iter().any(|r| r.hpel.get().is_some()),
2250 "Balanced never built the half-pel cache"
2251 );
2252 }
2253
2254 #[test]
2255 fn rejects_unsupported_config() {
2256 // R6 made High + 8x8 + CABAC a SUPPORTED combination (transform_size_8x8_flag
2257 // at both syntax positions + the ctxBlockCat-5 residual writer), verified
2258 // pixel-identical against ffmpeg. This test used to assert it was rejected;
2259 // asserting it is ACCEPTED is what keeps the capability from silently
2260 // regressing behind a re-added guard.
2261 let mut cfg = EncoderConfig::new(16, 16);
2262 cfg.profile = Profile::High;
2263 cfg.transform_8x8 = true;
2264 cfg.cabac = true;
2265 assert!(
2266 Encoder::new(cfg).is_ok(),
2267 "High + 8x8 + CABAC must be accepted"
2268 );
2269 // Narrowing the profile CLAMPS the 8x8 transform rather than failing: it is
2270 // default-on, and `EncoderConfig::new()` + `profile = Main` must stay a valid
2271 // pair. Assert the encoder builds AND that the PPS does not advertise a tool
2272 // Main cannot carry.
2273 let mut cfg = EncoderConfig::new(16, 16);
2274 cfg.profile = Profile::Main;
2275 cfg.transform_8x8 = true;
2276 let enc = Encoder::new(cfg).expect("Main + 8x8 must clamp, not fail");
2277 assert!(
2278 !enc.cfg.transform_8x8,
2279 "8x8 must be cleared when the profile cannot signal it"
2280 );
2281 }
2282
2283 #[test]
2284 fn encodes_access_unit_with_sps_pps_idr() {
2285 use rusty_h264_common::nal::split_annex_b;
2286 let cfg = EncoderConfig::new(32, 32);
2287 let mut enc = Encoder::new(cfg).unwrap();
2288 let frame = YuvFrame::black(32, 32);
2289 // mb-tree defaults ON: the streaming path buffers one GOP, so a single
2290 // frame's access unit arrives on `flush` (same shape the fuzz seeds use).
2291 let mut au = enc.encode(&frame);
2292 au.extend(enc.flush());
2293
2294 let nals = split_annex_b(&au);
2295 assert_eq!(nals.len(), 3);
2296 assert_eq!(NalUnitType::from_id(nals[0][0]), NalUnitType::Sps);
2297 assert_eq!(NalUnitType::from_id(nals[1][0]), NalUnitType::Pps);
2298 assert_eq!(NalUnitType::from_id(nals[2][0]), NalUnitType::IdrSlice);
2299 }
2300
2301 #[test]
2302 fn encode_all_matches_sequential_cqp() {
2303 // GOP-parallel batch encoding must be byte-identical to frame-by-frame
2304 // sequential encoding at constant QP (GOPs are independent).
2305 let (w, h) = (48usize, 32usize);
2306 let mut cfg = EncoderConfig::new(w, h);
2307 cfg.gop_size = 4; // 10 frames → 3 GOPs (4,4,2)
2308 let frames: Vec<YuvFrame> = (0..10u8)
2309 .map(|t| YuvFrame {
2310 width: w,
2311 height: h,
2312 y: (0..w * h)
2313 .map(|i| (i as u8).wrapping_add(t.wrapping_mul(7)))
2314 .collect(),
2315 u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
2316 v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
2317 })
2318 .collect();
2319 let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
2320 let mut seq: Vec<u8> = frames.iter().flat_map(|f| seq_enc.encode(f)).collect();
2321 seq.extend_from_slice(&seq_enc.flush()); // end of stream (lookahead tail)
2322 let par: Vec<u8> = Encoder::new(cfg)
2323 .unwrap()
2324 .encode_all(&frames)
2325 .unwrap()
2326 .concat();
2327 assert_eq!(seq, par, "GOP-parallel must equal sequential+flush at CQP");
2328 }
2329
2330 #[test]
2331 fn encode_all_matches_sequential_quality_preset() {
2332 // Same invariant on the QUALITY preset, whose per-frame dispatch decisions
2333 // (b2_mgain SAD/mv-cost routing) once lived in a process-global and RACED
2334 // across GOP workers — divergence only appears with >1 GOP in flight, which
2335 // the single-GOP hash harness never exercised. Content varies per frame so
2336 // the per-frame routing decisions actually differ between GOPs.
2337 let (w, h) = (48usize, 32usize);
2338 let mut cfg = EncoderConfig::new(w, h);
2339 cfg.gop_size = 3; // 12 frames → 4 GOPs, several workers in flight
2340 cfg.preset = crate::config::Preset::Quality;
2341 let frames: Vec<YuvFrame> = (0..12u8)
2342 .map(|t| YuvFrame {
2343 width: w,
2344 height: h,
2345 y: (0..w * h)
2346 .map(|i| {
2347 // alternate calm and busy frames so the mgain probe flips
2348 let base = (i as u8).wrapping_add(t.wrapping_mul(3));
2349 if t % 2 == 0 {
2350 base
2351 } else {
2352 base.wrapping_mul(37).wrapping_add(i as u8)
2353 }
2354 })
2355 .collect(),
2356 u: vec![128u8.wrapping_add(t); (w / 2) * (h / 2)],
2357 v: vec![128u8.wrapping_sub(t); (w / 2) * (h / 2)],
2358 })
2359 .collect();
2360 let mut seq_enc = Encoder::new(cfg.clone()).unwrap();
2361 let mut seq: Vec<u8> = frames.iter().flat_map(|f| seq_enc.encode(f)).collect();
2362 seq.extend_from_slice(&seq_enc.flush());
2363 let par: Vec<u8> = Encoder::new(cfg)
2364 .unwrap()
2365 .encode_all(&frames)
2366 .unwrap()
2367 .concat();
2368 assert_eq!(
2369 seq, par,
2370 "quality-preset GOP-parallel must equal sequential+flush"
2371 );
2372 }
2373
2374 #[test]
2375 fn rejects_mismatched_frame() {
2376 let cfg = EncoderConfig::new(16, 16);
2377 let mut enc = Encoder::new(cfg).unwrap();
2378 let frame = YuvFrame::black(32, 16);
2379 assert_eq!(enc.try_encode(&frame), Err(EncodeError::FrameMismatch));
2380 }
2381}