Skip to main content

oxideav_h265/
sequence.rs

1//! Whole-bitstream decode driver — Annex B byte stream to output
2//! pictures.
3//!
4//! This is the §8.1 general decoding process expressed over the crate's
5//! subsystems: the [`crate::nal`] Annex B demux feeds parameter-set
6//! activation ([`crate::sps`] / [`crate::pps`]), each coded picture's
7//! slice segments are parsed ([`crate::slice`]) and their
8//! `slice_segment_data()` CABAC-decoded through the §7.3.8 syntax walk
9//! ([`crate::slice_data`]), and the decoded coding tree units are handed
10//! to the picture-level reconstruction + in-loop-filter driver
11//! ([`crate::inter_recon::reconstruct_inter_picture`]) with the §8.3
12//! POC / RPS / reference-list cycle threaded by
13//! [`crate::decode::PictureSequenceState`]. Decoded pictures are
14//! returned in output order (§8.3.1 `PicOrderCntVal` order within each
15//! coded video sequence).
16
17use std::collections::BTreeMap;
18
19use crate::availability::{PictureTiling, TilingParams};
20use crate::bitreader::BitReader;
21use crate::cabac::{init_type, CabacEngine};
22use crate::ctx_init::SliceContexts;
23use crate::decode::{PictureHeaderInfo, PictureSequenceState, SliceRefParams};
24use crate::dpb::{LongTermEntry, RefPicLists};
25use crate::inter_pred::WpListWeights;
26use crate::inter_recon::{
27    reconstruct_inter_picture, InterSliceContext, PlacedInterCtu, RefListAccess, SliceWpTables,
28};
29use crate::nal::{NalError, NalIter, NalUnit};
30use crate::picture::Picture;
31use crate::poc::NalKind;
32use crate::pps::{PicParameterSet, PpsError};
33use crate::recon::{ReconError, ReconParams};
34use crate::residual::ResidualCodingError;
35use crate::slice::{SliceError, SliceLongTermRefPicSource, SliceSegmentHeader, SliceType};
36use crate::slice_data::{
37    decode_coding_tree_unit_in_picture, end_of_slice_segment_flag, CodingTreeUnit,
38    PictureParseState, SliceDataParams,
39};
40use crate::sps::{
41    MaterializedShortTermRefPicSet, SeqParameterSet, ShortTermRefPicSetMaterializeError, SpsError,
42};
43
44/// NAL unit type: video parameter set (Table 7-1).
45const NAL_VPS: u8 = 32;
46/// NAL unit type: sequence parameter set.
47const NAL_SPS: u8 = 33;
48/// NAL unit type: picture parameter set.
49const NAL_PPS: u8 = 34;
50
51/// Errors from the whole-bitstream decode driver.
52#[derive(Debug)]
53pub enum SequenceError {
54    /// Annex B demux / NAL header error.
55    Nal(NalError),
56    /// SPS parse error.
57    Sps(SpsError),
58    /// PPS parse error.
59    Pps(PpsError),
60    /// Slice-segment-header parse error.
61    Slice(SliceError),
62    /// §7.3.8 slice-data CABAC walk error.
63    SliceData(ResidualCodingError),
64    /// Picture reconstruction error.
65    Recon(ReconError),
66    /// §7.4.8 short-term-RPS materialization error.
67    Rps(ShortTermRefPicSetMaterializeError),
68    /// A referenced parameter set was never activated.
69    MissingParameterSet {
70        /// `"sps"` or `"pps"`.
71        kind: &'static str,
72        /// The referenced parameter-set id.
73        id: u8,
74    },
75    /// A structural bitstream-conformance failure.
76    Malformed(&'static str),
77    /// A conformant configuration this driver does not decode yet.
78    Unsupported(&'static str),
79}
80
81impl core::fmt::Display for SequenceError {
82    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
83        match self {
84            Self::Nal(e) => write!(f, "NAL demux error: {e}"),
85            Self::Sps(e) => write!(f, "SPS parse error: {e}"),
86            Self::Pps(e) => write!(f, "PPS parse error: {e}"),
87            Self::Slice(e) => write!(f, "slice header parse error: {e}"),
88            Self::SliceData(e) => write!(f, "slice data decode error: {e}"),
89            Self::Recon(e) => write!(f, "picture reconstruction error: {e}"),
90            Self::Rps(e) => write!(f, "short-term RPS materialization error: {e:?}"),
91            Self::MissingParameterSet { kind, id } => {
92                write!(f, "referenced {kind} id {id} was never activated")
93            }
94            Self::Malformed(what) => write!(f, "malformed bitstream: {what}"),
95            Self::Unsupported(what) => write!(f, "unsupported configuration: {what}"),
96        }
97    }
98}
99
100impl std::error::Error for SequenceError {}
101
102impl From<NalError> for SequenceError {
103    fn from(e: NalError) -> Self {
104        Self::Nal(e)
105    }
106}
107impl From<SpsError> for SequenceError {
108    fn from(e: SpsError) -> Self {
109        Self::Sps(e)
110    }
111}
112impl From<PpsError> for SequenceError {
113    fn from(e: PpsError) -> Self {
114        Self::Pps(e)
115    }
116}
117impl From<SliceError> for SequenceError {
118    fn from(e: SliceError) -> Self {
119        Self::Slice(e)
120    }
121}
122impl From<ResidualCodingError> for SequenceError {
123    fn from(e: ResidualCodingError) -> Self {
124        Self::SliceData(e)
125    }
126}
127impl From<ReconError> for SequenceError {
128    fn from(e: ReconError) -> Self {
129        Self::Recon(e)
130    }
131}
132impl From<ShortTermRefPicSetMaterializeError> for SequenceError {
133    fn from(e: ShortTermRefPicSetMaterializeError) -> Self {
134        Self::Rps(e)
135    }
136}
137
138/// One decoded picture with its output-ordering keys.
139#[derive(Debug, Clone)]
140pub struct DecodedFrame {
141    /// Index of the coded video sequence this picture belongs to
142    /// (incremented at each IRAP with `NoRaslOutputFlag == 1`).
143    pub cvs_index: u32,
144    /// `PicOrderCntVal` (§8.3.1).
145    pub poc: i32,
146    /// `pic_output_flag` — `false` pictures are decoded (they may be
147    /// referenced) but not output.
148    pub output: bool,
149    /// The reconstructed, in-loop-filtered picture.
150    pub picture: Picture,
151}
152
153/// One slice segment of the picture being assembled.
154#[derive(Debug)]
155struct SegmentData {
156    nal_type: u8,
157    temporal_id: u8,
158    layer_id: u8,
159    rbsp: Vec<u8>,
160    /// Coded (escaped) payload — the §7.4.7.1 entry-point offsets are
161    /// expressed in this byte space.
162    escaped: Vec<u8>,
163    header: SliceSegmentHeader,
164}
165
166/// The whole-bitstream decoder: parameter-set activation + per-picture
167/// slice-data decode + the §8.3 reference cycle.
168#[derive(Debug, Default)]
169pub struct SequenceDecoder {
170    sps: BTreeMap<u8, SeqParameterSet>,
171    pps: BTreeMap<u8, PicParameterSet>,
172    state: PictureSequenceState,
173    frames: Vec<DecodedFrame>,
174    pending: Vec<SegmentData>,
175    cvs_index: u32,
176    seen_picture: bool,
177    /// Debug: tolerate an end_of_slice_segment_flag mismatch (decode
178    /// as much as possible instead of erroring).
179    tolerant: bool,
180}
181
182impl SequenceDecoder {
183    /// A fresh decoder with no activated parameter sets.
184    #[must_use]
185    pub fn new() -> Self {
186        Self::default()
187    }
188
189    /// Debug: keep decoding past an `end_of_slice_segment_flag`
190    /// mismatch. Not part of the stable API.
191    #[doc(hidden)]
192    pub fn set_tolerant(&mut self, tolerant: bool) {
193        self.tolerant = tolerant;
194    }
195
196    /// Feed a whole Annex B byte stream, decoding every access unit.
197    ///
198    /// # Errors
199    /// Any demux / parse / decode error; the decoder state is
200    /// unspecified after an error.
201    pub fn push_annexb(&mut self, data: &[u8]) -> Result<(), SequenceError> {
202        for unit in NalIter::new(data) {
203            self.push_nal_unit(unit?)?;
204        }
205        Ok(())
206    }
207
208    /// Feed one demuxed NAL unit.
209    ///
210    /// # Errors
211    /// Any parse / decode error.
212    pub fn push_nal_unit(&mut self, unit: NalUnit) -> Result<(), SequenceError> {
213        let NalUnit {
214            header,
215            rbsp,
216            escaped,
217        } = unit;
218        if header.is_vcl() {
219            // §7.4.2.4.4: a VCL NAL with first_slice_segment_in_pic_flag
220            // set starts a new access unit.
221            let first_in_pic = rbsp.first().is_some_and(|b| b & 0x80 != 0);
222            if first_in_pic {
223                self.finish_picture()?;
224            }
225            let pps_id = peek_slice_pps_id(&rbsp, header.nal_unit_type)?;
226            let pps = self
227                .pps
228                .get(&pps_id)
229                .ok_or(SequenceError::MissingParameterSet {
230                    kind: "pps",
231                    id: pps_id,
232                })?;
233            let sps = self
234                .sps
235                .get(&pps.sps_id)
236                .ok_or(SequenceError::MissingParameterSet {
237                    kind: "sps",
238                    id: pps.sps_id,
239                })?;
240            let parsed = SliceSegmentHeader::parse(&rbsp, header.nal_unit_type, sps, pps)?;
241            self.pending.push(SegmentData {
242                nal_type: header.nal_unit_type,
243                temporal_id: header.temporal_id,
244                layer_id: header.nuh_layer_id,
245                rbsp,
246                escaped,
247                header: parsed,
248            });
249            return Ok(());
250        }
251        match header.nal_unit_type {
252            NAL_VPS => {
253                // §7.4.2.4.4: a VPS / SPS / PPS NAL unit (nuh_layer_id
254                // 0) succeeding a VCL NAL unit starts a NEW access
255                // unit — the pending picture is complete. Decode it
256                // BEFORE the arriving parameter set can overwrite the
257                // sets it was coded against (streams legally re-send
258                // a parameter set with the same id and new content for
259                // the next CVS / picture).
260                if header.nuh_layer_id == 0 {
261                    self.finish_picture()?;
262                }
263                // The VPS carries no fields the single-layer decode
264                // needs; activation is otherwise a no-op.
265            }
266            NAL_SPS => {
267                if header.nuh_layer_id == 0 {
268                    self.finish_picture()?;
269                }
270                let sps = SeqParameterSet::parse(&rbsp)?;
271                self.sps.insert(sps.sps_id, sps);
272            }
273            NAL_PPS => {
274                if header.nuh_layer_id == 0 {
275                    self.finish_picture()?;
276                }
277                let pps = PicParameterSet::parse(&rbsp)?;
278                self.pps.insert(pps.pps_id, pps);
279            }
280            // AUD / EOS / EOB / FD / SEI: nothing to activate.
281            _ => {}
282        }
283        Ok(())
284    }
285
286    /// Decode any picture still being assembled (a flush point — call
287    /// when the input stream ends but the decoder object lives on).
288    ///
289    /// # Errors
290    /// Any decode error from the pending picture.
291    pub fn flush(&mut self) -> Result<(), SequenceError> {
292        self.finish_picture()
293    }
294
295    /// Drain the pictures decoded so far, in decode order. The caller
296    /// owns output reordering (the streaming [`crate::decoder`] holds a
297    /// `sps_max_num_reorder_pics`-deep queue; [`Self::finish`] sorts a
298    /// whole sequence at once).
299    pub fn take_decoded(&mut self) -> Vec<DecodedFrame> {
300        std::mem::take(&mut self.frames)
301    }
302
303    /// `sps_max_num_reorder_pics` of the highest sub-layer of the most
304    /// recently activated SPS (`None` before any SPS).
305    #[must_use]
306    pub fn max_num_reorder_pics(&self) -> Option<u32> {
307        self.sps.values().next_back().map(|sps| {
308            let idx =
309                usize::from(sps.max_sub_layers_minus1).min(sps.sub_layer_ordering_info.len() - 1);
310            sps.sub_layer_ordering_info[idx].max_num_reorder_pics
311        })
312    }
313
314    /// Decode any picture still being assembled and return every
315    /// decoded frame in output order.
316    ///
317    /// # Errors
318    /// Any decode error from the final pending picture.
319    pub fn finish(mut self) -> Result<Vec<DecodedFrame>, SequenceError> {
320        self.finish_picture()?;
321        // §C.5.2.2 output order: `PicOrderCntVal` order within each
322        // coded video sequence, sequences in decode order.
323        let mut frames = self.frames;
324        frames.sort_by_key(|f| (f.cvs_index, f.poc));
325        Ok(frames)
326    }
327
328    /// Decode the pending picture's slice segments into a picture.
329    fn finish_picture(&mut self) -> Result<(), SequenceError> {
330        if self.pending.is_empty() {
331            return Ok(());
332        }
333        let segs = std::mem::take(&mut self.pending);
334        self.decode_picture(&segs)
335    }
336
337    fn decode_picture(&mut self, segs: &[SegmentData]) -> Result<(), SequenceError> {
338        let indep = &segs[0];
339        if indep.header.dependent_slice_segment_flag {
340            return Err(SequenceError::Malformed(
341                "first slice segment of a picture is dependent",
342            ));
343        }
344        let pps = self
345            .pps
346            .get(&indep.header.slice_pic_parameter_set_id)
347            .ok_or(SequenceError::MissingParameterSet {
348                kind: "pps",
349                id: indep.header.slice_pic_parameter_set_id,
350            })?;
351        let sps = self
352            .sps
353            .get(&pps.sps_id)
354            .ok_or(SequenceError::MissingParameterSet {
355                kind: "sps",
356                id: pps.sps_id,
357            })?;
358
359        let geom = Geometry::derive(sps, pps)?;
360
361        // §7.4.2.4.4 CVS bookkeeping: an IRAP with NoRaslOutputFlag
362        // starts a new coded video sequence (for output ordering).
363        let nal_kind = NalKind::new(indep.nal_type);
364        let no_rasl_output =
365            nal_kind.is_idr() || nal_kind.is_bla() || (nal_kind.is_irap() && !self.seen_picture);
366        if nal_kind.is_irap() && no_rasl_output && self.seen_picture {
367            self.cvs_index += 1;
368        }
369        self.seen_picture = true;
370
371        // ---- §7.3.8 slice-data CABAC decode of every slice segment ----
372        let pic_size_in_ctbs = (geom.pic_w_ctbs * geom.pic_h_ctbs) as usize;
373        let mut decoded: Vec<(u32, u32, CodingTreeUnit)> = Vec::new();
374        let mut slice_addr_of: Vec<Option<u32>> = vec![None; pic_size_in_ctbs];
375        let first_slice_type = segs[0]
376            .header
377            .slice_type
378            .ok_or(SequenceError::Malformed("independent slice without type"))?;
379        let mut parse_state = PictureParseState::new(&build_slice_data_params(
380            &segs[0].header,
381            sps,
382            pps,
383            &geom,
384            first_slice_type,
385        ));
386
387        // §7.4.7.1 — a dependent slice segment inherits the slice-level
388        // header values (and SliceAddrRs) from the preceding independent
389        // slice segment; §9.3.2.2 restores its CABAC context variables
390        // from the state stored at the end of the previous segment
391        // (TableStateIdxDs, §9.3.2.4).
392        let mut cur_indep: &SegmentData = indep;
393        let mut ds_stored: Option<SliceContexts> = None;
394        // §9.3.2.4 WPP snapshot — ONE picture-wide storage: a CTU row
395        // started by a later slice segment of the same slice
396        // synchronizes from the state stored while an earlier segment
397        // decoded the row above (§9.3.2.5, T-availability gated).
398        let mut wpp_stored: Option<SliceContexts> = None;
399        for seg in segs {
400            if seg.header.dependent_slice_segment_flag {
401                if ds_stored.is_none() {
402                    return Err(SequenceError::Malformed(
403                        "dependent slice segment without a preceding segment's context state",
404                    ));
405                }
406            } else {
407                cur_indep = seg;
408            }
409            decode_slice_segment_data(
410                seg,
411                &cur_indep.header,
412                sps,
413                pps,
414                &geom,
415                &mut parse_state,
416                &mut decoded,
417                &mut slice_addr_of,
418                &mut ds_stored,
419                &mut wpp_stored,
420                self.tolerant,
421            )?;
422        }
423
424        // ---- §8.3 reference cycle + §8.4/§8.5/§8.7 reconstruction ----
425        let slice_type = indep
426            .header
427            .slice_type
428            .ok_or(SequenceError::Malformed("independent slice without type"))?;
429        let header_info = self.build_header_info(indep, sps, nal_kind, no_rasl_output)?;
430        let slice_ref = build_slice_ref_params(&indep.header, pps, slice_type, &header_info);
431
432        let ref_state = self.state.begin_picture(&header_info, &slice_ref);
433        let lists = ref_state.ref_pic_lists.clone().unwrap_or(RefPicLists {
434            list0: Vec::new(),
435            list1: None,
436        });
437        let entries = self.state.dpb().entries();
438        let col_field = ref_state.col_pic.map(|idx| &entries[idx].motion);
439        let col_poc = ref_state
440            .col_pic
441            .map(|idx| entries[idx].poc)
442            .unwrap_or_default();
443        let refs = RefListAccess {
444            lists: &lists,
445            entries,
446        };
447
448        let recon_params = build_recon_params(&indep.header, sps, pps, &geom)?;
449        let slice_ctx = build_inter_slice_context(
450            &indep.header,
451            sps,
452            pps,
453            &geom,
454            &recon_params,
455            ref_state.poc.val,
456            col_poc,
457            ref_state.no_backward_pred,
458            slice_type,
459        );
460
461        // Per-slice slice_loop_filter_across_slices_enabled_flag
462        // (§7.4.7.1: inferred from the PPS flag when absent).
463        let mut across_of_slice: BTreeMap<u32, bool> = BTreeMap::new();
464        for seg in segs {
465            if !seg.header.dependent_slice_segment_flag {
466                across_of_slice.insert(
467                    seg.header.slice_segment_address,
468                    seg.header
469                        .slice_loop_filter_across_slices_enabled_flag
470                        .unwrap_or(pps.pps_loop_filter_across_slices_enabled_flag),
471                );
472            }
473        }
474        let placed: Vec<PlacedInterCtu<'_>> = decoded
475            .iter()
476            .map(|(x, y, ctu)| {
477                let rs = (y >> geom.ctb_log2) * geom.pic_w_ctbs + (x >> geom.ctb_log2);
478                let slice_addr_rs = slice_addr_of[rs as usize].unwrap_or(0);
479                PlacedInterCtu {
480                    x_ctb: *x,
481                    y_ctb: *y,
482                    slice_addr_rs,
483                    filter_across_slices: across_of_slice
484                        .get(&slice_addr_rs)
485                        .copied()
486                        .unwrap_or(pps.pps_loop_filter_across_slices_enabled_flag),
487                    ctu,
488                }
489            })
490            .collect();
491
492        let (picture, motion) = reconstruct_inter_picture(
493            geom.width as usize,
494            geom.height as usize,
495            &recon_params,
496            &slice_ctx,
497            &geom.tiles,
498            &placed,
499            &refs,
500            col_field,
501        )?;
502
503        let output = indep.header.pic_output_flag;
504        let poc = ref_state.poc;
505        self.frames.push(DecodedFrame {
506            cvs_index: self.cvs_index,
507            poc: poc.val,
508            output,
509            picture: picture.clone(),
510        });
511        self.state
512            .store_picture(poc, indep.layer_id, picture, motion);
513        Ok(())
514    }
515
516    /// Assemble the §8.3 [`PictureHeaderInfo`] from the independent
517    /// slice segment header.
518    fn build_header_info(
519        &self,
520        seg: &SegmentData,
521        sps: &SeqParameterSet,
522        nal_kind: NalKind,
523        no_rasl_output: bool,
524    ) -> Result<PictureHeaderInfo, SequenceError> {
525        let max_poc_lsb = 1u32 << (sps.log2_max_pic_order_cnt_lsb_minus4 + 4);
526        let short_term_rps = materialize_slice_rps(&seg.header, sps)?;
527        let mut long_term = Vec::new();
528        for lt in &seg.header.long_term_ref_pics {
529            let poc_lsb_lt = match lt.source {
530                SliceLongTermRefPicSource::Sps { lt_idx_sps } => {
531                    sps.long_term_ref_pics
532                        .get(lt_idx_sps as usize)
533                        .ok_or(SequenceError::Malformed(
534                            "lt_idx_sps out of range of the SPS long-term table",
535                        ))?
536                        .poc_lsb
537                }
538                SliceLongTermRefPicSource::InSlice { poc_lsb_lt, .. } => poc_lsb_lt,
539            };
540            let used = lt.used_by_curr_pic_lt(sps).ok_or(SequenceError::Malformed(
541                "lt_idx_sps out of range of the SPS long-term table",
542            ))?;
543            long_term.push(LongTermEntry {
544                poc_lsb_lt,
545                used_by_curr_pic_lt: used,
546                delta_poc_msb_present: lt.delta_poc_msb_present_flag,
547                delta_poc_msb_cycle_lt: lt.delta_poc_msb_cycle_lt,
548            });
549        }
550        Ok(PictureHeaderInfo {
551            nal_kind,
552            temporal_id: seg.temporal_id,
553            layer_id: seg.layer_id,
554            no_rasl_output,
555            poc_lsb: seg.header.slice_pic_order_cnt_lsb.unwrap_or(0),
556            max_poc_lsb,
557            short_term_rps,
558            long_term,
559        })
560    }
561}
562
563/// Decode a whole Annex B byte stream to its output-order frames.
564///
565/// # Errors
566/// Any demux / parse / decode error.
567pub fn decode_annexb_sequence(data: &[u8]) -> Result<Vec<DecodedFrame>, SequenceError> {
568    let mut dec = SequenceDecoder::new();
569    dec.push_annexb(data)?;
570    dec.finish()
571}
572
573/// Debug helper: CABAC-decode the FIRST picture's slice-segment data and
574/// return whatever CTUs were decoded, even when the walk diverges (the
575/// `end_of_slice_segment_flag` never fires). Not part of the stable API.
576#[doc(hidden)]
577pub fn decode_annexb_sequence_debug(
578    data: &[u8],
579) -> Result<Vec<(u32, u32, CodingTreeUnit)>, SequenceError> {
580    // Which picture (1-based) to CABAC-decode; earlier pictures are
581    // fully decoded through the normal driver so the DPB and parse
582    // state are real.
583    let target: usize = std::env::var("H265_DEBUG_PIC")
584        .ok()
585        .and_then(|v| v.parse().ok())
586        .unwrap_or(1);
587    let mut dec = SequenceDecoder::new();
588    let mut segs: Vec<SegmentData> = Vec::new();
589    let mut pic_no = 0usize;
590    for unit in NalIter::new(data) {
591        let unit = unit?;
592        if unit.header.is_vcl() {
593            let first = unit.rbsp.first().is_some_and(|b| b & 0x80 != 0);
594            if first {
595                pic_no += 1;
596                if pic_no > target && !segs.is_empty() {
597                    break;
598                }
599            }
600            if pic_no < target {
601                dec.push_nal_unit(unit)?;
602                continue;
603            }
604            let pps_id = peek_slice_pps_id(&unit.rbsp, unit.header.nal_unit_type)?;
605            let pps = dec.pps.get(&pps_id).unwrap();
606            let sps = dec.sps.get(&pps.sps_id).unwrap();
607            let parsed =
608                SliceSegmentHeader::parse(&unit.rbsp, unit.header.nal_unit_type, sps, pps)?;
609            segs.push(SegmentData {
610                nal_type: unit.header.nal_unit_type,
611                temporal_id: unit.header.temporal_id,
612                layer_id: unit.header.nuh_layer_id,
613                rbsp: unit.rbsp,
614                escaped: unit.escaped,
615                header: parsed,
616            });
617        } else {
618            dec.push_nal_unit(unit)?;
619        }
620    }
621    let seg = &segs[0];
622    let pps = dec.pps.get(&seg.header.slice_pic_parameter_set_id).unwrap();
623    let sps = dec.sps.get(&pps.sps_id).unwrap();
624    let geom = Geometry::derive(sps, pps)?;
625    let pic_size = (geom.pic_w_ctbs * geom.pic_h_ctbs) as usize;
626    let mut decoded = Vec::new();
627    let mut slice_addr_of = vec![None; pic_size];
628    let st = seg.header.slice_type.unwrap();
629    let mut parse_state =
630        PictureParseState::new(&build_slice_data_params(&seg.header, sps, pps, &geom, st));
631    let mut ds_stored = None;
632    let mut wpp_stored = None;
633    let res = decode_slice_segment_data(
634        seg,
635        &seg.header,
636        sps,
637        pps,
638        &geom,
639        &mut parse_state,
640        &mut decoded,
641        &mut slice_addr_of,
642        &mut ds_stored,
643        &mut wpp_stored,
644        false,
645    );
646    if let Err(e) = res {
647        eprintln!("(walk error: {e})");
648    }
649    Ok(decoded)
650}
651
652/// Debug helper: reconstruct the FIRST picture even when the CABAC walk
653/// diverges, returning the in-loop-filtered picture. Not stable API.
654#[doc(hidden)]
655pub fn decode_annexb_first_picture_tolerant(data: &[u8]) -> Result<Picture, SequenceError> {
656    let mut dec = SequenceDecoder::new();
657    let mut segs: Vec<SegmentData> = Vec::new();
658    for unit in NalIter::new(data) {
659        let unit = unit?;
660        if unit.header.is_vcl() {
661            let first = unit.rbsp.first().is_some_and(|b| b & 0x80 != 0);
662            if first && !segs.is_empty() {
663                break;
664            }
665            let pps_id = peek_slice_pps_id(&unit.rbsp, unit.header.nal_unit_type)?;
666            let pps = dec.pps.get(&pps_id).unwrap();
667            let sps = dec.sps.get(&pps.sps_id).unwrap();
668            let parsed =
669                SliceSegmentHeader::parse(&unit.rbsp, unit.header.nal_unit_type, sps, pps)?;
670            segs.push(SegmentData {
671                nal_type: unit.header.nal_unit_type,
672                temporal_id: unit.header.temporal_id,
673                layer_id: unit.header.nuh_layer_id,
674                rbsp: unit.rbsp,
675                escaped: unit.escaped,
676                header: parsed,
677            });
678        } else {
679            dec.push_nal_unit(unit)?;
680        }
681    }
682    let seg = &segs[0];
683    let pps = dec
684        .pps
685        .get(&seg.header.slice_pic_parameter_set_id)
686        .unwrap()
687        .clone();
688    let sps = dec.sps.get(&pps.sps_id).unwrap().clone();
689    let geom = Geometry::derive(&sps, &pps)?;
690    let pic_size = (geom.pic_w_ctbs * geom.pic_h_ctbs) as usize;
691    let mut decoded = Vec::new();
692    let mut slice_addr_of = vec![None; pic_size];
693    let st0 = seg.header.slice_type.unwrap();
694    let mut parse_state = PictureParseState::new(&build_slice_data_params(
695        &seg.header,
696        &sps,
697        &pps,
698        &geom,
699        st0,
700    ));
701    let mut ds_stored = None;
702    let mut wpp_stored = None;
703    if let Err(e) = decode_slice_segment_data(
704        seg,
705        &seg.header,
706        &sps,
707        &pps,
708        &geom,
709        &mut parse_state,
710        &mut decoded,
711        &mut slice_addr_of,
712        &mut ds_stored,
713        &mut wpp_stored,
714        true,
715    ) {
716        eprintln!("(walk error: {e})");
717    }
718    let slice_type = seg.header.slice_type.unwrap();
719    let recon_params = build_recon_params(&seg.header, &sps, &pps, &geom)?;
720    let slice_ctx = build_inter_slice_context(
721        &seg.header,
722        &sps,
723        &pps,
724        &geom,
725        &recon_params,
726        0,
727        0,
728        true,
729        slice_type,
730    );
731    let placed: Vec<PlacedInterCtu<'_>> = decoded
732        .iter()
733        .map(|(x, y, ctu)| PlacedInterCtu {
734            x_ctb: *x,
735            y_ctb: *y,
736            slice_addr_rs: 0,
737            filter_across_slices: true,
738            ctu,
739        })
740        .collect();
741    let lists = RefPicLists {
742        list0: Vec::new(),
743        list1: None,
744    };
745    let refs = RefListAccess {
746        lists: &lists,
747        entries: &[],
748    };
749    let (picture, _) = reconstruct_inter_picture(
750        geom.width as usize,
751        geom.height as usize,
752        &recon_params,
753        &slice_ctx,
754        &geom.tiles,
755        &placed,
756        &refs,
757        None,
758    )?;
759    Ok(picture)
760}
761
762/// The per-SPS geometry constants (§7.4.3.2.1 derived variables).
763struct Geometry {
764    width: u32,
765    height: u32,
766    ctb_log2: u32,
767    min_cb_log2: u32,
768    min_tb_log2: u32,
769    max_tb_log2: u32,
770    chroma_array_type: u8,
771    pic_w_ctbs: u32,
772    pic_h_ctbs: u32,
773    tiles: TilingParams,
774}
775
776impl Geometry {
777    fn derive(sps: &SeqParameterSet, pps: &PicParameterSet) -> Result<Self, SequenceError> {
778        let min_cb_log2 = u32::from(sps.log2_min_luma_coding_block_size_minus3) + 3;
779        let ctb_log2 = min_cb_log2 + u32::from(sps.log2_diff_max_min_luma_coding_block_size);
780        let min_tb_log2 = u32::from(sps.log2_min_luma_transform_block_size_minus2) + 2;
781        let max_tb_log2 = min_tb_log2 + u32::from(sps.log2_diff_max_min_luma_transform_block_size);
782        let width = sps.pic_width_in_luma_samples;
783        let height = sps.pic_height_in_luma_samples;
784        if width == 0 || height == 0 {
785            return Err(SequenceError::Malformed("zero picture dimensions"));
786        }
787        let ctb = 1u32 << ctb_log2;
788        let chroma_array_type = if sps.separate_colour_plane_flag {
789            0
790        } else {
791            sps.chroma_format_idc
792        };
793        let tiles = if pps.tiles_enabled_flag {
794            TilingParams {
795                num_tile_columns_minus1: pps.tiles.num_tile_columns_minus1,
796                num_tile_rows_minus1: pps.tiles.num_tile_rows_minus1,
797                uniform_spacing_flag: pps.tiles.uniform_spacing_flag,
798                column_width_minus1: pps.tiles.column_width_minus1.clone(),
799                row_height_minus1: pps.tiles.row_height_minus1.clone(),
800            }
801        } else {
802            TilingParams::single_tile()
803        };
804        Ok(Self {
805            width,
806            height,
807            ctb_log2,
808            min_cb_log2,
809            min_tb_log2,
810            max_tb_log2,
811            chroma_array_type,
812            pic_w_ctbs: width.div_ceil(ctb),
813            pic_h_ctbs: height.div_ceil(ctb),
814            tiles,
815        })
816    }
817
818    fn tiling(&self) -> Result<PictureTiling, SequenceError> {
819        PictureTiling::new(
820            self.pic_w_ctbs,
821            self.pic_h_ctbs,
822            self.width,
823            self.height,
824            self.ctb_log2,
825            self.min_tb_log2,
826            &self.tiles,
827        )
828        .map_err(|_| SequenceError::Malformed("invalid tile geometry"))
829    }
830}
831
832/// Pre-read `slice_pic_parameter_set_id` from a slice-segment RBSP (the
833/// two leading fields before it are fixed-width).
834fn peek_slice_pps_id(rbsp: &[u8], nal_unit_type: u8) -> Result<u8, SequenceError> {
835    let mut br = BitReader::new(rbsp);
836    let _first = br.u1().map_err(|_| {
837        SequenceError::Malformed("slice header truncated before first_slice_segment_in_pic_flag")
838    })?;
839    if (NalKind::BLA_W_LP..=NalKind::RSV_IRAP_VCL23).contains(&nal_unit_type) {
840        let _ = br.u1().map_err(|_| {
841            SequenceError::Malformed("slice header truncated at no_output_of_prior_pics_flag")
842        })?;
843    }
844    let pps_id = br.ue().map_err(|_| {
845        SequenceError::Malformed("slice header truncated at slice_pic_parameter_set_id")
846    })?;
847    if pps_id > 63 {
848        return Err(SequenceError::Malformed(
849            "slice_pic_parameter_set_id out of range",
850        ));
851    }
852    Ok(pps_id as u8)
853}
854
855/// §7.4.8 — materialize the slice's short-term RPS (SPS-indexed or
856/// slice-inline, explicit or inter-RPS-predicted).
857fn materialize_slice_rps(
858    header: &SliceSegmentHeader,
859    sps: &SeqParameterSet,
860) -> Result<MaterializedShortTermRefPicSet, SequenceError> {
861    // IDR: no RPS block — empty set.
862    let Some(sps_flag) = header.short_term_ref_pic_set_sps_flag else {
863        return Ok(MaterializedShortTermRefPicSet {
864            delta_poc_s0: Vec::new(),
865            used_by_curr_pic_s0: Vec::new(),
866            delta_poc_s1: Vec::new(),
867            used_by_curr_pic_s1: Vec::new(),
868        });
869    };
870    // Materialize the SPS chain (set i may inter-predict from an
871    // earlier set).
872    let mut chain: Vec<MaterializedShortTermRefPicSet> =
873        Vec::with_capacity(sps.short_term_ref_pic_sets.len());
874    for (idx, set) in sps.short_term_ref_pic_sets.iter().enumerate() {
875        let source = if set.inter_ref_pic_set_prediction_flag {
876            let ref_idx = idx
877                .checked_sub(set.delta_idx_minus1 as usize + 1)
878                .ok_or(SequenceError::Malformed("RefRpsIdx underflow"))?;
879            Some(&chain[ref_idx])
880        } else {
881            None
882        };
883        chain.push(set.materialize(source)?);
884    }
885    if sps_flag {
886        let idx = header.short_term_ref_pic_set_idx.unwrap_or(0) as usize;
887        chain
888            .into_iter()
889            .nth(idx)
890            .ok_or(SequenceError::Malformed("short_term_ref_pic_set_idx OOR"))
891    } else {
892        let set = header
893            .inline_short_term_ref_pic_set
894            .as_ref()
895            .ok_or(SequenceError::Malformed("missing inline st_ref_pic_set"))?;
896        let source = if set.inter_ref_pic_set_prediction_flag {
897            // stRpsIdx == num_short_term_ref_pic_sets for the inline set.
898            let ref_idx = chain
899                .len()
900                .checked_sub(set.delta_idx_minus1 as usize + 1)
901                .ok_or(SequenceError::Malformed("RefRpsIdx underflow"))?;
902            Some(&chain[ref_idx])
903        } else {
904            None
905        };
906        Ok(set.materialize(source)?)
907    }
908}
909
910/// `NumPicTotalCurr` (§7.4.7.2) from the already-resolved picture
911/// header info (single-layer; `pps_curr_pic_ref_enabled_flag`
912/// contributes the closing `NumPicTotalCurr++`).
913fn num_pic_total_curr(info: &PictureHeaderInfo, curr_pic_ref_enabled: bool) -> u32 {
914    let st = info
915        .short_term_rps
916        .used_by_curr_pic_s0
917        .iter()
918        .chain(info.short_term_rps.used_by_curr_pic_s1.iter())
919        .filter(|&&u| u)
920        .count();
921    let lt = info
922        .long_term
923        .iter()
924        .filter(|e| e.used_by_curr_pic_lt)
925        .count();
926    (st + lt) as u32 + u32::from(curr_pic_ref_enabled)
927}
928
929fn build_slice_ref_params(
930    header: &SliceSegmentHeader,
931    pps: &PicParameterSet,
932    slice_type: SliceType,
933    info: &PictureHeaderInfo,
934) -> SliceRefParams {
935    let is_b = slice_type == SliceType::B;
936    let is_inter = slice_type != SliceType::I;
937    let curr_pic_ref_enabled = pps
938        .pps_scc_extension
939        .as_ref()
940        .is_some_and(|s| s.pps_curr_pic_ref_enabled_flag);
941    SliceRefParams {
942        is_inter,
943        is_b,
944        num_ref_idx_l0_active_minus1: u32::from(
945            header
946                .num_ref_idx_l0_active_minus1
947                .unwrap_or(pps.num_ref_idx_l0_default_active_minus1),
948        ),
949        num_ref_idx_l1_active_minus1: u32::from(
950            header
951                .num_ref_idx_l1_active_minus1
952                .unwrap_or(pps.num_ref_idx_l1_default_active_minus1),
953        ),
954        num_pic_total_curr: num_pic_total_curr(info, curr_pic_ref_enabled),
955        temporal_mvp_enabled: header.slice_temporal_mvp_enabled_flag,
956        collocated_from_l0_flag: header.collocated_from_l0_flag.unwrap_or(true),
957        collocated_ref_idx: header.collocated_ref_idx.unwrap_or(0),
958        curr_pic_ref_enabled,
959    }
960}
961
962fn build_recon_params(
963    header: &SliceSegmentHeader,
964    sps: &SeqParameterSet,
965    pps: &PicParameterSet,
966    geom: &Geometry,
967) -> Result<ReconParams, SequenceError> {
968    let slice_qp_y = header
969        .slice_qp_y(pps)
970        .ok_or(SequenceError::Malformed("slice header without slice_qp"))?;
971    let range = sps.sps_range_extension.as_ref();
972    // §7.4.5: when scaling_list_enabled_flag == 1 the active
973    // scaling-list data is the PPS body if present, else the SPS body
974    // if present, else the default lists.
975    let scaling = if sps.scaling_list_enabled_flag {
976        let factors = match (&pps.scaling_list_data, &sps.scaling_list_data) {
977            (Some(d), _) => d.scaling_factors(geom.chroma_array_type),
978            (None, Some(d)) => d.scaling_factors(geom.chroma_array_type),
979            (None, None) => crate::scaling_list::ScalingListData::all_default()
980                .scaling_factors(geom.chroma_array_type),
981        };
982        Some(factors)
983    } else {
984        None
985    };
986    Ok(ReconParams {
987        chroma_array_type: geom.chroma_array_type,
988        bit_depth_luma: sps.bit_depth_luma_minus8 + 8,
989        bit_depth_chroma: sps.bit_depth_chroma_minus8 + 8,
990        intra_smoothing_disabled: range.is_some_and(|r| r.intra_smoothing_disabled_flag),
991        strong_intra_smoothing_enabled: sps.strong_intra_smoothing_enabled_flag,
992        slice_qp_y,
993        cb_qp_offset: i32::from(pps.pps_cb_qp_offset) + i32::from(header.slice_cb_qp_offset),
994        cr_qp_offset: i32::from(pps.pps_cr_qp_offset) + i32::from(header.slice_cr_qp_offset),
995        // §7.4.3.3.3: PpsActQpOffset{Y,Cb,Cr} = pps_act_{y,cb}_qp_offset_plus5 − 5
996        // / pps_act_cr_qp_offset_plus3 − 3; the slice offsets add on top
997        // (§7.4.7.1), each 0 when absent.
998        act_y_qp_offset: pps
999            .pps_scc_extension
1000            .as_ref()
1001            .map_or(-5, |s| s.pps_act_y_qp_offset_plus5 - 5)
1002            + header.slice_act_y_qp_offset,
1003        act_cb_qp_offset: pps
1004            .pps_scc_extension
1005            .as_ref()
1006            .map_or(-5, |s| s.pps_act_cb_qp_offset_plus5 - 5)
1007            + header.slice_act_cb_qp_offset,
1008        act_cr_qp_offset: pps
1009            .pps_scc_extension
1010            .as_ref()
1011            .map_or(-3, |s| s.pps_act_cr_qp_offset_plus3 - 3)
1012            + header.slice_act_cr_qp_offset,
1013        transform_skip_rotation_enabled: range
1014            .is_some_and(|r| r.transform_skip_rotation_enabled_flag),
1015        implicit_rdpcm_enabled: range.is_some_and(|r| r.implicit_rdpcm_enabled_flag),
1016        intra_boundary_filtering_disabled: sps
1017            .sps_scc_extension
1018            .as_ref()
1019            .is_some_and(|s| s.intra_boundary_filtering_disabled_flag),
1020        extended_precision: range.is_some_and(|r| r.extended_precision_processing_flag),
1021        scaling,
1022        chroma_qp_offset_list: pps
1023            .pps_range_extension
1024            .as_ref()
1025            .map(|r| {
1026                r.chroma_qp_offset_list
1027                    .iter()
1028                    .map(|e| (i32::from(e.cb_qp_offset), i32::from(e.cr_qp_offset)))
1029                    .collect()
1030            })
1031            .unwrap_or_default(),
1032        cu_qp_offset_c: core::cell::Cell::new((0, 0)),
1033    })
1034}
1035
1036#[allow(clippy::too_many_arguments)]
1037fn build_inter_slice_context(
1038    header: &SliceSegmentHeader,
1039    sps: &SeqParameterSet,
1040    pps: &PicParameterSet,
1041    geom: &Geometry,
1042    recon: &ReconParams,
1043    curr_poc: i32,
1044    col_poc: i32,
1045    no_backward_pred: bool,
1046    slice_type: SliceType,
1047) -> InterSliceContext {
1048    let pps_range = pps.pps_range_extension.as_ref();
1049    let deblock = header.deblocking.as_ref();
1050    // §8.5.3.3.4.1 — weightedPredFlag: weighted_pred_flag for P slices,
1051    // weighted_bipred_flag for B slices.
1052    let weighted_pred_flag = match slice_type {
1053        SliceType::P => pps.weighted_pred_flag,
1054        SliceType::B => pps.weighted_bipred_flag,
1055        SliceType::I => false,
1056    };
1057    let wp = if weighted_pred_flag {
1058        header
1059            .pred_weight_table
1060            .as_ref()
1061            .map(|pwt| build_slice_wp_tables(pwt, sps))
1062    } else {
1063        None
1064    };
1065    InterSliceContext {
1066        curr_poc,
1067        constrained_intra_pred: pps.constrained_intra_pred_flag,
1068        slice_is_b: slice_type == SliceType::B,
1069        ctb_log2_size_y: geom.ctb_log2,
1070        pic_width_luma: geom.width,
1071        pic_height_luma: geom.height,
1072        max_num_merge_cand: usize::from(header.max_num_merge_cand().unwrap_or(5)),
1073        num_ref_idx_l0_active: i32::from(
1074            header
1075                .num_ref_idx_l0_active_minus1
1076                .unwrap_or(pps.num_ref_idx_l0_default_active_minus1),
1077        ) + 1,
1078        num_ref_idx_l1_active: i32::from(
1079            header
1080                .num_ref_idx_l1_active_minus1
1081                .unwrap_or(pps.num_ref_idx_l1_default_active_minus1),
1082        ) + 1,
1083        log2_par_mrg_level: pps.log2_parallel_merge_level_minus2 + 2,
1084        temporal_mvp_enabled: header.slice_temporal_mvp_enabled_flag,
1085        collocated_from_l0_flag: header.collocated_from_l0_flag.unwrap_or(true),
1086        col_poc,
1087        no_backward_pred,
1088        min_tb_log2_size_y: geom.min_tb_log2,
1089        log2_min_cu_qp_delta_size: geom.ctb_log2 - pps.diff_cu_qp_delta_depth,
1090        wpp_qp_row_reset: pps.entropy_coding_sync_enabled_flag,
1091        filter_across_slices: header
1092            .slice_loop_filter_across_slices_enabled_flag
1093            .unwrap_or(pps.pps_loop_filter_across_slices_enabled_flag),
1094        filter_across_tiles: pps.loop_filter_across_tiles_enabled_flag,
1095        deblock_enabled: deblock.map_or(true, |d| !d.disabled_flag),
1096        beta_offset_div2: deblock.map_or(0, |d| i32::from(d.beta_offset_div2)),
1097        tc_offset_div2: deblock.map_or(0, |d| i32::from(d.tc_offset_div2)),
1098        slice_qp_y: recon.slice_qp_y,
1099        cb_qp_offset: recon.cb_qp_offset,
1100        cr_qp_offset: recon.cr_qp_offset,
1101        pps_cb_qp_offset: i32::from(pps.pps_cb_qp_offset),
1102        pps_cr_qp_offset: i32::from(pps.pps_cr_qp_offset),
1103        slice_sao_luma_flag: header.slice_sao_luma_flag,
1104        slice_sao_chroma_flag: header.slice_sao_chroma_flag,
1105        log2_sao_offset_scale_luma: pps_range.map_or(0, |r| r.log2_sao_offset_scale_luma as u8),
1106        log2_sao_offset_scale_chroma: pps_range.map_or(0, |r| r.log2_sao_offset_scale_chroma as u8),
1107        wp,
1108        pcm_loop_filter_disabled: sps
1109            .pcm
1110            .as_ref()
1111            .is_some_and(|p| p.loop_filter_disabled_flag),
1112        use_integer_mv: header.use_integer_mv_flag,
1113        // §7.4.3.3.3 eq. 7-40: TwoVersionsOfCurrDecPicFlag =
1114        // pps_curr_pic_ref_enabled_flag && ( sao enabled ||
1115        // !pps_deblocking_filter_disabled_flag ||
1116        // deblocking_filter_override_enabled_flag ).
1117        two_versions_curr_pic: pps
1118            .pps_scc_extension
1119            .as_ref()
1120            .is_some_and(|s| s.pps_curr_pic_ref_enabled_flag)
1121            && (sps.sample_adaptive_offset_enabled_flag
1122                || !pps.deblocking.disabled_flag
1123                || pps.deblocking.override_enabled_flag),
1124    }
1125}
1126
1127/// §7.4.7.3 — resolve a parsed `pred_weight_table()` into the
1128/// per-reference values the §8.5.3.3.4.3 combine reads: `LumaWeightLX[i]`
1129/// / `ChromaWeightLX[i][j]` (weight-flag inference included), the
1130/// `WpOffsetBdShiftY`- / `WpOffsetBdShiftC`-scaled offsets (equations
1131/// 7-31 / 7-32 + 8-268 / 8-269 / 8-273 / 8-274), and the equation-7-58
1132/// `ChromaOffsetLX` derivation.
1133fn build_slice_wp_tables(
1134    pwt: &crate::slice::PredWeightTable,
1135    sps: &SeqParameterSet,
1136) -> SliceWpTables {
1137    let hp = sps
1138        .sps_range_extension
1139        .as_ref()
1140        .is_some_and(|r| r.high_precision_offsets_enabled_flag);
1141    let bd_y = i32::from(sps.bit_depth_luma_minus8) + 8;
1142    let bd_c = i32::from(sps.bit_depth_chroma_minus8) + 8;
1143    // Equations 7-31 / 7-32 / 7-34.
1144    let bd_shift_y = if hp { 0 } else { bd_y - 8 };
1145    let bd_shift_c = if hp { 0 } else { bd_c - 8 };
1146    let half_range_c = 1i32 << (if hp { bd_c - 1 } else { 7 });
1147    let chroma_denom = pwt.chroma_log2_weight_denom();
1148
1149    let resolve = |l0: bool, n: usize| -> Vec<WpListWeights> {
1150        (0..n)
1151            .map(|i| {
1152                let (lw, lo, cw0, cw1, co0, co1) = if l0 {
1153                    (
1154                        pwt.luma_weight_l0(i),
1155                        pwt.entries_l0.get(i).map(|e| e.luma_offset),
1156                        pwt.chroma_weight_l0(i, 0),
1157                        pwt.chroma_weight_l0(i, 1),
1158                        pwt.chroma_offset_l0(i, 0, half_range_c),
1159                        pwt.chroma_offset_l0(i, 1, half_range_c),
1160                    )
1161                } else {
1162                    (
1163                        pwt.luma_weight_l1(i),
1164                        pwt.entries_l1.get(i).map(|e| e.luma_offset),
1165                        pwt.chroma_weight_l1(i, 0),
1166                        pwt.chroma_weight_l1(i, 1),
1167                        pwt.chroma_offset_l1(i, 0, half_range_c),
1168                        pwt.chroma_offset_l1(i, 1, half_range_c),
1169                    )
1170                };
1171                WpListWeights {
1172                    w_luma: lw.unwrap_or(1 << pwt.luma_log2_weight_denom),
1173                    o_luma: lo.unwrap_or(0) << bd_shift_y,
1174                    w_cb: cw0.unwrap_or(1 << chroma_denom),
1175                    o_cb: co0.unwrap_or(0) << bd_shift_c,
1176                    w_cr: cw1.unwrap_or(1 << chroma_denom),
1177                    o_cr: co1.unwrap_or(0) << bd_shift_c,
1178                }
1179            })
1180            .collect()
1181    };
1182
1183    SliceWpTables {
1184        luma_log2_weight_denom: pwt.luma_log2_weight_denom,
1185        chroma_log2_weight_denom: chroma_denom,
1186        l0: resolve(true, pwt.entries_l0.len()),
1187        l1: resolve(false, pwt.entries_l1.len()),
1188    }
1189}
1190
1191/// §7.4.3 — derive the [`SliceDataParams`] for one slice segment.
1192fn build_slice_data_params(
1193    header: &SliceSegmentHeader,
1194    sps: &SeqParameterSet,
1195    pps: &PicParameterSet,
1196    geom: &Geometry,
1197    slice_type: SliceType,
1198) -> SliceDataParams {
1199    let pps_range = pps.pps_range_extension.as_ref();
1200    let (log2_min_ipcm, log2_max_ipcm) = sps.pcm.as_ref().map_or((3, 5), |p| {
1201        let min = u32::from(p.log2_min_pcm_luma_coding_block_size_minus3) + 3;
1202        (
1203            min,
1204            min + u32::from(p.log2_diff_max_min_pcm_luma_coding_block_size),
1205        )
1206    });
1207    let cu_chroma_qp_offset_enabled = header.cu_chroma_qp_offset_enabled_flag;
1208    let log2_min_cu_chroma_qp_offset_size =
1209        geom.ctb_log2 - pps_range.map_or(0, |r| r.diff_cu_chroma_qp_offset_depth);
1210    let scc = sps.sps_scc_extension.as_ref();
1211    let palette_max_size = scc.map_or(0, |e| e.palette_max_size);
1212    SliceDataParams {
1213        ctb_log2_size_y: geom.ctb_log2,
1214        min_cb_log2_size_y: geom.min_cb_log2,
1215        max_tb_log2_size_y: geom.max_tb_log2,
1216        min_tb_log2_size_y: geom.min_tb_log2,
1217        pic_width_in_luma_samples: geom.width,
1218        pic_height_in_luma_samples: geom.height,
1219        chroma_array_type: geom.chroma_array_type,
1220        bit_depth_luma: u32::from(sps.bit_depth_luma_minus8) + 8,
1221        bit_depth_chroma: u32::from(sps.bit_depth_chroma_minus8) + 8,
1222        slice_type_is_i: slice_type == SliceType::I,
1223        slice_type_is_b: slice_type == SliceType::B,
1224        slice_sao_luma_flag: header.slice_sao_luma_flag,
1225        slice_sao_chroma_flag: header.slice_sao_chroma_flag,
1226        transquant_bypass_enabled_flag: pps.transquant_bypass_enabled_flag,
1227        cu_qp_delta_enabled_flag: pps.cu_qp_delta_enabled_flag,
1228        log2_min_cu_qp_delta_size: geom.ctb_log2 - pps.diff_cu_qp_delta_depth,
1229        cu_chroma_qp_offset_enabled_flag: cu_chroma_qp_offset_enabled,
1230        log2_min_cu_chroma_qp_offset_size,
1231        chroma_qp_offset_list_len_minus1: pps_range
1232            .map_or(0, |r| r.chroma_qp_offset_list_len_minus1),
1233        amp_enabled_flag: sps.amp_enabled_flag,
1234        pcm_enabled_flag: sps.pcm_enabled_flag,
1235        log2_min_ipcm_cb_size_y: log2_min_ipcm,
1236        log2_max_ipcm_cb_size_y: log2_max_ipcm,
1237        pcm_bit_depth_luma: sps
1238            .pcm
1239            .as_ref()
1240            .map_or(8, |p| u32::from(p.bit_depth_luma_minus1) + 1),
1241        pcm_bit_depth_chroma: sps
1242            .pcm
1243            .as_ref()
1244            .map_or(8, |p| u32::from(p.bit_depth_chroma_minus1) + 1),
1245        max_transform_hierarchy_depth_intra: u32::from(sps.max_transform_hierarchy_depth_intra),
1246        max_transform_hierarchy_depth_inter: u32::from(sps.max_transform_hierarchy_depth_inter),
1247        max_num_merge_cand: u32::from(header.max_num_merge_cand().unwrap_or(5)),
1248        num_ref_idx_l0_active_minus1: u32::from(
1249            header
1250                .num_ref_idx_l0_active_minus1
1251                .unwrap_or(pps.num_ref_idx_l0_default_active_minus1),
1252        ),
1253        num_ref_idx_l1_active_minus1: u32::from(
1254            header
1255                .num_ref_idx_l1_active_minus1
1256                .unwrap_or(pps.num_ref_idx_l1_default_active_minus1),
1257        ),
1258        mvd_l1_zero_flag: header.mvd_l1_zero_flag.unwrap_or(false),
1259        sign_data_hiding_enabled_flag: pps.sign_data_hiding_enabled_flag,
1260        cross_component_prediction_enabled_flag: pps_range
1261            .is_some_and(|r| r.cross_component_prediction_enabled_flag),
1262        residual_adaptive_colour_transform_enabled_flag: pps
1263            .pps_scc_extension
1264            .as_ref()
1265            .is_some_and(|s| s.residual_adaptive_colour_transform_enabled_flag),
1266        transform_skip_enabled_flag: pps.transform_skip_enabled_flag,
1267        log2_max_transform_skip_size: pps_range
1268            .map_or(2, |r| r.log2_max_transform_skip_block_size_minus2 + 2),
1269        implicit_rdpcm_enabled_flag: sps
1270            .sps_range_extension
1271            .as_ref()
1272            .is_some_and(|r| r.implicit_rdpcm_enabled_flag),
1273        explicit_rdpcm_enabled_flag: sps
1274            .sps_range_extension
1275            .as_ref()
1276            .is_some_and(|r| r.explicit_rdpcm_enabled_flag),
1277        transform_skip_context_enabled_flag: sps
1278            .sps_range_extension
1279            .as_ref()
1280            .is_some_and(|r| r.transform_skip_context_enabled_flag),
1281        persistent_rice_adaptation_enabled_flag: sps
1282            .sps_range_extension
1283            .as_ref()
1284            .is_some_and(|r| r.persistent_rice_adaptation_enabled_flag),
1285        cabac_bypass_alignment_enabled_flag: sps
1286            .sps_range_extension
1287            .as_ref()
1288            .is_some_and(|r| r.cabac_bypass_alignment_enabled_flag),
1289        extended_precision_processing_flag: sps
1290            .sps_range_extension
1291            .as_ref()
1292            .is_some_and(|r| r.extended_precision_processing_flag),
1293        palette_mode_enabled_flag: scc.is_some_and(|e| e.palette_mode_enabled_flag),
1294        palette_max_size,
1295        palette_max_predictor_size: palette_max_size
1296            + scc.map_or(0, |e| e.delta_palette_max_predictor_size),
1297    }
1298}
1299
1300/// §7.3.8.1 — CABAC-decode one slice segment's `slice_segment_data()`,
1301/// appending its CTUs (in tile-scan order) to `decoded` and recording
1302/// each CTB's `SliceAddrRs` in `slice_addr_of`.
1303///
1304/// `effective_header` supplies the slice-level values — for an
1305/// independent segment it is `seg.header` itself; for a dependent
1306/// segment it is the preceding independent segment's header (§7.4.7.1
1307/// inheritance). `ds_stored` is the picture's §9.3.2.4
1308/// `TableStateIdxDs` context store: read (synchronized, §9.3.2.5 /
1309/// §9.3.2.2) at a dependent segment's start, written at every
1310/// segment's `end_of_slice_segment_flag == 1` while
1311/// `dependent_slice_segments_enabled_flag` is set.
1312#[allow(clippy::too_many_arguments)]
1313fn decode_slice_segment_data(
1314    seg: &SegmentData,
1315    effective_header: &SliceSegmentHeader,
1316    sps: &SeqParameterSet,
1317    pps: &PicParameterSet,
1318    geom: &Geometry,
1319    state: &mut PictureParseState,
1320    decoded: &mut Vec<(u32, u32, CodingTreeUnit)>,
1321    slice_addr_of: &mut [Option<u32>],
1322    ds_stored: &mut Option<SliceContexts>,
1323    wpp_stored: &mut Option<SliceContexts>,
1324    tolerant: bool,
1325) -> Result<(), SequenceError> {
1326    let header = effective_header;
1327    let slice_type = header
1328        .slice_type
1329        .ok_or(SequenceError::Malformed("independent slice without type"))?;
1330    let params = build_slice_data_params(header, sps, pps, geom, slice_type);
1331    let slice_qp_y = header
1332        .slice_qp_y(pps)
1333        .ok_or(SequenceError::Malformed("slice header without slice_qp"))?;
1334
1335    let data_offset = seg
1336        .header
1337        .byte_offset_to_slice_data
1338        .ok_or(SequenceError::Malformed("slice header without data offset"))?;
1339    if data_offset >= seg.rbsp.len() {
1340        return Err(SequenceError::Malformed("slice data offset out of range"));
1341    }
1342
1343    // §7.4.7.1 — split `slice_segment_data( )` into its subsets. The
1344    // entry-point offsets count CODED bytes (emulation-prevention bytes
1345    // included), so map each escaped boundary onto the stripped RBSP.
1346    // (Entry points are per-segment syntax: read them from the
1347    // segment's own header even when it is dependent.)
1348    let substreams = split_substreams(
1349        &seg.escaped,
1350        seg.rbsp.len(),
1351        data_offset,
1352        seg.header.entry_point_offsets.as_ref(),
1353    )?;
1354
1355    // Table 9-4 initType: I => 0; P => cabac_init ? 2 : 1;
1356    // B => cabac_init ? 1 : 2 (crate::cabac::init_type on the raw
1357    // slice_type value).
1358    let raw_slice_type = match slice_type {
1359        SliceType::B => 0,
1360        SliceType::P => 1,
1361        SliceType::I => 2,
1362    };
1363    let it = init_type(raw_slice_type, header.cabac_init_flag.unwrap_or(false));
1364
1365    let tiling = geom.tiling()?;
1366    let tiles_on = pps.tiles_enabled_flag;
1367    let wpp = pps.entropy_coding_sync_enabled_flag;
1368    // §7.4.7.1: SliceAddrRs is the INDEPENDENT segment's address; a
1369    // dependent segment starts decoding at its own segment address but
1370    // its CTBs belong to the inherited slice.
1371    let slice_addr_rs = header.slice_segment_address;
1372    let mut ctb_addr_ts = tiling.ctb_addr_rs_to_ts(seg.header.slice_segment_address);
1373    let pic_size_in_ctbs = (geom.pic_w_ctbs * geom.pic_h_ctbs) as usize;
1374
1375    let sub_range = |idx: usize| -> Result<&[u8], SequenceError> {
1376        let &(a, b) = substreams
1377            .get(idx)
1378            .ok_or(SequenceError::Malformed("more CTB rows than substreams"))?;
1379        seg.rbsp
1380            .get(a..b)
1381            .ok_or(SequenceError::Malformed("substream range out of RBSP"))
1382    };
1383    let mut sub_idx = 0usize;
1384    let mut engine = CabacEngine::new(BitReader::new(sub_range(0)?))
1385        .map_err(|_| SequenceError::Malformed("slice data too short for CABAC init"))?;
1386    // §9.3.2.2 — a dependent slice segment synchronizes its context
1387    // variables from TableStateIdxDs (§9.3.2.5) instead of
1388    // re-initializing.
1389    // §9.3.2.3 — the palette predictor re-initialization value (the
1390    // PPS initializers if present, else the SPS initializers, else
1391    // empty), applied wherever §9.3.2.1 re-initializes the context
1392    // variables. A dependent segment SYNCHRONIZES the predictor from
1393    // the stored state instead (it travels inside SliceContexts).
1394    let num_comps = if geom.chroma_array_type == 0 { 1 } else { 3 };
1395    let base_palette_predictor = pps
1396        .pps_scc_extension
1397        .as_ref()
1398        .filter(|e| e.pps_palette_predictor_initializers_present_flag)
1399        .map(|e| {
1400            crate::palette::PalettePredictor::from_initializers(
1401                &e.pps_palette_predictor_initializer,
1402                num_comps,
1403            )
1404        })
1405        .or_else(|| {
1406            sps.sps_scc_extension
1407                .as_ref()
1408                .filter(|e| e.sps_palette_predictor_initializers_present_flag)
1409                .map(|e| {
1410                    crate::palette::PalettePredictor::from_initializers(
1411                        &e.sps_palette_predictor_initializer,
1412                        num_comps,
1413                    )
1414                })
1415        })
1416        .unwrap_or_default();
1417    let fresh_contexts = || {
1418        let mut c = SliceContexts::init(it, slice_qp_y);
1419        c.palette_predictor = base_palette_predictor.clone();
1420        c
1421    };
1422    // §6.4.1-gated availability of the spatial neighbour T (eq. 9-3,
1423    // the above-right CTB) for the §9.3.2.5 WPP synchronization: T
1424    // must exist, lie in the SAME slice (the stored snapshot may come
1425    // from an earlier slice segment of that slice) and the same tile.
1426    let t_available = |ctb_addr_ts: u32, slice_addr_of: &[Option<u32>]| {
1427        let rs = tiling.ctb_addr_ts_to_rs(ctb_addr_ts);
1428        let (rx, ry) = (rs % geom.pic_w_ctbs, rs / geom.pic_w_ctbs);
1429        ry > 0 && rx + 1 < geom.pic_w_ctbs && {
1430            let t_rs = (ry - 1) * geom.pic_w_ctbs + rx + 1;
1431            slice_addr_of[t_rs as usize] == Some(slice_addr_rs)
1432                && tiling.tile_id(tiling.ctb_addr_rs_to_ts(t_rs)) == tiling.tile_id(ctb_addr_ts)
1433        }
1434    };
1435    // §9.3.2.1 — the initial context state of this slice segment. For
1436    // a DEPENDENT segment the branch order matters: a segment whose
1437    // first CTU is the first CTU of a tile RE-INITIALIZES (§9.3.2.2 /
1438    // §9.3.2.3), one whose first CTU starts a CTU row of a tile under
1439    // entropy_coding_sync SYNCHRONIZES from the WPP snapshot
1440    // (§9.3.2.5, T-availability gated), and only otherwise does the
1441    // §9.3.2.5 dependent-segment synchronization from TableStateIdxDs
1442    // apply. An independent segment re-initializes.
1443    let mut ctx = if seg.header.dependent_slice_segment_flag {
1444        let first_rs = seg.header.slice_segment_address;
1445        let (rx, _ry) = (first_rs % geom.pic_w_ctbs, first_rs / geom.pic_w_ctbs);
1446        let tile_start = tiles_on
1447            && ctb_addr_ts > 0
1448            && tiling.tile_id(ctb_addr_ts) != tiling.tile_id(ctb_addr_ts - 1);
1449        let wpp_row_start = wpp
1450            && !tile_start
1451            && (rx == 0
1452                || tiling.tile_id(tiling.ctb_addr_rs_to_ts(first_rs - 1))
1453                    != tiling.tile_id(ctb_addr_ts));
1454        if tile_start {
1455            fresh_contexts()
1456        } else if wpp_row_start {
1457            match (&*wpp_stored, t_available(ctb_addr_ts, slice_addr_of)) {
1458                (Some(stored), true) => stored.clone(),
1459                _ => fresh_contexts(),
1460            }
1461        } else {
1462            ds_stored.clone().ok_or(SequenceError::Malformed(
1463                "dependent segment without Ds state",
1464            ))?
1465        }
1466    } else {
1467        fresh_contexts()
1468    };
1469    let mut first_ctu = true;
1470    // Set after the row-final CTU's end_of_subset_one_bit: the next CTU
1471    // starts a new substream.
1472    let mut advance_substream = false;
1473
1474    loop {
1475        if (ctb_addr_ts as usize) >= pic_size_in_ctbs {
1476            return Err(SequenceError::Malformed(
1477                "slice segment runs past the last CTB of the picture",
1478            ));
1479        }
1480        let ctb_addr_rs = tiling.ctb_addr_ts_to_rs(ctb_addr_ts);
1481        let rx = ctb_addr_rs % geom.pic_w_ctbs;
1482        let ry = ctb_addr_rs / geom.pic_w_ctbs;
1483        let x_ctb = rx << geom.ctb_log2;
1484        let y_ctb = ry << geom.ctb_log2;
1485        slice_addr_of[ctb_addr_rs as usize] = Some(slice_addr_rs);
1486
1487        // §9.3.1 / §9.3.2.1 — subset-boundary context handling inside
1488        // one slice segment. Item 2: the first CTU of a tile
1489        // re-initializes the context variables (§9.3.2.2). Item 3
1490        // (WPP): the first luma CTB of a CTU row of a tile either
1491        // synchronizes from the stored above-right state (§9.3.2.5) or
1492        // re-initializes (§9.3.2.2). Either way the next entry-point
1493        // substream starts here. The first CTU of the segment keeps
1494        // the §9.3.1-item-1 slice initialization done above.
1495        if !first_ctu {
1496            let tile_start =
1497                tiles_on && tiling.tile_id(ctb_addr_ts) != tiling.tile_id(ctb_addr_ts - 1);
1498            // §9.3.2.1: CtbAddrInRs % PicWidthInCtbsY == 0, or the
1499            // raster-left neighbour lies in a different tile.
1500            let wpp_row_start = wpp
1501                && !tile_start
1502                && (rx == 0
1503                    || tiling.tile_id(tiling.ctb_addr_rs_to_ts(ctb_addr_rs - 1))
1504                        != tiling.tile_id(ctb_addr_ts));
1505            if tile_start || wpp_row_start {
1506                if !advance_substream {
1507                    return Err(SequenceError::Malformed(
1508                        "subset start without end_of_subset_one_bit",
1509                    ));
1510                }
1511                sub_idx += 1;
1512                engine = CabacEngine::new(BitReader::new(sub_range(sub_idx)?))
1513                    .map_err(|_| SequenceError::Malformed("substream too short for CABAC init"))?;
1514                if tile_start {
1515                    // §9.3.2.2 / §9.3.2.3 — fresh contexts (and
1516                    // re-initialized palette predictor) at the tile
1517                    // start.
1518                    ctx = fresh_contexts();
1519                } else {
1520                    // Spatial neighbour T = the CTB at ( x0 + CtbSizeY,
1521                    // y0 − CtbSizeY ) (eq. 9-3), §6.4.1-gated.
1522                    ctx = match (&*wpp_stored, t_available(ctb_addr_ts, slice_addr_of)) {
1523                        (Some(stored), true) => stored.clone(),
1524                        _ => fresh_contexts(),
1525                    };
1526                }
1527            }
1528        }
1529        advance_substream = false;
1530        first_ctu = false;
1531
1532        // §7.3.8.3 SAO merge-candidate availability: the left / above
1533        // CTB must exist, lie in the same slice segment sequence
1534        // (same SliceAddrRs) and the same tile.
1535        let tile_here = tiling.tile_id(ctb_addr_ts);
1536        let merge_left = rx > 0 && {
1537            let left_rs = ctb_addr_rs - 1;
1538            slice_addr_of[left_rs as usize] == Some(slice_addr_rs)
1539                && tiling.tile_id(tiling.ctb_addr_rs_to_ts(left_rs)) == tile_here
1540        };
1541        let merge_up = ry > 0 && {
1542            let up_rs = ctb_addr_rs - geom.pic_w_ctbs;
1543            slice_addr_of[up_rs as usize] == Some(slice_addr_rs)
1544                && tiling.tile_id(tiling.ctb_addr_rs_to_ts(up_rs)) == tile_here
1545        };
1546
1547        let ctu = decode_coding_tree_unit_in_picture(
1548            &mut engine,
1549            &mut ctx,
1550            &params,
1551            state,
1552            x_ctb,
1553            y_ctb,
1554            slice_addr_rs,
1555            tile_here,
1556            merge_left,
1557            merge_up,
1558        )?;
1559        decoded.push((x_ctb, y_ctb, ctu));
1560
1561        // §9.3.1 / §9.3.2.4 — store the context state after the SECOND
1562        // CTB of a CTU row of a tile: CtbAddrInRs % PicWidthInCtbsY
1563        // == 1, or CtbAddrInRs > 1 and the CTB two to the raster-left
1564        // lies in a different tile.
1565        if wpp
1566            && (rx == 1
1567                || (ctb_addr_rs > 1
1568                    && tiles_on
1569                    && tiling.tile_id(ctb_addr_ts)
1570                        != tiling.tile_id(tiling.ctb_addr_rs_to_ts(ctb_addr_rs - 2))))
1571        {
1572            *wpp_stored = Some(ctx.clone());
1573        }
1574
1575        let eos = end_of_slice_segment_flag(&mut engine)
1576            .map_err(|_| SequenceError::Malformed("CABAC underrun at end_of_slice_segment"))?;
1577        ctb_addr_ts += 1;
1578        if eos {
1579            // §9.3.1 / §9.3.2.4 — store the context variables into
1580            // TableStateIdxDs for a following dependent slice segment.
1581            if pps.dependent_slice_segments_enabled_flag {
1582                *ds_stored = Some(ctx.clone());
1583            }
1584            break;
1585        }
1586        if (ctb_addr_ts as usize) >= pic_size_in_ctbs {
1587            if tolerant {
1588                eprintln!("(tolerant: end_of_slice_segment_flag not set on the last CTB)");
1589                break;
1590            }
1591            return Err(SequenceError::Malformed(
1592                "end_of_slice_segment_flag not set on the last CTB",
1593            ));
1594        }
1595        // §7.3.8.1 — end_of_subset_one_bit + byte_alignment( ) when
1596        // the NEXT CTB (CtbAddrInTs already incremented) starts a new
1597        // tile, or (WPP) a new CTU row of a tile; the next CTU reads
1598        // from the following substream.
1599        let next_rs = tiling.ctb_addr_ts_to_rs(ctb_addr_ts);
1600        let tile_boundary =
1601            tiles_on && tiling.tile_id(ctb_addr_ts) != tiling.tile_id(ctb_addr_ts - 1);
1602        let wpp_boundary = wpp
1603            && (next_rs % geom.pic_w_ctbs == 0
1604                || tiling.tile_id(ctb_addr_ts)
1605                    != tiling.tile_id(tiling.ctb_addr_rs_to_ts(next_rs - 1)));
1606        if tile_boundary || wpp_boundary {
1607            let one = end_of_slice_segment_flag(&mut engine)
1608                .map_err(|_| SequenceError::Malformed("CABAC underrun at end_of_subset_one_bit"))?;
1609            if !one && !tolerant {
1610                return Err(SequenceError::Malformed("end_of_subset_one_bit not set"));
1611            }
1612            advance_substream = true;
1613        }
1614    }
1615    Ok(())
1616}
1617
1618/// §7.4.7.1 — the stripped-RBSP byte ranges of the
1619/// `num_entry_point_offsets + 1` subsets of `slice_segment_data( )`.
1620///
1621/// The wire offsets count coded (escaped) bytes from the first byte of
1622/// the slice segment data, so walk the escaped payload with the
1623/// §7.4.1.1 emulation state machine and translate each boundary into
1624/// the stripped-RBSP index space.
1625fn split_substreams(
1626    escaped: &[u8],
1627    rbsp_len: usize,
1628    stripped_data_offset: usize,
1629    entry_points: Option<&crate::slice::EntryPointOffsets>,
1630) -> Result<Vec<(usize, usize)>, SequenceError> {
1631    let n_offsets = entry_points.map_or(0, |e| e.entry_point_offset_minus1.len());
1632    if n_offsets == 0 {
1633        return Ok(vec![(stripped_data_offset, rbsp_len)]);
1634    }
1635    // stripped index -> escaped index of the slice-data start.
1636    let mut stripped_of_escaped = vec![0usize; escaped.len() + 1];
1637    let mut zeros = 0u32;
1638    let mut stripped = 0usize;
1639    for (i, &b) in escaped.iter().enumerate() {
1640        stripped_of_escaped[i] = stripped;
1641        if zeros >= 2 && b == 0x03 {
1642            // Emulation-prevention byte: consumed, not emitted.
1643            zeros = 0;
1644            continue;
1645        }
1646        if b == 0 {
1647            zeros += 1;
1648        } else {
1649            zeros = 0;
1650        }
1651        stripped += 1;
1652    }
1653    stripped_of_escaped[escaped.len()] = stripped;
1654    // Escaped index of the slice-data start.
1655    let escaped_start = stripped_of_escaped
1656        .iter()
1657        .position(|&sidx| sidx == stripped_data_offset)
1658        .ok_or(SequenceError::Malformed("slice data offset unmappable"))?;
1659
1660    let entry_points = entry_points.expect("checked above");
1661    let mut ranges = Vec::with_capacity(n_offsets + 1);
1662    let mut first_escaped = escaped_start;
1663    let mut first_stripped = stripped_data_offset;
1664    for &off_m1 in &entry_points.entry_point_offset_minus1 {
1665        let len = off_m1 as usize + 1;
1666        let last_escaped = first_escaped
1667            .checked_add(len)
1668            .filter(|&e| e <= escaped.len())
1669            .ok_or(SequenceError::Malformed("entry point past slice data"))?;
1670        let last_stripped = stripped_of_escaped[last_escaped];
1671        ranges.push((first_stripped, last_stripped));
1672        first_escaped = last_escaped;
1673        first_stripped = last_stripped;
1674    }
1675    ranges.push((first_stripped, rbsp_len));
1676    Ok(ranges)
1677}