Skip to main content

rusty_h264_encoder/
mbtree.rs

1//! Macroblock-tree lookahead adaptive quantization (temporal AQ).
2//!
3//! A cheap forward pass over a GOP's SOURCE frames estimates, per macroblock, how
4//! much of the *future's* coding cost depends on it, then lowers the QP of
5//! heavily-referenced macroblocks — investing bits where they pay off across many
6//! later frames (a sharp reference makes every frame that predicts from it cheaper).
7//! This is the temporal complement to the spatial AQ (`aq_qp_map`): AQ moves bits
8//! by texture *within* a frame; mb-tree moves them by reference *importance across*
9//! frames.
10//!
11//! Method (x264's mb-tree, adapted to our CQP GOP):
12//!   1. Per frame, per MB: `intra` = spatial AC SATD; `inter` = best small-search
13//!      motion-compensated residual SATD to the previous SOURCE frame (capped at
14//!      `intra`), plus the winning MV. Source-domain (like x264's lowres lookahead)
15//!      so no reconstruction is needed — it's a pure pre-pass.
16//!   2. Backward propagation: walk frames last→first. Each MB's total importance is
17//!      `intra + propagate_in`; the fraction its predictor earned — `(intra-inter)/
18//!      intra` — is credited to the reference MBs it points to (bilinear by MV,
19//!      area-weighted over the up-to-4 overlapped MBs) in the previous frame.
20//!   3. QP offset `= -strength · log2((intra + propagate_in) / intra)` (≤ 0:
21//!      heavily-referenced MBs get finer QP; leaves get 0). CENTERED per GOP
22//!      (subtract the GOP-mean offset) so the average QP — hence the rate — is
23//!      preserved and the effect is a pure redistribution of bits toward the MBs
24//!      the future depends on.
25
26use crate::config::{EncoderConfig, LookaheadMode};
27use rusty_h264_common::inter::mc_luma;
28use rusty_h264_common::transform::hadamard_4x4;
29use rusty_h264_common::YuvFrame;
30
31/// Per-MB lookahead cost + motion for one frame.
32#[derive(Clone, Copy)]
33struct MbCost {
34    intra: i32,     // spatial AC SATD, >= 1
35    inter: i32,     // best MC-residual SATD to the previous frame, capped at `intra`
36    mv: (i32, i32), // winning MV (quarter-pel) — propagation direction
37}
38
39/// SATD of a 4×4 residual (sum of |Hadamard coeffs|).
40fn satd4(res: &[i32; 16]) -> i64 {
41    hadamard_4x4(res).iter().map(|&v| v.unsigned_abs() as i64).sum()
42}
43
44/// Edge-clamped coded-size luma (matches the encoder's source preparation).
45pub(crate) fn coded_luma(cfg: &EncoderConfig, frame: &YuvFrame) -> Vec<u8> {
46    let (cw, ch) = (cfg.mb_width() * 16, cfg.mb_height() * 16);
47    let (w, h) = (frame.width, frame.height);
48    let mut y = vec![0u8; cw * ch];
49    // Row-slice copy + right-edge fill instead of a per-pixel double-`min`
50    // gather — the `load_mb` shape (write into a PRE-SIZED destination), which
51    // is the panic-round shape that works, as distinct from the append shape
52    // that regressed. Interior rows become a `memcpy` + `memset`; the bottom
53    // padding re-copies row `h-1`. Bit-identical to the per-pixel form by the
54    // `coded_luma_matches_per_pixel` oracle.
55    let wc = w.min(cw);
56    for j in 0..ch {
57        let src = &frame.y[j.min(h - 1) * w..][..w];
58        let dst = &mut y[j * cw..][..cw];
59        dst[..wc].copy_from_slice(&src[..wc]);
60        dst[wc..].fill(src[w - 1]);
61    }
62    y
63}
64
65/// 2×2-average downsample of a luma plane to half resolution (both dims are MB
66/// multiples → stay even). The Hybrid/HalfRes lookahead runs the MV search on this:
67/// 4× fewer pixels, ~4× cheaper. The MV direction survives; only the COST accuracy
68/// suffers on blurred detail — which the Hybrid mode fixes by re-scoring at full-res.
69fn downsample2x(y: &[u8], cw: usize, ch: usize) -> (Vec<u8>, usize, usize) {
70    let (hw, hh) = (cw / 2, ch / 2);
71    let mut out = vec![0u8; hw * hh];
72    // Row-pair slices + `chunks_exact(2)` instead of five multiplied indexings
73    // per output pixel: the iterator shape carries the extents, so the four
74    // source loads and the store lose their per-pixel bounds checks and the
75    // row-base multiplies hoist to once per row. Same sums, same `(s+2)/4`
76    // rounding: BIT-IDENTICAL.
77    for j in 0..hh {
78        let r0 = &y[2 * j * cw..][..cw];
79        let r1 = &y[(2 * j + 1) * cw..][..cw];
80        let dst = &mut out[j * hw..][..hw];
81        for ((o, p0), p1) in dst.iter_mut().zip(r0.chunks_exact(2)).zip(r1.chunks_exact(2)) {
82            let s = p0[0] as u32 + p0[1] as u32 + p1[0] as u32 + p1[1] as u32;
83            *o = ((s + 2) / 4) as u8;
84        }
85    }
86    (out, hw, hh)
87}
88
89/// Spatial AC SATD of a `bs`×`bs` block at pixel `(bx0, by0)` (DC excluded, summed
90/// over 4×4 sub-blocks). The intra "cost" floor — how expensive with no prediction.
91fn intra_cost(sy: &[u8], cw: usize, bx0: usize, by0: usize, bs: usize) -> i32 {
92    let mut s = 0i64;
93    for by in 0..bs / 4 {
94        for bx in 0..bs / 4 {
95            let mut blk = [0i32; 16];
96            for dy in 0..4 {
97                for dx in 0..4 {
98                    blk[dy * 4 + dx] = sy[(by0 + by * 4 + dy) * cw + bx0 + bx * 4 + dx] as i32;
99                }
100            }
101            let h = hadamard_4x4(&blk);
102            s += h[1..].iter().map(|&v| v.unsigned_abs() as i64).sum::<i64>();
103        }
104    }
105    (s.min(i32::MAX as i64) as i32).max(1)
106}
107
108/// Full-pel MC-residual SATD of a `bs`×`bs` block at a given (plane) quarter-pel MV.
109/// DETERMINISTIC cost instrument for the lookahead (H-36). Wall-clock on this box
110/// swings ±40 points run-to-run on an IDENTICAL config, which is far larger than
111/// the content effect being measured — so the lookahead's cost is judged by its
112/// WORK COUNT (candidate evaluations), which is exactly reproducible.
113pub(crate) static SATD_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
114
115fn mc_satd(sy: &[u8], cw: usize, ch: usize, ref_y: &[u8], bx0: usize, by0: usize, bs: usize, mv: (i32, i32)) -> i64 {
116    SATD_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
117    // H-35: the diamond below only ever probes FULL-PEL vectors (`dx * 4` in
118    // quarter-pel units), so nearly every call can read the reference IN PLACE and
119    // hand both planes to the vendored asm SATD — instead of copying a bs×bs block
120    // out of `mc_luma` and then running a SCALAR per-4×4 Hadamard. ("Exported ≠
121    // wired": the main ME has used the asm kernel for months; the lookahead had its
122    // own scalar twin, which is why mb-tree cost far more than its half-res search
123    // should.) Identical value: `satd_px`'s scalar arm IS this function's old sum,
124    // and its asm arm is pinned byte-exact to that by the accel oracles.
125    let (ix, iy) = (bx0 as isize + (mv.0 >> 2) as isize, by0 as isize + (mv.1 >> 2) as isize);
126    if mv.0 & 3 == 0
127        && mv.1 & 3 == 0
128        && ix >= 0
129        && iy >= 0
130        && ix as usize + bs <= cw
131        && iy as usize + bs <= ch
132        && by0 + bs <= ch
133        && bx0 + bs <= cw
134    {
135        return crate::mb16::satd_px(
136            &sy[by0 * cw + bx0..],
137            cw,
138            &ref_y[iy as usize * cw + ix as usize..],
139            cw,
140            bs,
141            bs,
142        );
143    }
144    // Sub-pel seed probe or an edge-overhanging vector: the general path.
145    let mut pred = [0u8; 256]; // bs ≤ 16 → fits; stride = bs
146    mc_luma(ref_y, cw, ch, bx0, by0, bs, bs, mv.0, mv.1, &mut pred);
147    let mut s = 0i64;
148    for by in 0..bs / 4 {
149        for bx in 0..bs / 4 {
150            let mut res = [0i32; 16];
151            for dy in 0..4 {
152                for dx in 0..4 {
153                    res[dy * 4 + dx] = sy[(by0 + by * 4 + dy) * cw + bx0 + bx * 4 + dx] as i32
154                        - pred[(by * 4 + dy) * bs + (bx * 4 + dx)] as i32;
155                }
156            }
157            s += satd4(&res);
158        }
159    }
160    s
161}
162
163/// Best MC-residual SATD of a `bs`×`bs` block and its winning (plane) MV, via a
164/// full-pel diamond search SEEDED from a predictor (the neighbour's MV, for pan
165/// coherence). The diamond (step 8→1 full-pel) tracks large motion a fixed ±2px set
166/// missed — a wrong MV gives mb-tree a wrong propagation DIRECTION (misdirects bits).
167fn inter_cost(sy: &[u8], cw: usize, ch: usize, ref_y: &[u8], bx0: usize, by0: usize, bs: usize, seed: (i32, i32), max_step: i32) -> (i32, (i32, i32)) {
168    let mut best_mv = (0, 0);
169    let mut best = mc_satd(sy, cw, ch, ref_y, bx0, by0, bs, (0, 0));
170    // H-45: a PROVABLY byte-identical early-out. SATD is a sum of absolute values,
171    // so every candidate is ≥ 0; once `best == 0` the guard `s < best` can never
172    // fire again, and both the winning MV and the cost are already final. Without
173    // this the diamond still pays its fixed floor — 4 probes at each of the 4 step
174    // levels (8→4→2→1), because the round that TERMINATES a level costs 4 probes
175    // that do not move — which is why the lookahead's eval count came out nearly
176    // content-invariant (16–18 /MB/frame on flat AND busy clips) and why mb-tree's
177    // relative cost was WORST on the static content it helps most (+18.8% akiyo).
178    if best == 0 {
179        return (0, best_mv);
180    }
181    if seed != (0, 0) {
182        let s = mc_satd(sy, cw, ch, ref_y, bx0, by0, bs, seed);
183        if s < best {
184            best = s;
185            best_mv = seed;
186        }
187    }
188    // `max_step` bounds the initial diamond hop: 8 for a from-scratch search, small
189    // (2) for the hybrid's full-res refine around an already-good coarse MV.
190    let mut step = max_step;
191    while step >= 1 {
192        loop {
193            let mut moved = false;
194            for &(dx, dy) in &[(step, 0), (-step, 0), (0, step), (0, -step)] {
195                let mv = (best_mv.0 + dx * 4, best_mv.1 + dy * 4); // quarter-pel units
196                let s = mc_satd(sy, cw, ch, ref_y, bx0, by0, bs, mv);
197                if s < best {
198                    best = s;
199                    best_mv = mv;
200                    moved = true;
201                }
202            }
203            if !moved {
204                break;
205            }
206        }
207        step >>= 1;
208    }
209    (best.min(i32::MAX as i64) as i32, best_mv)
210}
211
212/// Per-MB costs for one frame, at the lookahead `mode`'s resolution(s). MVs are
213/// always returned in FULL-res quarter-pel (propagation is resolution-independent).
214/// - `FullRes`: search + score at full-res (16×16).
215/// - `HalfRes`: search + score at half-res (8×8) — MV scaled ×2.
216/// - `Hybrid`: search the MV on half-res, then REFINE + score intra/inter at
217///   full-res (the cost accuracy the pure-half-res path lost, at ~its speed).
218///
219/// `ref_full`/`ref_half` are `None` for the IDR (intra-only, nothing to propagate).
220#[allow(clippy::too_many_arguments)]
221fn frame_costs(
222    full: &[u8],
223    cwf: usize,
224    chf: usize,
225    half: &[u8],
226    cwh: usize,
227    chh: usize,
228    mb_w: usize,
229    mb_h: usize,
230    ref_full: Option<&[u8]>,
231    ref_half: Option<&[u8]>,
232    mode: LookaheadMode,
233) -> Vec<MbCost> {
234    let mut out: Vec<MbCost> = Vec::with_capacity(mb_w * mb_h);
235    for mb_y in 0..mb_h {
236        for mb_x in 0..mb_w {
237            // Neighbour MV seed (full-res quarter-pel), for pan coherence.
238            let seed_full = if mb_x > 0 {
239                out[mb_y * mb_w + mb_x - 1].mv
240            } else if mb_y > 0 {
241                out[(mb_y - 1) * mb_w + mb_x].mv
242            } else {
243                (0, 0)
244            };
245            // Intra cost at the scoring resolution (full for FullRes/Hybrid, half for HalfRes).
246            let intra = if mode == LookaheadMode::HalfRes {
247                intra_cost(half, cwh, mb_x * 8, mb_y * 8, 8)
248            } else {
249                intra_cost(full, cwf, mb_x * 16, mb_y * 16, 16)
250            };
251            let (inter, mv) = match (mode, ref_full, ref_half) {
252                (LookaheadMode::FullRes, Some(rf), _) => {
253                    let (ic, mv) = inter_cost(full, cwf, chf, rf, mb_x * 16, mb_y * 16, 16, seed_full, 8);
254                    (ic.min(intra), mv)
255                }
256                (LookaheadMode::HalfRes, _, Some(rh)) => {
257                    let seed = (seed_full.0 / 2, seed_full.1 / 2);
258                    let (ic, mvp) = inter_cost(half, cwh, chh, rh, mb_x * 8, mb_y * 8, 8, seed, 8);
259                    (ic.min(intra), (mvp.0 * 2, mvp.1 * 2))
260                }
261                (LookaheadMode::Hybrid, Some(rf), Some(rh)) => {
262                    // Coarse MV from the cheap half-res search…
263                    let seed = (seed_full.0 / 2, seed_full.1 / 2);
264                    let (_, mvp) = inter_cost(half, cwh, chh, rh, mb_x * 8, mb_y * 8, 8, seed, 8);
265                    let coarse = (mvp.0 * 2, mvp.1 * 2); // → full-res quarter-pel
266                    // …then a SMALL full-res refine that also gives the accurate cost.
267                    let (ic, mv) = inter_cost(full, cwf, chf, rf, mb_x * 16, mb_y * 16, 16, coarse, 2);
268                    (ic.min(intra), mv)
269                }
270                _ => (intra, (0, 0)), // IDR (no reference)
271            };
272            out.push(MbCost { intra, inter, mv });
273        }
274    }
275    out
276}
277
278/// Distribute `amount` from frame `f`'s MB (referencing the previous frame at MV
279/// `mv`) into `prev`'s per-MB propagation accumulator, area-weighted over the up-to-4
280/// macroblocks the referenced 16×16 block overlaps (edge-clamped).
281fn propagate_to(prev: &mut [f64], mb_w: usize, mb_h: usize, mb_x: usize, mb_y: usize, mv: (i32, i32), amount: f64) {
282    if amount <= 0.0 {
283        return;
284    }
285    // Referenced block top-left in pixels (integer part of the quarter-pel MV),
286    // clamped so it stays inside the frame.
287    let rx = (mb_x as i32 * 16 + (mv.0 >> 2)).clamp(0, (mb_w as i32 - 1) * 16);
288    let ry = (mb_y as i32 * 16 + (mv.1 >> 2)).clamp(0, (mb_h as i32 - 1) * 16);
289    let cx0 = (rx / 16) as usize;
290    let cy0 = (ry / 16) as usize;
291    // Overlap widths with the left/top MB column/row (the remaining area spills into
292    // the right/bottom neighbour when the block isn't MB-aligned).
293    let fx = (rx % 16) as f64;
294    let fy = (ry % 16) as f64;
295    let wl = 16.0 - fx; // area in column cx0
296    let wt = 16.0 - fy; // area in row    cy0
297    for (dy, wy) in [(0usize, wt), (1, fy)] {
298        if wy <= 0.0 {
299            continue;
300        }
301        let cy = (cy0 + dy).min(mb_h - 1);
302        for (dx, wx) in [(0usize, wl), (1, fx)] {
303            if wx <= 0.0 {
304                continue;
305            }
306            let cx = (cx0 + dx).min(mb_w - 1);
307            prev[cy * mb_w + cx] += amount * (wx * wy) / 256.0;
308        }
309    }
310}
311
312/// One frame's detector preparation: coded-size luma + its half-res plane —
313/// the per-FRAME half of the scene-cut pair ratio, split out so pair scorers
314/// can prepare each frame ONCE. Scoring a batch through `windows(2)` prepped
315/// every interior frame TWICE (as one pair's `cur`, then the next pair's
316/// `prev`); the rolling caches in `lookahead::segment_gops` and the streaming
317/// encoder hand the previous pair's `cur` prep across, halving the dominant
318/// detector cost. Same planes from the same functions: BIT-IDENTICAL ratios.
319pub(crate) struct PairPrep {
320    full: Vec<u8>,
321    half: Vec<u8>,
322}
323
324// Manual Debug: the derived form would dump both planes byte-by-byte through
325// the `Encoder` derive that requires this.
326impl std::fmt::Debug for PairPrep {
327    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
328        write!(f, "PairPrep({} + {} bytes)", self.full.len(), self.half.len())
329    }
330}
331
332pub(crate) fn pair_prep(cfg: &EncoderConfig, f: &YuvFrame) -> PairPrep {
333    let (mb_w, mb_h) = (cfg.mb_width(), cfg.mb_height());
334    let (cwf, chf) = (mb_w * 16, mb_h * 16);
335    let full = coded_luma(cfg, f);
336    let (half, _, _) = downsample2x(&full, cwf, chf);
337    PairPrep { full, half }
338}
339
340/// Scene-cut pair ratio `Σ min(inter, intra) / Σ intra` on the LOOKAHEAD
341/// estimator: half-res intra SATD vs half-res diamond-searched MC residual.
342/// The diamond (step 8→1) is the load-bearing part — the cheap ±2px activity
343/// set was already refuted IN THIS FILE for exactly the failure a scene-cut
344/// detector cannot afford: a fast pan reads as "unpredictable" when the
345/// probe cannot reach the true vector, and every pair looks like a cut.
346pub(crate) fn pair_ratio_prepped(cfg: &EncoderConfig, cur: &PairPrep, prev: &PairPrep) -> f64 {
347    let (mb_w, mb_h) = (cfg.mb_width(), cfg.mb_height());
348    let (cwf, chf) = (mb_w * 16, mb_h * 16);
349    let (cwh, chh) = (cwf / 2, chf / 2);
350    let costs = frame_costs(
351        &cur.full, cwf, chf, &cur.half, cwh, chh, mb_w, mb_h,
352        Some(&prev.full), Some(&prev.half), LookaheadMode::HalfRes,
353    );
354    let (mut num, mut den) = (0i64, 0i64);
355    for c in &costs {
356        num += c.inter.min(c.intra) as i64;
357        den += c.intra as i64;
358    }
359    num as f64 / den.max(1) as f64
360}
361
362/// mb-tree per-frame per-MB QP offsets for a GOP of SOURCE frames (display order,
363/// the IDR first). `strength <= 0` returns all-zero (no-op / byte-identical). The
364/// offsets are centered per GOP so the mean QP — hence the rate — is preserved.
365/// Per-GOP gate telemetry — the Front-B harvest seam.
366///
367/// The mb-tree latch decides PER GOP, so fitting a law on per-clip signals is a
368/// unit mismatch: within one clip the GOPs straddle the threshold (football
369/// measured 0.747 / 0.480 / 0.402). This records one row per GOP so the
370/// refinery can pair signals with a per-GOP objective.
371///
372/// Observe-only and off unless `RFF_MBTREE_GOPSTATS` is set.
373pub mod gopstats {
374    use std::sync::Mutex;
375    // GLOBAL, not thread_local: `encode_all` encodes GOPs IN PARALLEL, each on
376    // its own worker thread with a fresh encoder. Thread-local rows land on the
377    // worker and are invisible to the caller — the first version of this seam
378    // silently harvested ZERO rows for exactly that reason.
379    //
380    // Rows therefore arrive in worker-completion order, NOT GOP order. Harvest
381    // with `RUSTY_THREADS=1` so the order is the GOP order the objective is
382    // keyed by; `take()` refuses to guess otherwise.
383    static ROWS: Mutex<Vec<GopRow>> = Mutex::new(Vec::new());
384    /// One GOP's gate inputs and its latch decision.
385    #[derive(Debug, Clone, Copy)]
386    pub struct GopRow {
387        /// Strength-invariant offset dispersion (the latch signal).
388        pub sd: f64,
389        /// Raw RMS before normalisation (scales with mbtree_strength).
390        pub sd_raw: f64,
391        /// 1 - mean predictability; the back-off axis.
392        pub residual_frac: f64,
393        /// Effective strength after the back-off latch (0 = latched off).
394        pub eff_strength: f64,
395        /// Did the differentiation latch ZERO this GOP's offsets?
396        pub latched_off: bool,
397    }
398    pub fn on() -> bool {
399        static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
400        *V.get_or_init(|| std::env::var_os("RFF_MBTREE_GOPSTATS").is_some())
401    }
402    pub(crate) fn push(r: GopRow) {
403        if on() {
404            if let Ok(mut g) = ROWS.lock() {
405                g.push(r);
406            }
407        }
408    }
409    /// Drain the rows recorded so far.
410    ///
411    /// ⚠ In GOP order ONLY when the encode ran single-threaded
412    /// (`RUSTY_THREADS=1`). With parallel GOP encoding the rows interleave by
413    /// completion, and pairing them positionally with a per-GOP objective would
414    /// silently mismatch signals to outcomes.
415    pub fn take() -> Vec<GopRow> {
416        ROWS.lock().map(|mut g| std::mem::take(&mut *g)).unwrap_or_default()
417    }
418}
419
420/// Minimum propagation-offset dispersion for mb-tree to apply at all.
421/// `RFF_MBTREE_SDMIN=0` restores the ungated behaviour exactly.
422fn mbtree_spread_min(cfg: &EncoderConfig) -> f64 {
423    static V: std::sync::OnceLock<Option<f64>> = std::sync::OnceLock::new();
424    V.get_or_init(|| {
425        std::env::var("RFF_MBTREE_SDMIN")
426            .ok()
427            .and_then(|v| v.parse().ok())
428    })
429    .unwrap_or(cfg.mbtree_spread_min)
430}
431
432pub fn gop_qp_offsets(cfg: &EncoderConfig, frames: &[YuvFrame], strength: f64) -> Vec<Vec<i32>> {
433    gop_qp_offsets_refs(cfg, &frames.iter().collect::<Vec<_>>(), strength)
434}
435
436/// [`gop_qp_offsets`] over BORROWED frames. The bframes path's mb-tree runs on
437/// each GOP's anchor SUB-SEQUENCE — non-contiguous display indices — and used
438/// to materialize every window by DEEP-CLONING the anchor frames (a full
439/// Y+U+V copy per anchor per window, ~150 KB each at CIF, ~3 MB at 1080p).
440/// A slice of references carries the same frames with zero copies; the owned
441/// wrapper above keeps the contiguous callers unchanged.
442pub fn gop_qp_offsets_refs(cfg: &EncoderConfig, frames: &[&YuvFrame], strength: f64) -> Vec<Vec<i32>> {
443    let (mb_w, mb_h) = (cfg.mb_width(), cfg.mb_height());
444    let n = frames.len();
445    if strength <= 0.0 || n == 0 || mb_w * mb_h == 0 {
446        gopstats::push(gopstats::GopRow {
447            sd: 0.0, sd_raw: 0.0, residual_frac: 0.0, eff_strength: 0.0, latched_off: true,
448        });
449        return vec![vec![0i32; mb_w * mb_h]; n];
450    }
451    // Lookahead resolution mode (Hybrid default: half-res MV search + full-res cost
452    // scoring). `RFF_MBTREE_LA=full|hybrid|half` overrides for A/B.
453    let mode = match std::env::var("RFF_MBTREE_LA").as_deref() {
454        Ok("full") => LookaheadMode::FullRes,
455        Ok("hybrid") => LookaheadMode::Hybrid,
456        Ok("half") => LookaheadMode::HalfRes,
457        _ => cfg.mbtree_lookahead,
458    };
459    let (cwf, chf) = (mb_w * 16, mb_h * 16);
460    let (cwh, chh) = (mb_w * 8, mb_h * 8);
461    let full: Vec<Vec<u8>> = frames.iter().map(|f| coded_luma(cfg, f)).collect();
462    // GRAIN LATCH (Great Gate P3 item 1 — docs/gate-ledger.md mbtree-grain-veto):
463    // propagation credit is FICTION on noise (nothing persists), so mb-tree
464    // redistributes on false gradients — measured +4.41% BD-SSIM on grain once
465    // the AQ grain veto stopped masking it. Reuse the aq-grain-veto conjunction
466    // ("unexplained temporal residual: not texture, not motion → noise"),
467    // SOURCE-vs-SOURCE on the GOP's first frame pair — grain is stationary
468    // (per-frame floor spread was 7.7–8.3 over 120 frames), and both probe
469    // frames sit INSIDE the GOP, so a boundary scene cut cannot sit between
470    // them. Fires → zero offsets, byte-identical to mb-tree off for the GOP.
471    // Thresholds transfer from the AQ fit: its IDR arm already validated this
472    // exact conjunction on source-vs-source signals. `RFF_MBTREE_GRAIN=0`
473    // disables (bisection anchor). Single-frame GOPs fail open (no pair).
474    let grain_veto = {
475        static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
476        *ON.get_or_init(|| std::env::var("RFF_MBTREE_GRAIN").map(|s| s != "0").unwrap_or(true))
477    };
478    if grain_veto && n >= 2 {
479        let sig = crate::signals::FrameSignals::new(&full[1], cwf, mb_w, mb_h, Some(&full[0]));
480        let grain = sig.grain_signature();
481        crate::signals::census::bump(crate::signals::census::MBTREE_GRAIN, grain);
482        if grain {
483            if std::env::var("RFF_MBTREE_DBG").is_ok() {
484                eprintln!("MBTREE_DBG grain latch: eff=0.000 (zero offsets)");
485            }
486            // The grain veto IS a gate decision — record it, or the harvest
487            // drops the GOP and every later row pairs with the wrong objective.
488            gopstats::push(gopstats::GopRow {
489                sd: 0.0, sd_raw: 0.0, residual_frac: 0.0, eff_strength: 0.0, latched_off: true,
490            });
491            return vec![vec![0i32; mb_w * mb_h]; n];
492        }
493    }
494    // Half-res planes needed for Hybrid + HalfRes (the MV search); FullRes skips them.
495    let need_half = mode != LookaheadMode::FullRes;
496    let half: Vec<Vec<u8>> = if need_half {
497        full.iter().map(|f| downsample2x(f, cwf, chf).0).collect()
498    } else {
499        Vec::new()
500    };
501    let empty: Vec<u8> = Vec::new();
502    // 1. per-frame per-MB costs (frame 0 = IDR, intra-only).
503    let costs: Vec<Vec<MbCost>> = (0..n)
504        .map(|f| {
505            let ref_full = if f == 0 { None } else { Some(full[f - 1].as_slice()) };
506            let ref_half = if f == 0 || !need_half { None } else { Some(half[f - 1].as_slice()) };
507            let hf = if need_half { half[f].as_slice() } else { &empty[..] };
508            frame_costs(&full[f], cwf, chf, hf, cwh, chh, mb_w, mb_h, ref_full, ref_half, mode)
509        })
510        .collect();
511    // 2. backward propagation: each MB credits the fraction its predictor earned to
512    //    the previous frame's referenced MBs.
513    let mbs = mb_w * mb_h;
514    // ONE flat accumulator instead of `Vec<Vec<f64>>` (n+1 allocations → 1);
515    // the same `split_at_mut` cur/prev discipline, the same cells, the same
516    // values. `frac_buf` records each MB's `(intra−inter)/intra` as it is
517    // computed here, because the residual-fraction pass below was recomputing
518    // the IDENTICAL divide for every MB of every frame — reading the stored
519    // value in the same forward order keeps `fsum` bit-identical while
520    // removing (n−1)·mbs float divides. (Rows are frames 1..n.)
521    let mut propagate: Vec<f64> = vec![0.0; n * mbs];
522    let mut frac_buf: Vec<f64> = vec![0.0; n.saturating_sub(1) * mbs];
523    for f in (1..n).rev() {
524        let (head, tail) = propagate.split_at_mut(f * mbs);
525        let cur = &tail[..mbs];
526        let prev = &mut head[(f - 1) * mbs..];
527        for mb_y in 0..mb_h {
528            for mb_x in 0..mb_w {
529                let m = mb_y * mb_w + mb_x;
530                let c = costs[f][m];
531                let total = c.intra as f64 + cur[m];
532                // Fraction of this MB's cost the previous frame's reference "carries".
533                let frac = (c.intra - c.inter) as f64 / c.intra as f64; // in [0,1]
534                frac_buf[(f - 1) * mbs + m] = frac;
535                propagate_to(prev, mb_w, mb_h, mb_x, mb_y, c.mv, total * frac);
536            }
537        }
538    }
539    // CONTENT-ADAPTIVE STRENGTH (codec-content-adaptive-dispatch): mb-tree's benefit
540    // scales with how many future bits it can redistribute, ∝ the mean residual
541    // fraction `1 − pred` (pred = mean predictability `(intra−inter)/intra` over inter
542    // frames). When prediction is near-perfect (pred → 1, frames near-free/mostly
543    // skip — a slow, smooth pan) there is nothing to gain and QP perturbation only adds
544    // noise, so mb-tree REGRESSES; ramp strength to 0 as the residual fraction falls
545    // below `MBTREE_RES_MIN`. Natural detailed/mixed content sits well above it.
546    // PREDICTABILITY BACK-OFF, re-fitted 2026-08-06 (Great Gate P2 —
547    // docs/gate-ledger.md mbtree-backoff-refit). The original linear ramp
548    // (eff = strength·min(rf/0.10, 1)) had the right AXIS and the wrong SHAPE
549    // AND THRESHOLD: it throttled mb-tree's biggest real-content WINNERS
550    // (akiyo_qcif rf 0.046 → eff 0.44, forgoing −5.09→−9.24% BD-SSIM;
551    // screen_text rf 0.04 → eff 0.35, forgoing −4.53→−7.06) while the one
552    // class that genuinely regresses at full strength — tsrc-class synthetic
553    // (rf 0.023–0.025, +3.53% BD-SSIM unthrottled) — still leaked +0.50
554    // through the ramp's partial strength. The measured populations are
555    // disjoint with a 1.56× natural gap (tsrc ≤ 0.025 | winners ≥ 0.039), so
556    // the honest form is the single-sided LATCH: OFF below `res_min` (a zero
557    // qpo — byte-identical to mb-tree off for the GOP), FULL strength above.
558    // Per-GOP steps are safe: each GOP's offsets are independent and centered.
559    // `RFF_MBTREE_RESMIN` overrides (0 = no back-off, always full strength).
560    let res_min: f64 = {
561        static E: std::sync::OnceLock<f64> = std::sync::OnceLock::new();
562        *E.get_or_init(|| {
563            std::env::var("RFF_MBTREE_RESMIN").ok().and_then(|v| v.parse().ok()).unwrap_or(0.03)
564        })
565    };
566    // `frac_buf` holds frames 1..n in exactly this pass's iteration order, so
567    // the running sum sees the SAME summands in the SAME order — bit-identical
568    // `fsum` with zero divides. `fc` was `+= 1.0` per iteration: a count, and
569    // every increment is exact (integers < 2^53), so the closed form is the
570    // same number.
571    let fsum: f64 = frac_buf.iter().sum();
572    let fc = (n.saturating_sub(1) * mbs) as f64;
573    let residual_frac = 1.0 - if fc > 0.0 { fsum / fc } else { 0.0 };
574    let eff_strength = if res_min > 0.0 && residual_frac < res_min { 0.0 } else { strength };
575    crate::signals::census::bump(
576        crate::signals::census::MBTREE_BACKOFF,
577        eff_strength == 0.0,
578    );
579    if eff_strength == 0.0 {
580        gopstats::push(gopstats::GopRow {
581            sd: 0.0, sd_raw: 0.0, residual_frac, eff_strength: 0.0, latched_off: true,
582        });
583        // Latched off: zero offsets are byte-identical to mb-tree off.
584        if std::env::var("RFF_MBTREE_DBG").is_ok() {
585            eprintln!("MBTREE_DBG spread=0.000 residual_frac={residual_frac:.3} eff=0.000 (latched off)");
586        }
587        return vec![vec![0i32; mb_w * mb_h]; n];
588    }
589    // 3. QP offset per MB (≤ 0), then center per GOP to preserve the mean QP.
590    //
591    // A2 (fast-transcendentals addendum): `propagate == 0` ⇒ `total == intra`
592    // ⇒ the ratio is EXACTLY 1.0 (`intra >= 1` by construction, and `x/x` is
593    // exact) ⇒ `log2(1.0)` is exactly +0.0 — so the libm `log2` call AND the
594    // divide are skipped and the surviving multiply is the ORIGINAL expression
595    // evaluated at its exact value (`-eff_strength * 0.0` keeps the -0.0 the
596    // original produced). The whole LAST frame of every GOP takes this path —
597    // backward propagation never credits it — plus every MB the future never
598    // references; that is ≥ 1/n of all calls by construction, more on content
599    // with dead regions. Flat `offs` (n allocations → 1) with the same
600    // frame-major order everywhere downstream.
601    // Site 6's ★★ arm (Round 10): poly `log2` behind the same switch as the
602    // AQ pipeline; the zero-propagate shortcut is exact either way.
603    let poly = crate::fastmath::polytier_on();
604    let mut offs: Vec<f64> = Vec::with_capacity(n * mbs);
605    for f in 0..n {
606        for m in 0..mbs {
607            let p = propagate[f * mbs + m];
608            offs.push(if p == 0.0 {
609                -eff_strength * 0.0
610            } else {
611                let intra = costs[f][m].intra as f64;
612                let total = intra + p;
613                let l = if poly {
614                    crate::fastmath::log2_poly(total / intra)
615                } else {
616                    (total / intra).log2()
617                };
618                -eff_strength * l
619            });
620        }
621    }
622    // Per-GOP CENTERING: subtract the GOP-mean offset so the average QP — hence the
623    // rate — is preserved. This is the right rate-neutralization in BOTH modes: in CQP
624    // it holds the fixed QP; in RC mode (mb-tree runs per-GOP over the anchor chain, the
625    // controller picks the frame base) it keeps the offsets rate-neutral per GOP so the
626    // controller's model is undisturbed. MEASURED: routing the cross-frame allocation
627    // through the RC's `complexity` instead (uncentered per-MB + a per-frame multiplier)
628    // was WORSE — it destroyed the cross-frame redistribution the centered offsets carry
629    // (tsrc −1.5% → +1.9%). Centering stays; the RC just supplies the base QP.
630    let cnt = (n * mbs) as f64;
631    // Same frame-major summation order as the nested `flatten` had.
632    let mean: f64 = offs.iter().sum::<f64>() / cnt;
633    for o in offs.iter_mut() {
634        *o -= mean;
635    }
636    // DIFFERENTIATION LATCH (Great Gate P3 item 4 — the pan loser).
637    //
638    // mb-tree lowers QP on blocks whose quality PROPAGATES to the frames that
639    // reference them. That only buys anything when propagation is DIFFERENTIAL
640    // — some blocks matter much more than others (a static scene with a moving
641    // subject; screen content with dead regions). On a smooth pan EVERY block
642    // propagates about equally, so these centered offsets carry no information
643    // and mb-tree just redistributes rate for no perceptual reason. It is the
644    // clip class mb-tree has always lost on: stockholm +3.10, ducks +1.20,
645    // shields +1.06, bus +0.68, in_to_tree +0.69, crowd_run +0.35, city +0.27
646    // BD-SSIM, all REGRESSIONS IN THE SHIPPED DEFAULT before this latch.
647    //
648    // `sd` is the dispersion of mb-tree's OWN output. That is the whole point:
649    // it is not a content proxy standing in for the phenomenon (the fit that
650    // was refused here twice used `lv_spread`/`flat_run`, spatial statistics
651    // that merely correlate with panning on this corpus), it is the tool
652    // measuring whether it has anything to say. Below the line its offsets are
653    // noise around a centered mean, and zeroing them is EXACTLY mb-tree off.
654    //
655    // Measured over 21 clips: fires on 5, net +26.82 with ZERO regressions,
656    // versus +23.25 with SEVEN for always-on. Costs four modest forgone wins
657    // (tempete -1.61, football -1.39, soccer -0.59, crew -0.19).
658    //
659    // The refutation arm that killed the alternative: a second clause
660    // (`headroom>10 && tdecay>1.3`) recovered football/soccer at perfect fit on
661    // the 16-clip table, then fired on bus_cif -- a FAST PAN -- and LOST +0.68.
662    // Dropped. See gate-ledger `mbtree-dispatch`.
663    // STRENGTH-INVARIANT. Each offset is `-eff_strength * log2(total/intra)`, so
664    // a raw RMS scales LINEARLY with `mbtree_strength` — and the threshold was
665    // fitted at the default 0.9. At `--mbtree-strength 0.5` every sd shrinks 44%
666    // and clips that should fire latch OFF; at 2.0 the losers start firing.
667    // Dividing by eff_strength measures the DIFFERENTIATION itself (RMS of the
668    // log2 importance ratio) — what the gate is actually about, and invariant to
669    // how hard the offsets are then applied.
670    //
671    // Same defect class as the CAVLC bits/MB bug: a threshold on a signal whose
672    // SCALE depends on an axis the fitting corpus never varied.
673    let sd_raw = (offs.iter().map(|o| o * o).sum::<f64>() / cnt).sqrt();
674    let sd = sd_raw / eff_strength.max(1e-9);
675    if std::env::var("RFF_MBTREE_DBG").is_ok() {
676        eprintln!("MBTREE_DBG spread={sd:.3} raw={sd_raw:.3} residual_frac={residual_frac:.3} eff={eff_strength:.3}");
677    }
678    let sd_min = mbtree_spread_min(cfg);
679    gopstats::push(gopstats::GopRow {
680        sd, sd_raw, residual_frac, eff_strength, latched_off: sd < sd_min,
681    });
682    crate::signals::census::bump(crate::signals::census::MBTREE_SPREAD_LATCH, sd < sd_min);
683    if sd < sd_min {
684        return vec![vec![0i32; mb_w * mb_h]; n];
685    }
686    // Round + clamp to a sane per-MB QP swing. (Indexed slicing, not
687    // `chunks_exact` — the latter computes `len / mbs`, a runtime-divisor
688    // `div` this campaign exists to remove.) Site 8's ★★ arm: magic-number
689    // round behind the poly-tier switch.
690    const MBTREE_DQP_MAX: i32 = 6;
691    (0..n)
692        .map(|f| {
693            offs[f * mbs..][..mbs]
694                .iter()
695                .map(|&o| {
696                    let r = if poly { crate::fastmath::round_ties_even_fast(o) } else { o.round() };
697                    (r as i32).clamp(-MBTREE_DQP_MAX, MBTREE_DQP_MAX)
698                })
699                .collect()
700        })
701        .collect()
702}
703
704#[cfg(test)]
705mod tests {
706    use super::*;
707
708    /// Deterministic synthetic GOP: textured static background, a textured
709    /// 16x16 block translating 4px/frame, and a strip whose texture scrolls
710    /// 1px/frame — motion the diamond can track, with enough residual that the
711    /// predictability back-off would not zero it even if it were live.
712    fn synth_frames(w: usize, h: usize, n: usize) -> Vec<YuvFrame> {
713        (0..n)
714            .map(|f| {
715                let mut y = vec![0u8; w * h];
716                for j in 0..h {
717                    for i in 0..w {
718                        // Static checker+gradient background.
719                        let mut v = (((i / 4 + j / 4) % 2) as i32 * 60 + (i as i32 * 3 + j as i32 * 2) % 90 + 40) as i32;
720                        // Scrolling strip (1px/frame): partial predictability.
721                        if (16..24).contains(&j) {
722                            v = 30 + (((i + f) * 13) % 200) as i32;
723                        }
724                        // Moving textured square (4px/frame).
725                        let sx = 4 + f * 4;
726                        if (sx..sx + 16).contains(&i) && (28..44).contains(&j) {
727                            v = 200 - ((i * 7 + j * 11) % 120) as i32;
728                        }
729                        y[j * w + i] = v.clamp(0, 255) as u8;
730                    }
731                }
732                YuvFrame { width: w, height: h, y, u: vec![128; (w / 2) * (h / 2)], v: vec![128; (w / 2) * (h / 2)] }
733            })
734            .collect()
735    }
736
737    fn fnv1a_i32s(offs: &[Vec<i32>]) -> u64 {
738        let mut hsh: u64 = 0xcbf2_9ce4_8422_2325;
739        for fr in offs {
740            for &v in fr {
741                for b in v.to_le_bytes() {
742                    hsh ^= b as u64;
743                    hsh = hsh.wrapping_mul(0x100_0000_01b3);
744                }
745            }
746        }
747        hsh
748    }
749
750    /// The row-slice `coded_luma` must equal the original per-pixel
751    /// double-`min` gather (kept here as the oracle), including on frames
752    /// whose display size is not an MB multiple (live right/bottom padding).
753    #[test]
754    fn coded_luma_matches_per_pixel() {
755        for (w, h) in [(64usize, 48usize), (20, 13), (33, 17), (16, 16)] {
756            let cfg = EncoderConfig::new(w, h);
757            let frame = &synth_frames(w, h, 1)[0];
758            let (cw, ch) = (cfg.mb_width() * 16, cfg.mb_height() * 16);
759            let mut want = vec![0u8; cw * ch];
760            for j in 0..ch {
761                for i in 0..cw {
762                    want[j * cw + i] = frame.y[j.min(h - 1) * w + i.min(w - 1)];
763                }
764            }
765            assert_eq!(coded_luma(&cfg, frame), want, "{w}x{h}");
766        }
767    }
768
769    /// End-to-end golden gate for `gop_qp_offsets`: the exact `Vec<Vec<i32>>`
770    /// output is pinned by hash across all three lookahead modes, so any edit
771    /// claiming bit-identity has a whole-pipeline oracle (costs → propagation
772    /// → offsets → centering → latch → rounding). The grain veto and the
773    /// residual back-off are pinned OPEN via their documented env anchors
774    /// (`RFF_MBTREE_GRAIN=0`, `RFF_MBTREE_RESMIN=0` — both OnceLock-cached, so
775    /// they are set before the first call) so the golden pins the ARITHMETIC,
776    /// not a veto's early-out zeros.
777    #[test]
778    fn gop_qp_offsets_golden() {
779        std::env::set_var("RFF_MBTREE_GRAIN", "0");
780        std::env::set_var("RFF_MBTREE_RESMIN", "0");
781        let (w, h, n) = (64usize, 48usize, 6usize);
782        let frames = synth_frames(w, h, n);
783        let mut cfg = EncoderConfig::new(w, h);
784
785        crate::fastmath::TEST_POLYTIER.with(|c| c.set(Some(false)));
786        cfg.mbtree_lookahead = LookaheadMode::HalfRes;
787        let a = gop_qp_offsets(&cfg, &frames, 0.9);
788        // The gate must prove the tool ran: an all-zero output would pin only
789        // a latch, not the arithmetic under edit.
790        assert!(a.iter().flatten().any(|&v| v != 0), "HalfRes offsets all zero");
791        assert_eq!(fnv1a_i32s(&a), 1359955132549194384, "HalfRes golden");
792
793        cfg.mbtree_lookahead = LookaheadMode::Hybrid;
794        let b = gop_qp_offsets(&cfg, &frames, 0.9);
795        assert!(b.iter().flatten().any(|&v| v != 0), "Hybrid offsets all zero");
796        assert_eq!(fnv1a_i32s(&b), 7391805242828194773, "Hybrid golden");
797
798        cfg.mbtree_lookahead = LookaheadMode::FullRes;
799        let c = gop_qp_offsets(&cfg, &frames, 2.0);
800        assert!(c.iter().flatten().any(|&v| v != 0), "FullRes offsets all zero");
801        assert_eq!(fnv1a_i32s(&c), 17244099396955043453, "FullRes golden");
802
803        // Round 10 decision-identity: the POLY arm (poly log2 in the offs
804        // loop + ties-even final round) must yield the IDENTICAL i32 offsets
805        // in every mode — asserted, not argued.
806        crate::fastmath::TEST_POLYTIER.with(|c| c.set(Some(true)));
807        cfg.mbtree_lookahead = LookaheadMode::HalfRes;
808        assert_eq!(gop_qp_offsets(&cfg, &frames, 0.9), a, "poly arm HalfRes differs");
809        cfg.mbtree_lookahead = LookaheadMode::Hybrid;
810        assert_eq!(gop_qp_offsets(&cfg, &frames, 0.9), b, "poly arm Hybrid differs");
811        cfg.mbtree_lookahead = LookaheadMode::FullRes;
812        assert_eq!(gop_qp_offsets(&cfg, &frames, 2.0), c, "poly arm FullRes differs");
813        crate::fastmath::TEST_POLYTIER.with(|c| c.set(None));
814    }
815}