Skip to main content

rusty_h264_decoder/
lib.rs

1//! Pure-Rust H.264 decoder — Constrained Baseline **+ B-slices + most of High
2//! profile**, CAVLC and CABAC.
3//!
4//! Validated **bit-exact against Cisco's `h264dec`** on 35 of 35 clean streams
5//! from openh264's conformance corpus; the CABAC paths were brought up
6//! symbol-by-symbol against an instrumented openh264 oracle and are gated
7//! **pixel-exact vs ffmpeg**. The reconstruction path is shared with the encoder
8//! (via `rusty_h264-common`), so the two halves agree bit-for-bit by
9//! construction.
10//!
11//! Decodes: full intra (`I_16x16`/`I_4x4`/`I_8x8`/`I_PCM`), inter
12//! (`P_Skip`/16×16/16×8/8×16/`P_8x8`) with quarter-pel motion compensation,
13//! B-slices (temporal + spatial direct, implicit + explicit weighted
14//! prediction), the 8×8 transform and 8×8 intra prediction, scaling matrices,
15//! in-loop deblocking, and a multi-reference DPB with POC reordering and MMCO.
16//! CABAC covers I, P and B slices (not yet: `I_PCM`, High-profile 8×8 residual).
17//!
18//! This crate is `#![forbid(unsafe_code)]` and is **fuzzed to never panic or
19//! hang** on malformed input — errors surface as [`DecodeError`].
20//!
21//! [`Decoder::decode_stream`] is the one-call entry point (frames in display
22//! order); [`Decoder::decode`] is the streaming form (one picture per access
23//! unit, in decode order — pair it with [`Decoder::last_poc`]).
24
25mod cabac;
26/// Profile-only re-export of the CABAC bin census for benchmarking harnesses.
27#[cfg(feature = "profile")]
28pub use cabac::bin_census;
29mod mb16;
30mod params;
31
32pub use params::{Pps, Sps};
33pub use mb16::{MvField, MV_DUMP};
34
35/// Test-only re-export of the CABAC arithmetic *decoder* so the encoder crate can
36/// round-trip-validate its CABAC *encoder* against the exact reference engine.
37#[doc(hidden)]
38/// MEASUREMENT KNOB — `RFF_ABL_DEBLOCK=1` skips the loop filter so it can be
39/// priced by ablation on the UNINSTRUMENTED binary. Read once; inert when unset.
40fn abl_deblock() -> bool {
41    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
42    *ON.get_or_init(|| std::env::var("RFF_ABL_DEBLOCK").map_or(false, |v| v != "0"))
43}
44
45pub mod cabac_test {
46    pub use crate::cabac::Cabac;
47    pub use crate::mb16::b_inter_shape;
48    pub use crate::mb16::parse_mb_type_b;
49}
50
51use mb16::{FrameDecoder, GridPool, WeightTable};
52use rusty_h264_common::bit_reader::OutOfData;
53use rusty_h264_common::nal::{emulation_unprevent, split_annex_b};
54use rusty_h264_common::{BitReader, NalUnitType, YuvFrame};
55
56/// Decode errors.
57#[derive(Debug, Clone, PartialEq, Eq)]
58pub enum DecodeError {
59    /// Bitstream ended unexpectedly.
60    Truncated,
61    /// A required parameter set was missing before a slice.
62    MissingParameterSet,
63    /// A coding tool outside the implemented subset appeared in the stream.
64    Unsupported(&'static str),
65}
66
67impl From<OutOfData> for DecodeError {
68    fn from(_: OutOfData) -> Self {
69        DecodeError::Truncated
70    }
71}
72
73impl core::fmt::Display for DecodeError {
74    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
75        match self {
76            DecodeError::Truncated => f.write_str("bitstream truncated"),
77            DecodeError::MissingParameterSet => f.write_str("slice before SPS/PPS"),
78            DecodeError::Unsupported(s) => write!(f, "unsupported coding tool: {s}"),
79        }
80    }
81}
82
83impl std::error::Error for DecodeError {}
84
85/// A reference picture: deblocked reconstruction at coded resolution.
86/// Stored now (4a); read by motion compensation in 4b.
87/// Shared handle to a reference picture. The DPB and every per-slice reference
88/// list hold `Arc`s: list construction used to DEEP-CLONE each entry's planes +
89/// motion grids per slice (H-32 found ~600 KB+/slice of pure memcpy on B
90/// streams); an `Arc` clone is a refcount bump reading the same bytes, so the
91/// change is byte-identical by construction. `Arc::make_mut` covers the one
92/// mutation site (MMCO long-term marking).
93pub(crate) type Ref = std::sync::Arc<RefFrame>;
94
95#[derive(Debug, Clone, Default)]
96#[allow(dead_code)]
97pub(crate) struct RefFrame {
98    /// EDGE-PADDED planes (openh264 `ExpandPicture`): built ONCE per reference
99    /// frame so motion compensation reads them in place — the per-MC-call
100    /// clamped-tile extraction (~400 B copied per call, ~100 MB/clip on real
101    /// streams) dies with this. Luma pad [`LPAD`], chroma pad [`CPAD`]; strides
102    /// via [`RefFrame::lstride`]/[`RefFrame::cstride`].
103    pub py: Vec<u8>,
104    pub pu: Vec<u8>,
105    pub pv: Vec<u8>,
106    pub cw: usize,
107    pub ch: usize,
108    /// `frame_num` of the picture, for PicNum-based reference-list reordering.
109    pub frame_num: u32,
110    /// `PicOrderCnt` of the picture, for B-slice reference-list ordering.
111    pub poc: i32,
112    /// Per-4×4-block List-0 motion field (motion vector + reference index, `-1`
113    /// for intra), and the block-grid width. Read as the *co-located* picture's
114    /// motion for B-slice direct prediction (`colZeroFlag`, temporal direct).
115    pub mv: Vec<(i32, i32)>,
116    pub ref_idx: Vec<i32>,
117    /// Per-4×4-block **List-1** motion. Needed because the co-located motion
118    /// derivation (spec §8.4.1.2.1) falls back to List-1 when the co-located block
119    /// has no List-0 prediction (`predFlagL0Col == 0`). A co-located picture only
120    /// contains L1-only blocks when it is itself a B picture — which is precisely
121    /// what b-pyramid produces, so this stayed unexercised until B-references
122    /// appeared.
123    pub mv1: Vec<(i32, i32)>,
124    pub ref_idx1: Vec<i32>,
125    /// Per-4×4-block POC of the List-0 picture each block referenced (`i32::MIN`
126    /// for intra). Used by temporal direct's `MapColToList0` (the co-located
127    /// reference index alone is meaningless in the current list).
128    pub ref_poc: Vec<i32>,
129    pub w4: usize,
130    /// Long-term reference state. Long-term refs sit after short-term ones in
131    /// `RefPicList0` (ordered by `long_term_idx` ascending) and survive the
132    /// sliding window until explicitly unmarked (spec §8.2.4).
133    pub long_term: bool,
134    pub long_term_idx: u32,
135}
136
137/// Luma / chroma pad of every [`RefFrame`] plane. Luma 16 serves MVs overshooting
138/// the picture by up to ~14 px in place (chroma: half that, matching); wilder MVs
139/// take `mc_*_padded`'s clamped-halo fallback — correct, just slower.
140pub(crate) const LPAD: usize = 16;
141pub(crate) const CPAD: usize = 8;
142
143impl RefFrame {
144    #[inline]
145    pub fn lstride(&self) -> usize {
146        self.cw + 2 * LPAD
147    }
148    #[inline]
149    pub fn cstride(&self) -> usize {
150        self.cw / 2 + 2 * CPAD
151    }
152}
153
154/// A memory-management control operation (`dec_ref_pic_marking`, spec §7.4.3.3).
155#[derive(Clone, Copy)]
156enum Mmco {
157    /// 1: mark a short-term reference (by PicNum) as unused.
158    Unref(u32),
159    /// 2: mark a long-term reference (by LongTermPicNum) as unused.
160    UnrefLong(u32),
161    /// 3: assign a short-term reference (by PicNum) a LongTermFrameIdx.
162    AssignLong(u32, u32),
163    /// 4: drop long-term references with idx ≥ max_long_term_frame_idx_plus1.
164    MaxLong(u32),
165    /// 5: empty the DPB (and reset the current picture's frame_num to 0).
166    Reset,
167    /// 6: mark the current picture long-term with this LongTermFrameIdx.
168    CurrentLong(u32),
169}
170
171/// A picture being assembled from one or more slices (spec allows a picture to
172/// be split into multiple slices). Finalized — deblocked, output, and entered
173/// into the DPB — once all its macroblocks are decoded.
174struct PendingPic {
175    fd: mb16::FrameDecoder,
176    frame_num: u32,
177    poc: i32,
178    next_mb: usize,
179    total_mb: usize,
180    slice_count: u16,
181    deblock: bool,
182    filter_offset_a: i32,
183    filter_offset_b: i32,
184    crop_r: usize,
185    crop_b: usize,
186    max_refs: usize,
187    log2_max_frame_num: u32,
188    /// `false` for a non-reference picture (nal_ref_idc == 0): output it but do
189    /// not enter it into the DPB.
190    is_reference: bool,
191    idr_long_term: bool,
192    mmco_ops: Vec<Mmco>,
193}
194
195/// Measurement knob: disable grid pooling, restoring per-picture allocation.
196fn no_pool() -> bool {
197    use std::sync::atomic::{AtomicU8, Ordering};
198    static ON: AtomicU8 = AtomicU8::new(0);
199    match ON.load(Ordering::Relaxed) {
200        0 => {
201            let v = std::env::var_os("RS_H264_NO_POOL").is_some_and(|v| v == "1");
202            ON.store(if v { 1 } else { 2 }, Ordering::Relaxed);
203            v
204        }
205        n => n == 1,
206    }
207}
208
209/// A Constrained Baseline H.264 decoder. Holds the most recent parameter sets
210/// and the previous decoded picture (the inter reference) across calls.
211#[derive(Default)]
212pub struct Decoder {
213    /// Active parameter sets, keyed by id — a stream may carry several and switch
214    /// between them per slice (spec §7.3.2.1/.2).
215    sps: std::collections::HashMap<u32, Sps>,
216    pps: std::collections::HashMap<u32, Pps>,
217    /// Decoded-picture buffer (most-recent first); `ref_idx` indexes into this.
218    refs: Vec<Ref>,
219    /// The picture currently being assembled from its slices, if any.
220    cur: Option<PendingPic>,
221    /// Picture-order-count state (spec §8.2.1). Tracks the previous reference
222    /// picture's MSB/LSB (type 0) and frame-num offset (types 1/2) so display
223    /// order can be recovered — needed once B-pictures (out-of-order) land.
224    poc: PocState,
225    /// `PicOrderCnt` of the most recently returned picture (display-order key).
226    last_poc: i32,
227    /// `frame_num` of the previous short-term reference picture, for detecting
228    /// gaps in `frame_num` (spec §8.2.5.2).
229    prev_ref_frame_num: u32,
230    /// Per-picture grid allocations, handed from the finished picture to the next
231    /// one instead of being freed and re-allocated. See `mb16::GridPool`.
232    grid_pool: GridPool,
233    /// Recycled padded-plane buffers from evicted reference frames, drawn by
234    /// `as_reference_pooled`. Bounded (see `reclaim_retired`).
235    plane_pool: Vec<Vec<u8>>,
236    /// Reference frames evicted from the DPB whose planes have not been
237    /// reclaimed yet. Reclamation must wait until the evicting picture's
238    /// `FrameDecoder` is consumed — while it lives it still holds `Arc` clones
239    /// of its ref lists, so `Arc::try_unwrap` would fail at eviction time.
240    retired: Vec<Ref>,
241}
242
243/// Running picture-order-count derivation state.
244#[derive(Default)]
245struct PocState {
246    prev_msb: i32,
247    prev_lsb: i32,
248    prev_frame_num: u32,
249    prev_frame_num_offset: i64,
250}
251
252impl Decoder {
253    /// Creates a decoder with no parameter sets yet.
254    pub fn new() -> Self {
255        Self::default()
256    }
257
258    /// Decodes a complete Annex-B access unit, returning the reconstructed,
259    /// cropped frame if the access unit contained a coded picture.
260    pub fn decode(&mut self, annex_b: &[u8]) -> Result<Option<YuvFrame>, DecodeError> {
261        let _g = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Total);
262        let mut frame = None;
263        for nal in split_annex_b(annex_b) {
264            if nal.is_empty() {
265                continue;
266            }
267            let nal_type = NalUnitType::from_id(nal[0]);
268            let rbsp = emulation_unprevent(&nal[1..]);
269            match nal_type {
270                NalUnitType::Sps => {
271                    let s = Sps::parse(&rbsp)?;
272                    self.sps.insert(s.seq_parameter_set_id, s);
273                }
274                NalUnitType::Pps => {
275                    let p = Pps::parse(&rbsp)?;
276                    self.pps.insert(p.pic_parameter_set_id, p);
277                }
278                NalUnitType::IdrSlice | NalUnitType::NonIdrSlice => {
279                    let nal_ref_idc = (nal[0] >> 5) & 3;
280                    let is_idr = nal_type == NalUnitType::IdrSlice;
281                    if let Some(f) = self.decode_slice(&rbsp, is_idr, nal_ref_idc)? {
282                        frame = Some(f);
283                    }
284                }
285                _ => {} // SEI, AUD, etc. ignored
286            }
287        }
288        Ok(frame)
289    }
290
291    /// Decodes a complete Annex-B byte stream and returns every picture in
292    /// **display order** (`PicOrderCnt` within each GOP; an IDR ends a GOP).
293    ///
294    /// This is the convenient whole-stream entry point — it handles access-unit
295    /// splitting, multi-slice picture assembly, and B-picture reordering — versus
296    /// the lower-level per-access-unit [`Decoder::decode`], which returns pictures
297    /// in decode order.
298    pub fn decode_stream(&mut self, annex_b: &[u8]) -> Result<Vec<YuvFrame>, DecodeError> {
299        let mut out = Vec::new();
300        let mut gop: Vec<(i32, YuvFrame)> = Vec::new();
301        for au in split_access_units(annex_b) {
302            if au_is_idr(au) {
303                flush_gop(&mut gop, &mut out); // emit the prior GOP before the IDR
304            }
305            if let Some(frame) = self.decode(au)? {
306                gop.push((self.last_poc, frame));
307            }
308        }
309        flush_gop(&mut gop, &mut out);
310        Ok(out)
311    }
312
313    fn decode_slice(
314        &mut self,
315        rbsp: &[u8],
316        is_idr: bool,
317        nal_ref_idc: u8,
318    ) -> Result<Option<YuvFrame>, DecodeError> {
319        let mut r = BitReader::new(rbsp);
320        // --- slice_header ---
321        let first_mb_in_slice = r.read_ue()? as usize;
322        let slice_type = r.read_ue()?;
323        let is_p = matches!(slice_type, 0 | 5);
324        let is_b = matches!(slice_type, 1 | 6);
325        let is_i = matches!(slice_type, 2 | 7);
326        if !is_p && !is_b && !is_i {
327            return Err(DecodeError::Unsupported("SP/SI slices"));
328        }
329        // Resolve the parameter sets this slice references (by id).
330        let pic_parameter_set_id = r.read_ue()?;
331        let pps = self.pps.get(&pic_parameter_set_id).cloned().ok_or(DecodeError::MissingParameterSet)?;
332        let sps = self.sps.get(&pps.seq_parameter_set_id).cloned().ok_or(DecodeError::MissingParameterSet)?;
333        let sps = &sps;
334        let pps = &pps;
335        // CABAC (entropy_coding_mode_flag=1) has an entirely different slice-data parse
336        // (docs/cabac-decode-plan.md). I-slice CABAC is being brought up; the CABAC MB
337        // loop gates P/B until Phase 3. `cabac_init_idc` (P/B only) is read below.
338        let cabac = pps.entropy_coding_mode_flag;
339        let frame_num = r.read_bits(sps.log2_max_frame_num)?;
340        if is_idr {
341            let _idr_pic_id = r.read_ue()?;
342        }
343        // pic_order_cnt fields (spec §7.3.3). `field_pic_flag` is always 0
344        // (frame_mbs_only). Captured to derive PicOrderCnt for display ordering.
345        let mut poc_lsb = 0u32;
346        let mut delta_poc_bottom = 0i32;
347        if sps.pic_order_cnt_type == 0 {
348            poc_lsb = r.read_bits(sps.log2_max_pic_order_cnt_lsb)?;
349            if pps.bottom_field_pic_order_present {
350                delta_poc_bottom = r.read_se()?;
351            }
352        } else if sps.pic_order_cnt_type == 1 && !sps.delta_pic_order_always_zero {
353            let _delta_pic_order_cnt_0 = r.read_se()?;
354            if pps.bottom_field_pic_order_present {
355                let _delta_pic_order_cnt_1 = r.read_se()?;
356            }
357        }
358        // PicOrderCnt is determined by the first slice of the picture; later
359        // slices share it (and must not re-advance the POC state).
360        let pic_poc = if first_mb_in_slice == 0 {
361            self.compute_poc(sps, is_idr, nal_ref_idc, frame_num, poc_lsb, delta_poc_bottom)
362        } else {
363            self.cur.as_ref().map_or(0, |p| p.poc)
364        };
365        // redundant_pic_cnt: a non-zero value marks a *redundant* coded picture
366        // (an alternative representation of the primary picture). A primary
367        // decoder discards it (spec §7.4.3, §8.2.5 note). Must be read here or the
368        // rest of the slice header desyncs.
369        if pps.redundant_pic_cnt_present_flag {
370            let redundant_pic_cnt = r.read_ue()?;
371            if redundant_pic_cnt != 0 {
372                return Ok(None);
373            }
374        }
375        if std::env::var_os("RH264_DUMP_MB").is_some() {
376            eprintln!(
377                "SLICE fn={frame_num} poc={pic_poc} nal_ref_idc={nal_ref_idc} is_p={is_p} is_b={is_b} first_mb={first_mb_in_slice}"
378            );
379        }
380        // B slices choose direct-mode derivation here (spec §7.3.3).
381        let direct_spatial = if is_b { r.read_bit()? } else { true };
382        let mut num_ref_idx_l0 = pps.num_ref_idx_l0_default as usize;
383        let mut num_ref_idx_l1 = pps.num_ref_idx_l1_default as usize;
384        let mut reorder_l0: Vec<(u32, u32)> = Vec::new();
385        let mut reorder_l1: Vec<(u32, u32)> = Vec::new();
386        if is_p || is_b {
387            // num_ref_idx_active_override_flag
388            if r.read_bit()? {
389                num_ref_idx_l0 = (r.read_ue()? + 1) as usize;
390                if is_b {
391                    num_ref_idx_l1 = (r.read_ue()? + 1) as usize;
392                }
393            }
394            // ref_pic_list_modification_flag_l0
395            if r.read_bit()? {
396                parse_ref_pic_list_modification(&mut r, &mut reorder_l0)?;
397            }
398            if is_b && r.read_bit()? {
399                // ref_pic_list_modification_flag_l1
400                parse_ref_pic_list_modification(&mut r, &mut reorder_l1)?;
401            }
402        }
403        // Explicit weighted prediction carries a pred_weight_table() here. P
404        // (weighted_pred) uses single-list weights; B explicit bipred (idc 1) is
405        // not yet wired into the bi-pred averaging, so refuse that. Implicit
406        // bipred (idc 2) carries no table.
407        let weights = if is_p && pps.weighted_pred {
408            Some(parse_pred_weight_table(&mut r, num_ref_idx_l0, 0, false)?)
409        } else if is_b && pps.weighted_bipred_idc == 1 {
410            return Err(DecodeError::Unsupported("explicit B weighted prediction"));
411        } else {
412            None
413        };
414        // dec_ref_pic_marking (spec §7.3.3.3) — present only for reference
415        // pictures (nal_ref_idc != 0). Reading it for a non-reference slice would
416        // desync the rest of the header.
417        let mut idr_long_term = false;
418        let mut mmco_ops: Vec<Mmco> = Vec::new();
419        if nal_ref_idc == 0 {
420            // non-reference picture: no marking syntax
421        } else if is_idr {
422            let _no_output_of_prior_pics = r.read_bit()?;
423            idr_long_term = r.read_bit()?; // long_term_reference_flag
424        } else if r.read_bit()? {
425            // adaptive_ref_pic_marking_mode_flag
426            loop {
427                let op = r.read_ue()?;
428                match op {
429                    0 => break,
430                    1 => mmco_ops.push(Mmco::Unref(r.read_ue()?)),
431                    2 => mmco_ops.push(Mmco::UnrefLong(r.read_ue()?)),
432                    3 => {
433                        let diff = r.read_ue()?;
434                        let idx = r.read_ue()?;
435                        mmco_ops.push(Mmco::AssignLong(diff, idx));
436                    }
437                    4 => mmco_ops.push(Mmco::MaxLong(r.read_ue()?)),
438                    5 => mmco_ops.push(Mmco::Reset),
439                    6 => mmco_ops.push(Mmco::CurrentLong(r.read_ue()?)),
440                    _ => return Err(DecodeError::Unsupported("invalid MMCO")),
441                }
442                if mmco_ops.len() > 128 {
443                    return Err(DecodeError::Truncated);
444                }
445            }
446        }
447        // cabac_init_idc (spec §7.3.3) — CABAC context-model preset, P/B slices only.
448        // Spec range [0,2]; a larger (corrupt) value would index the 4-model context-init
449        // table out of bounds, so reject it here.
450        let cabac_init_idc = if cabac && !is_i {
451            let v = r.read_ue()?;
452            if v > 2 {
453                return Err(DecodeError::Unsupported("invalid cabac_init_idc"));
454            }
455            v
456        } else {
457            0
458        };
459        let slice_qp_delta = r.read_se()?;
460        // When deblocking_filter_control_present_flag is 0 the slice carries no
461        // disable_deblocking_filter_idc and it is inferred 0 — i.e. the in-loop
462        // filter is ON by default (spec §7.4.3). (Our own encoder always signals
463        // the control explicitly, so this default was previously untested.)
464        let mut deblock = true;
465        let (mut filter_offset_a, mut filter_offset_b) = (0i32, 0i32);
466        if pps.deblocking_filter_control_present_flag {
467            let disable_deblocking_filter_idc = r.read_ue()?;
468            // idc 1 = filter off; idc 0 = on; idc 2 = on but not across slice
469            // boundaries (equivalent to on for single-slice pictures).
470            deblock = disable_deblocking_filter_idc != 1;
471            if disable_deblocking_filter_idc != 1 {
472                // FilterOffset = slice_*_offset_div2 × 2 (spec §7.4.3).
473                filter_offset_a = r.read_se()? * 2;
474                filter_offset_b = r.read_se()? * 2;
475            }
476        }
477        let slice_qp = (pps.pic_init_qp + slice_qp_delta).clamp(0, 51) as u8;
478
479        // Synthesize placeholder short-term references for any gap in frame_num
480        // (spec §8.2.5.2) so the DPB / PicNum mapping stays correct.
481        if first_mb_in_slice == 0 && !is_idr && sps.gaps_in_frame_num_allowed {
482            self.insert_frame_num_gaps(
483                frame_num,
484                1u32 << sps.log2_max_frame_num,
485                sps.max_num_ref_frames.max(1) as usize,
486                sps.pic_width_in_mbs * 16,
487                sps.pic_height_in_mbs * 16,
488            );
489        }
490
491        // Build the reference list(s) for this slice. P uses RefPicList0 only;
492        // B uses RefPicList0 and RefPicList1 (POC-ordered).
493        let max_fn = 1u32 << sps.log2_max_frame_num;
494        let (ref_list0, ref_list1) = if is_b {
495            build_ref_list_b(
496                &self.refs, pic_poc, frame_num, max_fn,
497                num_ref_idx_l0, num_ref_idx_l1, &reorder_l0, &reorder_l1,
498            )?
499        } else if is_p {
500            (build_ref_list_p(&self.refs, frame_num, max_fn, num_ref_idx_l0, &reorder_l0)?, Vec::new())
501        } else {
502            (Vec::new(), Vec::new())
503        };
504        // --- picture assembly ---
505        // first_mb_in_slice == 0 starts a new picture; otherwise this slice
506        // continues the one in flight. An IDR clears the DPB at its first slice.
507        if first_mb_in_slice == 0 {
508            if is_idr {
509                self.refs.clear();
510            }
511            // H-49: the CABAC macroblock loop never decodes `transform_size_8x8_flag`
512            // — both reads of it sit on the CAVLC `BitReader`, and `decode_i8x8` only
513            // accepts one. A PPS with `transform_8x8_mode_flag` set therefore desyncs
514            // the arithmetic decoder within a few macroblocks, and the failure surfaces
515            // as a bogus `CABAC I_PCM` far from its cause (the mb_type parse lands on
516            // 25 out of garbage). Fail fast and accurately instead: a wrong error that
517            // points at the wrong feature costs more than a missing feature does.
518            // Removing this guard requires the CABAC 8×8 residual path — see H-49.
519            // DecSetup was declared in the Stage enum but never actually scoped, so
520            // the per-picture grid allocation had been invisible in every profile.
521            let _g_setup = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DecSetup);
522            let mut fd = FrameDecoder::with_pool(
523                sps.pic_width_in_mbs,
524                sps.pic_height_in_mbs,
525                slice_qp,
526                pps.chroma_qp_index_offset,
527                ref_list0,
528                num_ref_idx_l0,
529                pps.constrained_intra_pred_flag,
530                pps.transform_8x8_mode_flag,
531                sps.profile_idc != 66, // b_possible: Baseline/Constrained Baseline (66) forbid B
532                // `RS_H264_NO_POOL=1` reproduces the pre-pool behaviour exactly (a
533                // fresh allocation per picture) so the pool can be A/B'd paired on
534                // one binary, with no rebuild between arms.
535                if no_pool() { GridPool::default() } else { std::mem::take(&mut self.grid_pool) },
536            );
537            if is_b {
538                fd.set_b_context(
539                    ref_list1,
540                    num_ref_idx_l1,
541                    direct_spatial,
542                    pic_poc,
543                    pps.weighted_bipred_idc,
544                    sps.direct_8x8_inference,
545                );
546            }
547            if sps.has_scaling || pps.pic_scaling_matrix_present {
548                let (s4, s8) = resolve_scaling(sps, pps);
549                fd.set_scaling(s4, s8);
550            }
551            if let Some(w) = weights {
552                fd.set_weights(w);
553            }
554            // A pending picture still here means the previous one never reached
555            // total_mb and a new picture is now displacing it. That is a DECODER
556            // desync, and dropping it silently is how the missing-B-slice-ref_idx
557            // defect stayed hidden: the picture simply never entered the DPB, and
558            // the failure surfaced hundreds of macroblocks later as "bitstream
559            // truncated" from a reference-list modification asking for it. Refuse
560            // to swallow it -- an incomplete picture must announce itself.
561            if let Some(prev) = self.cur.take() {
562                if prev.next_mb < prev.total_mb {
563                    return Err(DecodeError::Truncated);
564                }
565            }
566            self.cur = Some(PendingPic {
567                fd,
568                frame_num,
569                poc: pic_poc,
570                next_mb: 0,
571                total_mb: sps.pic_width_in_mbs * sps.pic_height_in_mbs,
572                slice_count: 0,
573                deblock,
574                filter_offset_a,
575                filter_offset_b,
576                crop_r: sps.frame_crop_right as usize,
577                crop_b: sps.frame_crop_bottom as usize,
578                max_refs: sps.max_num_ref_frames.max(1) as usize,
579                log2_max_frame_num: sps.log2_max_frame_num,
580                is_reference: nal_ref_idc != 0,
581                idr_long_term,
582                mmco_ops,
583            });
584        } else {
585            // Continuation slice: reset the per-slice QP + reference list.
586            let Some(pic) = self.cur.as_mut() else {
587                return Err(DecodeError::Unsupported("slice continues a missing picture"));
588            };
589            pic.fd.begin_slice(slice_qp, ref_list0, num_ref_idx_l0);
590            if is_b {
591                pic.fd.set_b_context(
592                    ref_list1,
593                    num_ref_idx_l1,
594                    direct_spatial,
595                    pic.poc,
596                    pps.weighted_bipred_idc,
597                    sps.direct_8x8_inference,
598                );
599            }
600            if sps.has_scaling || pps.pic_scaling_matrix_present {
601                let (s4, s8) = resolve_scaling(sps, pps);
602                pic.fd.set_scaling(s4, s8);
603            }
604            if let Some(w) = weights {
605                pic.fd.set_weights(w);
606            }
607            // Latest slice's marking/deblock parameters win at finalization.
608            pic.deblock = deblock;
609            pic.filter_offset_a = filter_offset_a;
610            pic.filter_offset_b = filter_offset_b;
611            pic.idr_long_term |= idr_long_term;
612            pic.mmco_ops.extend(mmco_ops);
613        }
614
615        let pic = self.cur.as_mut().expect("pending picture set above");
616        // Row-interleave (mb16::row_hook) needs the CURRENT slice's deblock
617        // parameters during decode; `abl_deblock` resolved here so mb16 stays
618        // knob-agnostic.
619        pic.fd.set_deblock_params(deblock && !abl_deblock(), filter_offset_a, filter_offset_b);
620        let first = first_mb_in_slice.min(pic.total_mb);
621        let next = if cabac {
622            // cabac_alignment_one_bit → the slice data is byte-aligned from here.
623            r.align_to_byte().map_err(|_| DecodeError::Truncated)?;
624            let (data, start) = (r.data(), r.bit_pos() / 8);
625            pic.fd
626                .decode_slice_data_cabac(data, start, slice_qp, cabac_init_idc, is_i, is_p, first)
627        } else {
628            pic.fd.decode_slice_data(&mut r, is_p, first)
629        }
630        .map_err(|e| match e {
631            mb16::MbError::Truncated => DecodeError::Truncated,
632            mb16::MbError::Unsupported(s) => DecodeError::Unsupported(s),
633        })?;
634        pic.next_mb = next;
635        pic.slice_count += 1;
636        if std::env::var_os("RH264_DUMP_MB").is_some() {
637            eprintln!(
638                "  slice decoded {}/{} MBs{}",
639                next,
640                pic.total_mb,
641                if next < pic.total_mb { "   <-- INCOMPLETE" } else { "" }
642            );
643        }
644
645        if pic.next_mb < pic.total_mb {
646            return Ok(None); // picture not yet complete
647        }
648
649        // --- finalize the completed picture ---
650        let pic = self.cur.take().expect("pending picture");
651        let PendingPic {
652            mut fd,
653            frame_num,
654            poc,
655            deblock,
656            filter_offset_a,
657            filter_offset_b,
658            crop_r,
659            crop_b,
660            max_refs,
661            log2_max_frame_num,
662            is_reference,
663            idr_long_term,
664            mmco_ops,
665            ..
666        } = pic;
667        self.last_poc = poc;
668        // MEASUREMENT KNOB (`RFF_ABL_DEBLOCK=1`): skip the loop filter to price it
669        // with ZERO instrument tax. The scope-based profiler charges an rdtsc pair
670        // per scope, and at ~20M per-MB scopes that tax reached 1.3-1.4x of the
671        // whole decode -- so a per-MB stage's share cannot be read off it. Ablation
672        // on the UNINSTRUMENTED binary is the honest price. Output is wrong while
673        // set; decode WORK is unchanged (the filter reads and writes samples but
674        // decides nothing), so the timing stays comparable.
675        if deblock && !abl_deblock() {
676            fd.deblock(filter_offset_a, filter_offset_b);
677        }
678        // The necessary DPB plane clone (rec_y/u/v → RefFrame) — measured as its own
679        // stage, OUTSIDE the Finalize scope so the two don't double-count.
680        let reference = if is_reference {
681            let _dg = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::DpbClone);
682            Some(fd.as_reference_pooled(&mut self.plane_pool))
683        } else {
684            // A non-reference picture is output but never enters the DPB.
685            None
686        };
687        let _fg = rusty_h264_common::prof::scope(rusty_h264_common::prof::Stage::Finalize);
688        if let Some(mut reference) = reference {
689            reference.frame_num = frame_num;
690            reference.poc = poc;
691            if std::env::var_os("RH264_DUMP_MB").is_some() {
692                eprintln!("DPB-ADD fn={frame_num} poc={poc}");
693            }
694            if idr_long_term {
695                reference.long_term = true;
696                reference.long_term_idx = 0;
697            }
698            // Track the reference frame_num for gap detection (0 after MMCO 5).
699            self.prev_ref_frame_num =
700                self.apply_ref_marking(reference, &mmco_ops, frame_num, log2_max_frame_num, max_refs);
701        }
702        let (frame, pool) = fd.into_frame_recycle(crop_r, crop_b);
703        self.grid_pool = pool;
704        self.reclaim_retired();
705        Ok(Some(frame))
706    }
707
708    /// Moves the padded planes of retired (DPB-evicted) reference frames into
709    /// the recycle pool. Called after the current picture's `FrameDecoder` is
710    /// consumed, at which point a retired frame's `Arc` is normally unique; a
711    /// frame something still holds (it shouldn't) is simply dropped un-recycled.
712    fn reclaim_retired(&mut self) {
713        for arc in self.retired.drain(..) {
714            if let Ok(rf) = std::sync::Arc::try_unwrap(arc) {
715                self.plane_pool.push(rf.py);
716                self.plane_pool.push(rf.pu);
717                self.plane_pool.push(rf.pv);
718            }
719        }
720        // Bound the pool: 6 pictures' worth of planes (3 each) covers any
721        // realistic ref churn; beyond that we'd just be hoarding memory.
722        self.plane_pool.truncate(18);
723    }
724
725    /// Inserts "non-existing" short-term reference frames for each `frame_num`
726    /// skipped since the previous reference picture (spec §8.2.5.2). Their samples
727    /// are unspecified (a conformant stream never references them); we use mid-grey
728    /// so any accidental reference is benign. They occupy DPB slots and advance the
729    /// sliding window, keeping PicNum/ref-list derivation correct.
730    fn insert_frame_num_gaps(&mut self, frame_num: u32, max_fn: u32, max_refs: usize, w: usize, h: usize) {
731        if max_fn == 0 {
732            return;
733        }
734        let start = (self.prev_ref_frame_num + 1) % max_fn;
735        let gap = (frame_num + max_fn - start) % max_fn;
736        if gap == 0 {
737            return;
738        }
739        // Each placeholder is inserted at the front then the DPB is truncated to
740        // `max_refs`, so for a gap larger than that only the most recent `max_refs`
741        // placeholders can survive. Materialise just those — a malformed stream can
742        // declare a gap of MaxFrameNum-1 (up to 65535), and allocating that many
743        // full frames would be a CPU/memory DoS.
744        let cap = max_refs.max(1);
745        let n = (gap as usize).min(cap);
746        let (cw, ch) = (w, h);
747        let mut expected = (frame_num + max_fn - n as u32) % max_fn;
748        for _ in 0..n {
749            self.refs.insert(
750                0,
751                std::sync::Arc::new(RefFrame {
752                    // Uniform grey: the padded plane of a uniform plane is itself.
753                    py: vec![128; (cw + 2 * LPAD) * (ch + 2 * LPAD)],
754                    pu: vec![128; (cw / 2 + 2 * CPAD) * (ch / 2 + 2 * CPAD)],
755                    pv: vec![128; (cw / 2 + 2 * CPAD) * (ch / 2 + 2 * CPAD)],
756                    cw,
757                    ch,
758                    frame_num: expected,
759                    poc: 0,
760                    mv: Vec::new(),
761                    ref_idx: Vec::new(),
762                    mv1: Vec::new(),
763                    ref_idx1: Vec::new(),
764                    ref_poc: Vec::new(),
765                    w4: 0,
766                    long_term: false,
767                    long_term_idx: 0,
768                }),
769            );
770            self.refs.truncate(cap);
771            expected = (expected + 1) % max_fn;
772        }
773        self.prev_ref_frame_num = (frame_num + max_fn - 1) % max_fn;
774    }
775
776    /// The `PicOrderCnt` of the most recently returned picture. Pictures are
777    /// returned in decode order; sorting them by this value yields display order
778    /// (the only difference is reordered B-pictures).
779    pub fn last_poc(&self) -> i32 {
780        self.last_poc
781    }
782
783    /// Derives `PicOrderCnt` for the current picture (spec §8.2.1) and advances
784    /// the POC state. Types 0 and 2 are exact; type 1 is approximated by
785    /// frame-num order (no B-stream in scope uses it).
786    fn compute_poc(
787        &mut self,
788        sps: &Sps,
789        is_idr: bool,
790        nal_ref_idc: u8,
791        frame_num: u32,
792        poc_lsb: u32,
793        delta_bottom: i32,
794    ) -> i32 {
795        match sps.pic_order_cnt_type {
796            0 => {
797                let max_lsb = 1i32 << sps.log2_max_pic_order_cnt_lsb;
798                let (prev_msb, prev_lsb) =
799                    if is_idr { (0, 0) } else { (self.poc.prev_msb, self.poc.prev_lsb) };
800                let lsb = poc_lsb as i32;
801                let msb = if lsb < prev_lsb && prev_lsb - lsb >= max_lsb / 2 {
802                    prev_msb + max_lsb
803                } else if lsb > prev_lsb && lsb - prev_lsb > max_lsb / 2 {
804                    prev_msb - max_lsb
805                } else {
806                    prev_msb
807                };
808                let top = msb + lsb;
809                let poc = top.min(top + delta_bottom);
810                if nal_ref_idc != 0 {
811                    self.poc.prev_msb = msb;
812                    self.poc.prev_lsb = lsb;
813                }
814                poc
815            }
816            2 => {
817                let max_fn = 1i64 << sps.log2_max_frame_num;
818                let offset = if is_idr {
819                    0
820                } else if self.poc.prev_frame_num > frame_num {
821                    self.poc.prev_frame_num_offset + max_fn
822                } else {
823                    self.poc.prev_frame_num_offset
824                };
825                let poc = if is_idr {
826                    0
827                } else {
828                    2 * (offset + frame_num as i64) - i64::from(nal_ref_idc == 0)
829                };
830                self.poc.prev_frame_num_offset = offset;
831                self.poc.prev_frame_num = frame_num;
832                poc as i32
833            }
834            _ => {
835                self.poc.prev_frame_num = frame_num;
836                frame_num as i32 * 2
837            }
838        }
839    }
840
841    /// Inserts the just-decoded picture into the DPB and marks references
842    /// (spec §8.2.5). With no MMCO commands this is the sliding window (evict the
843    /// oldest short-term reference past capacity); with MMCO it is adaptive
844    /// marking, including long-term assignment.
845    ///
846    /// Takes `reference` BY VALUE and MOVES it into the DPB (the caller's local is
847    /// dropped right after) — the old `&mut` + `insert(0, reference.clone())` cloned
848    /// all three planes (~1.35 MB/frame) a second time, on top of `as_reference`'s
849    /// necessary clone. Returns the picture's final `frame_num` (0 after MMCO 5) for
850    /// the caller's gap-detection tracking, since `reference` is gone after the move.
851    fn apply_ref_marking(
852        &mut self,
853        mut reference: RefFrame,
854        ops: &[Mmco],
855        frame_num: u32,
856        log2_max_frame_num: u32,
857        max_refs: usize,
858    ) -> u32 {
859        let max = 1i64 << log2_max_frame_num;
860        let curr = frame_num as i64;
861        let pic_num = |rf: &RefFrame| -> i64 {
862            if (rf.frame_num as i64) > curr {
863                rf.frame_num as i64 - max
864            } else {
865                rf.frame_num as i64
866            }
867        };
868
869        if ops.is_empty() {
870            // Sliding window: insert the current (short-term) picture, then evict
871            // the oldest short-term reference while over capacity (long-term refs
872            // are retained).
873            let out_fn = reference.frame_num;
874            self.refs.insert(0, std::sync::Arc::new(reference));
875            while self.refs.len() > max_refs {
876                match self.refs.iter().rposition(|r| !r.long_term) {
877                    Some(pos) => {
878                        // Park the evicted frame; its planes are reclaimed on the
879                        // next picture boundary (see `reclaim_retired`).
880                        let evicted = self.refs.remove(pos);
881                        self.retired.push(evicted);
882                    }
883                    None => break,
884                }
885            }
886            return out_fn;
887        }
888
889        // Adaptive marking (MMCO), applied in order.
890        for &op in ops {
891            match op {
892                Mmco::Unref(diff) => {
893                    let target = curr - (diff as i64 + 1);
894                    self.refs.retain(|r| r.long_term || pic_num(r) != target);
895                }
896                Mmco::UnrefLong(ltpn) => {
897                    self.refs.retain(|r| !(r.long_term && r.long_term_idx == ltpn));
898                }
899                Mmco::AssignLong(diff, idx) => {
900                    let target = curr - (diff as i64 + 1);
901                    self.refs.retain(|r| !(r.long_term && r.long_term_idx == idx));
902                    for r in self.refs.iter_mut() {
903                        if !r.long_term && pic_num(r) == target {
904                            // Rare op; make_mut only copies if a slice still holds it.
905                            let r = std::sync::Arc::make_mut(r);
906                            r.long_term = true;
907                            r.long_term_idx = idx;
908                        }
909                    }
910                }
911                Mmco::MaxLong(max_plus1) => {
912                    self.refs.retain(|r| !(r.long_term && r.long_term_idx + 1 > max_plus1));
913                }
914                Mmco::Reset => {
915                    self.refs.clear();
916                    reference.frame_num = 0;
917                }
918                Mmco::CurrentLong(idx) => {
919                    self.refs.retain(|r| !(r.long_term && r.long_term_idx == idx));
920                    reference.long_term = true;
921                    reference.long_term_idx = idx;
922                }
923            }
924        }
925        let out_fn = reference.frame_num;
926        self.refs.insert(0, std::sync::Arc::new(reference));
927        // Safety net so a malformed marking stream can't grow the DPB unbounded.
928        let cap = max_refs.max(16);
929        if self.refs.len() > cap {
930            self.refs.truncate(cap);
931        }
932        out_fn
933    }
934}
935
936/// Emits a GOP's buffered pictures in display order (sorted by `PicOrderCnt`).
937fn flush_gop(gop: &mut Vec<(i32, YuvFrame)>, out: &mut Vec<YuvFrame>) {
938    gop.sort_by_key(|(poc, _)| *poc);
939    out.extend(gop.drain(..).map(|(_, f)| f));
940}
941
942/// Whether an access unit contains an IDR coded-slice NAL.
943fn au_is_idr(au: &[u8]) -> bool {
944    split_annex_b(au)
945        .iter()
946        .any(|n| !n.is_empty() && NalUnitType::from_id(n[0]) == NalUnitType::IdrSlice)
947}
948
949/// Splits an Annex-B byte stream into access units, each ending after a VCL
950/// (coded-slice) NAL with any preceding parameter-set/SEI NALs attached. Start
951/// codes are preserved so each unit can be passed straight to [`Decoder::decode`].
952///
953/// Public because [`Decoder::decode`] takes ONE access unit: a caller that wants
954/// decode-order pictures, or wants to drop each picture as it arrives instead of
955/// accumulating the stream like [`Decoder::decode_stream`], needs this to feed it.
956pub fn split_access_units(stream: &[u8]) -> Vec<&[u8]> {
957    // (offset of the start code, whether the NAL it begins is a VCL slice).
958    let mut codes: Vec<(usize, bool)> = Vec::new();
959    let mut i = 0;
960    while i + 3 <= stream.len() {
961        if stream[i] == 0 && stream[i + 1] == 0 && stream[i + 2] == 1 {
962            let nal_type = NalUnitType::from_id(stream.get(i + 3).copied().unwrap_or(0));
963            let is_vcl = matches!(nal_type, NalUnitType::IdrSlice | NalUnitType::NonIdrSlice);
964            // Include a leading zero (4-byte start code) in the unit boundary.
965            let sc = if i > 0 && stream[i - 1] == 0 { i - 1 } else { i };
966            codes.push((sc, is_vcl));
967            i += 3;
968        } else {
969            i += 1;
970        }
971    }
972    if codes.is_empty() {
973        return vec![stream];
974    }
975    let mut aus = Vec::new();
976    let mut start = codes[0].0;
977    for k in 0..codes.len() {
978        if codes[k].1 {
979            let end = codes.get(k + 1).map_or(stream.len(), |c| c.0);
980            aus.push(&stream[start..end]);
981            start = end;
982        }
983    }
984    aus
985}
986
987/// Parses a `pred_weight_table()` (spec §7.3.3.2) for the active reference lists
988/// (4:2:0 → chroma weights always present). List 1 is parsed only for B slices.
989fn parse_pred_weight_table(
990    r: &mut BitReader,
991    num_l0: usize,
992    num_l1: usize,
993    is_b: bool,
994) -> Result<WeightTable, DecodeError> {
995    let luma_log2_denom = r.read_ue()? as i32;
996    let chroma_log2_denom = r.read_ue()? as i32;
997    // Spec §7.4.3.2 constrains both weight denoms to [0, 7]; a malformed stream can
998    // carry any ue(v). Reject before `1 << denom` (which overflows for denom ≥ 31)
999    // so a corrupt bitstream is rejected gracefully, never panics.
1000    if !(0..=7).contains(&luma_log2_denom) || !(0..=7).contains(&chroma_log2_denom) {
1001        return Err(DecodeError::Unsupported("invalid weight denom"));
1002    }
1003    let mut wt = WeightTable {
1004        luma_log2_denom,
1005        chroma_log2_denom,
1006        ..Default::default()
1007    };
1008    let lists: &[(usize, usize)] = if is_b {
1009        &[(0, num_l0), (1, num_l1)]
1010    } else {
1011        &[(0, num_l0)]
1012    };
1013    for &(list, n) in lists {
1014        let mut luma = Vec::with_capacity(n);
1015        let mut chroma = Vec::with_capacity(n);
1016        for _ in 0..n {
1017            let (mut lw, mut lo) = (1 << luma_log2_denom, 0);
1018            if r.read_bit()? {
1019                lw = r.read_se()?;
1020                lo = r.read_se()?;
1021            }
1022            luma.push((lw, lo));
1023            let mut ch = [(1 << chroma_log2_denom, 0); 2];
1024            if r.read_bit()? {
1025                for slot in ch.iter_mut() {
1026                    *slot = (r.read_se()?, r.read_se()?);
1027                }
1028            }
1029            chroma.push(ch);
1030        }
1031        wt.luma[list] = luma;
1032        wt.chroma[list] = chroma;
1033    }
1034    Ok(wt)
1035}
1036
1037/// Resolves the effective scaling matrices for a slice from the SPS lists and
1038/// any PPS override (fall-back rule B), returning them un-zig-zagged to raster
1039/// order: six 4×4 lists [Y/Cb/Cr intra, Y/Cb/Cr inter] and two 8×8 luma lists
1040/// [Y-intra, Y-inter].
1041fn resolve_scaling(sps: &Sps, pps: &Pps) -> ([[i32; 16]; 6], [[i32; 64]; 2]) {
1042    use crate::params::{
1043        DEFAULT_4X4_INTER, DEFAULT_4X4_INTRA, DEFAULT_8X8_INTER, DEFAULT_8X8_INTRA,
1044    };
1045    const ZZ4: [usize; 16] = [0, 1, 4, 8, 5, 2, 3, 6, 9, 12, 13, 10, 7, 11, 14, 15];
1046    // 8×8 frame zig-zag scan → raster index (spec Table 8-12).
1047    const ZZ8: [usize; 64] = [
1048        0, 1, 8, 16, 9, 2, 3, 10, 17, 24, 32, 25, 18, 11, 4, 5, 12, 19, 26, 33, 40, 48, 41, 34, 27,
1049        20, 13, 6, 7, 14, 21, 28, 35, 42, 49, 56, 57, 50, 43, 36, 29, 22, 15, 23, 30, 37, 44, 51,
1050        58, 59, 52, 45, 38, 31, 39, 46, 53, 60, 61, 54, 47, 55, 62, 63,
1051    ];
1052    // Effective zig-zag lists: a PPS override (rule B) takes precedence; an absent
1053    // PPS list falls back to the SPS list (or the default / previous PPS list).
1054    let mut z4 = [[16u8; 16]; 6];
1055    for i in 0..6 {
1056        z4[i] = if pps.pic_scaling_matrix_present {
1057            if pps.scaling_present_4x4[i] {
1058                pps.scaling_4x4[i]
1059            } else {
1060                match i {
1061                    0 if sps.has_scaling => sps.scaling_4x4[0],
1062                    0 => DEFAULT_4X4_INTRA,
1063                    3 if sps.has_scaling => sps.scaling_4x4[3],
1064                    3 => DEFAULT_4X4_INTER,
1065                    _ => z4[i - 1],
1066                }
1067            }
1068        } else {
1069            sps.scaling_4x4[i]
1070        };
1071    }
1072    let mut z8 = [[16u8; 64]; 2];
1073    for (i, list) in z8.iter_mut().enumerate() {
1074        *list = if pps.pic_scaling_matrix_present {
1075            if pps.scaling_present_8x8[i] {
1076                pps.scaling_8x8[i]
1077            } else if sps.has_scaling {
1078                sps.scaling_8x8[i]
1079            } else if i == 0 {
1080                DEFAULT_8X8_INTRA
1081            } else {
1082                DEFAULT_8X8_INTER
1083            }
1084        } else {
1085            sps.scaling_8x8[i]
1086        };
1087    }
1088    let mut out4 = [[16i32; 16]; 6];
1089    for (li, list) in out4.iter_mut().enumerate() {
1090        for k in 0..16 {
1091            list[ZZ4[k]] = z4[li][k] as i32;
1092        }
1093    }
1094    let mut out8 = [[16i32; 64]; 2];
1095    for (li, list) in out8.iter_mut().enumerate() {
1096        for k in 0..64 {
1097            list[ZZ8[k]] = z8[li][k] as i32;
1098        }
1099    }
1100    (out4, out8)
1101}
1102
1103/// Parses a `ref_pic_list_modification` command list (spec §7.3.3.1) into
1104/// `(modification_of_pic_nums_idc, value)` pairs, stopping at idc 3.
1105fn parse_ref_pic_list_modification(
1106    r: &mut BitReader,
1107    out: &mut Vec<(u32, u32)>,
1108) -> Result<(), DecodeError> {
1109    loop {
1110        let idc = r.read_ue()?;
1111        if idc == 3 {
1112            break;
1113        }
1114        if idc > 3 {
1115            return Err(DecodeError::Unsupported("invalid ref_pic_list_modification"));
1116        }
1117        let val = r.read_ue()?; // abs_diff_pic_num_minus1 / long_term_pic_num
1118        out.push((idc, val));
1119        if out.len() > 64 {
1120            return Err(DecodeError::Truncated); // runaway / corrupt
1121        }
1122    }
1123    Ok(())
1124}
1125
1126/// Builds the P-slice `RefPicList0`: short-term references ordered by descending
1127/// `FrameNumWrap`, then long-term by ascending idx (spec §8.2.4.2.1), with any
1128/// `ref_pic_list_modification` applied.
1129fn build_ref_list_p(
1130    dpb: &[Ref],
1131    curr_frame_num: u32,
1132    max_frame_num: u32,
1133    num_active: usize,
1134    mods: &[(u32, u32)],
1135) -> Result<Vec<Ref>, DecodeError> {
1136    let curr = curr_frame_num as i64;
1137    let max = max_frame_num as i64;
1138    let pic_num = |fnum: u32| -> i64 {
1139        let f = fnum as i64;
1140        if f > curr { f - max } else { f }
1141    };
1142    let mut init: Vec<Ref> = dpb.iter().filter(|r| !r.long_term).cloned().collect();
1143    init.sort_by_key(|rf| core::cmp::Reverse(pic_num(rf.frame_num)));
1144    let mut long: Vec<Ref> = dpb.iter().filter(|r| r.long_term).cloned().collect();
1145    long.sort_by_key(|rf| rf.long_term_idx);
1146    init.extend(long);
1147    apply_list_modification(init, curr_frame_num, max_frame_num, num_active, mods)
1148}
1149
1150/// Builds the B-slice `RefPicList0` and `RefPicList1` (spec §8.2.4.2.3), ordered
1151/// by `PicOrderCnt` relative to the current picture: List0 leads with nearer
1152/// past pictures, List1 with nearer future pictures. Long-term references follow.
1153/// Per-list `ref_pic_list_modification` is then applied.
1154#[allow(clippy::too_many_arguments)]
1155fn build_ref_list_b(
1156    dpb: &[Ref],
1157    curr_poc: i32,
1158    curr_frame_num: u32,
1159    max_frame_num: u32,
1160    num0: usize,
1161    num1: usize,
1162    mods0: &[(u32, u32)],
1163    mods1: &[(u32, u32)],
1164) -> Result<(Vec<Ref>, Vec<Ref>), DecodeError> {
1165    let mut less: Vec<Ref> =
1166        dpb.iter().filter(|r| !r.long_term && r.poc < curr_poc).cloned().collect();
1167    let mut greater: Vec<Ref> =
1168        dpb.iter().filter(|r| !r.long_term && r.poc > curr_poc).cloned().collect();
1169    let mut long: Vec<Ref> = dpb.iter().filter(|r| r.long_term).cloned().collect();
1170    less.sort_by_key(|r| core::cmp::Reverse(r.poc)); // nearest past first
1171    greater.sort_by_key(|r| r.poc); // nearest future first
1172    long.sort_by_key(|r| r.long_term_idx);
1173
1174    let mut init0 = less.clone();
1175    init0.extend(greater.clone());
1176    init0.extend(long.clone());
1177    let mut init1 = greater;
1178    init1.extend(less);
1179    init1.extend(long);
1180
1181    // When List1 (truncated to its active length) equals List0 and has more than
1182    // one entry, swap its first two entries (spec §8.2.4.2.3).
1183    let eq_len = num0.min(num1).min(init0.len()).min(init1.len());
1184    if num1 > 1
1185        && init1.len() > 1
1186        && (0..eq_len).all(|i| same_picture(&init0[i], &init1[i]))
1187        && eq_len == num1.min(init1.len())
1188        && eq_len == num0.min(init0.len())
1189    {
1190        init1.swap(0, 1);
1191    }
1192
1193    let list0 = apply_list_modification(init0, curr_frame_num, max_frame_num, num0, mods0)?;
1194    let list1 = apply_list_modification(init1, curr_frame_num, max_frame_num, num1, mods1)?;
1195    Ok((list0, list1))
1196}
1197
1198/// Two DPB entries refer to the same picture (used for the List1 swap rule).
1199fn same_picture(a: &RefFrame, b: &RefFrame) -> bool {
1200    a.long_term == b.long_term
1201        && if a.long_term { a.long_term_idx == b.long_term_idx } else { a.poc == b.poc }
1202}
1203
1204/// Applies `ref_pic_list_modification` to an initialized reference list and
1205/// truncates it to `num_active` (spec §8.2.4.3). `init` is the full ordered list;
1206/// the result is `num_active` entries, possibly reordered. idc 0/1 reference
1207/// short-term pictures by PicNum, idc 2 long-term ones by LongTermFrameIdx.
1208fn apply_list_modification(
1209    init: Vec<Ref>,
1210    curr_frame_num: u32,
1211    max_frame_num: u32,
1212    num_active: usize,
1213    mods: &[(u32, u32)],
1214) -> Result<Vec<Ref>, DecodeError> {
1215    if mods.is_empty() {
1216        let mut init = init;
1217        init.truncate(num_active.max(1));
1218        return Ok(init);
1219    }
1220    let curr = curr_frame_num as i64;
1221    let max = max_frame_num as i64;
1222    let mut list = init.clone();
1223    let mut pic_num_pred = curr;
1224    let mut refidx = 0usize;
1225    for &(idc, val) in mods {
1226        let matches: Box<dyn Fn(&RefFrame) -> bool> = if idc == 2 {
1227            Box::new(move |r: &RefFrame| r.long_term && r.long_term_idx == val)
1228        } else {
1229            let abs_diff = (val as i64) + 1;
1230            let no_wrap = if idc == 0 {
1231                let x = pic_num_pred - abs_diff;
1232                if x < 0 { x + max } else { x }
1233            } else {
1234                let x = pic_num_pred + abs_diff;
1235                if x >= max { x - max } else { x }
1236            };
1237            pic_num_pred = no_wrap;
1238            let target = if no_wrap > curr { no_wrap - max } else { no_wrap };
1239            Box::new(move |r: &RefFrame| {
1240                let pn = if r.frame_num as i64 > curr {
1241                    r.frame_num as i64 - max
1242                } else {
1243                    r.frame_num as i64
1244                };
1245                !r.long_term && pn == target
1246            })
1247        };
1248        let found = init.iter().find(|r| matches(r)).cloned();
1249        let Some(found) = found else {
1250            if std::env::var_os("RH264_DUMP_MB").is_some() {
1251                let cand: Vec<String> = init
1252                    .iter()
1253                    .map(|r| {
1254                        let pn = if r.frame_num as i64 > curr {
1255                            r.frame_num as i64 - max
1256                        } else {
1257                            r.frame_num as i64
1258                        };
1259                        format!("(fn={} poc={} lt={} picnum={})", r.frame_num, r.poc, r.long_term, pn)
1260                    })
1261                    .collect();
1262                eprintln!(
1263                    "MODFAIL idc={idc} val={val}  curr_frame_num={curr} max={max}  init={}",
1264                    cand.join(" ")
1265                );
1266            }
1267            return Err(DecodeError::Truncated); // references a picture not in the DPB
1268        };
1269        if refidx > list.len() {
1270            break;
1271        }
1272        list.insert(refidx, found);
1273        if let Some(dup) = list.iter().enumerate().skip(refidx + 1).find(|(_, r)| matches(r)).map(|(i, _)| i) {
1274            list.remove(dup);
1275        }
1276        refidx += 1;
1277        if refidx >= num_active {
1278            break;
1279        }
1280    }
1281    list.truncate(num_active.max(1));
1282    Ok(list)
1283}
1284
1285#[cfg(test)]
1286mod tests {
1287    use super::*;
1288
1289    fn ref_at(poc: i32, fnum: u32) -> Ref {
1290        std::sync::Arc::new(RefFrame {
1291            py: vec![],
1292            pu: vec![],
1293            pv: vec![],
1294            cw: 0,
1295            ch: 0,
1296            frame_num: fnum,
1297            poc,
1298            mv: Vec::new(),
1299            ref_idx: Vec::new(),
1300            mv1: Vec::new(),
1301            ref_idx1: Vec::new(),
1302            ref_poc: Vec::new(),
1303            w4: 0,
1304            long_term: false,
1305            long_term_idx: 0,
1306        })
1307    }
1308
1309    #[test]
1310    fn b_ref_lists_ordered_by_poc() {
1311        // Current POC 4; DPB has past (0,2) and future (6,8) references.
1312        let dpb = vec![ref_at(8, 4), ref_at(6, 3), ref_at(2, 1), ref_at(0, 0)];
1313        let (l0, l1) = build_ref_list_b(&dpb, 4, 5, 16, 4, 4, &[], &[]).unwrap();
1314        // List0: nearer past first (desc), then nearer future (asc).
1315        assert_eq!(l0.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![2, 0, 6, 8]);
1316        // List1: nearer future first (asc), then nearer past (desc).
1317        assert_eq!(l1.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![6, 8, 2, 0]);
1318    }
1319
1320    #[test]
1321    fn b_ref_list1_swap_when_equal() {
1322        // Only past references -> List0 and List1 initialize identically, so
1323        // List1's first two entries are swapped (spec §8.2.4.2.3).
1324        let dpb = vec![ref_at(4, 2), ref_at(2, 1), ref_at(0, 0)];
1325        let (l0, l1) = build_ref_list_b(&dpb, 6, 3, 16, 3, 3, &[], &[]).unwrap();
1326        assert_eq!(l0.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![4, 2, 0]);
1327        assert_eq!(l1.iter().map(|r| r.poc).collect::<Vec<_>>(), vec![2, 4, 0]);
1328    }
1329
1330    #[test]
1331    fn frame_num_gaps_insert_placeholders() {
1332        let mut d = Decoder::new();
1333        d.prev_ref_frame_num = 2;
1334        // frame_num jumps 2 -> 5: placeholders for the skipped 3 and 4.
1335        d.insert_frame_num_gaps(5, 16, 8, 16, 16);
1336        let fns: Vec<u32> = d.refs.iter().map(|r| r.frame_num).collect();
1337        assert_eq!(fns, vec![4, 3], "most-recent placeholder at the front");
1338        assert_eq!(d.prev_ref_frame_num, 4);
1339        assert!(d.refs.iter().all(|r| r.py.iter().all(|&p| p == 128)), "grey fill");
1340    }
1341
1342    #[test]
1343    fn frame_num_gaps_wrap_and_noop() {
1344        // Wrap across MaxFrameNum: prev 14, frame_num 1 (max 16) -> fill 15, 0.
1345        let mut d = Decoder::new();
1346        d.prev_ref_frame_num = 14;
1347        d.insert_frame_num_gaps(1, 16, 8, 16, 16);
1348        assert_eq!(d.refs.iter().map(|r| r.frame_num).collect::<Vec<_>>(), vec![0, 15]);
1349        // No gap (consecutive) inserts nothing.
1350        let mut d = Decoder::new();
1351        d.prev_ref_frame_num = 3;
1352        d.insert_frame_num_gaps(4, 16, 8, 16, 16);
1353        assert!(d.refs.is_empty());
1354    }
1355
1356    #[test]
1357    fn missing_param_sets_errors() {
1358        let mut d = Decoder::new();
1359        // A lone (fake) IDR slice header: first_mb_in_slice=0, slice_type=7 (I),
1360        // pic_parameter_set_id=0 — then the PPS lookup fails (none stored).
1361        let nal = rusty_h264_common::NalUnit::new(3, NalUnitType::IdrSlice, vec![0x88, 0x80]);
1362        let err = d.decode(&nal.to_annex_b()).unwrap_err();
1363        assert_eq!(err, DecodeError::MissingParameterSet);
1364    }
1365}