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