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