Skip to main content

h264_reader/nal/
pps.rs

1use super::sps::{self};
2use crate::nal::sps::{
3    ChromaInfo, ScalingLists4x4, ScalingLists8x8, ScalingLists8x8Resolved, SeqParameterSet,
4};
5use crate::nal::sps::{SeqParamSetId, SeqParamSetIdError};
6use crate::rbsp::BitRead;
7use crate::{rbsp, Context};
8
9#[derive(Debug)]
10pub enum PpsError {
11    RbspReaderError(rbsp::BitReaderError),
12    InvalidSliceGroupMapType(u32),
13    InvalidNumSliceGroupsMinus1(u32),
14    InvalidNumRefIdx(&'static str, u32),
15    InvalidSliceGroupChangeType(u32),
16    UnknownSeqParamSetId(SeqParamSetId),
17    BadPicParamSetId(PicParamSetIdError),
18    BadSeqParamSetId(SeqParamSetIdError),
19    ScalingMatrix(sps::ScalingMatrixError),
20    InvalidSecondChromaQpIndexOffset(i32),
21    InvalidPicInitQpMinus26(i32),
22    InvalidPicInitQsMinus26(i32),
23    InvalidChromaQpIndexOffset(i32),
24    InvalidRunLengthMinus1(u32),
25    InvalidTopLeft(u32),
26    InvalidBottomRight(u32),
27    InvalidSliceGroupChangeRateMinus1(u32),
28    InvalidPicSizeInMapUnitsMinus1 { actual: u32, max: u32 },
29}
30
31impl From<rbsp::BitReaderError> for PpsError {
32    fn from(e: rbsp::BitReaderError) -> Self {
33        PpsError::RbspReaderError(e)
34    }
35}
36
37#[derive(Debug, Clone)]
38pub enum SliceGroupChangeType {
39    BoxOut,
40    RasterScan,
41    WipeOut,
42}
43impl SliceGroupChangeType {
44    fn from_id(id: u32) -> Result<SliceGroupChangeType, PpsError> {
45        match id {
46            3 => Ok(SliceGroupChangeType::BoxOut),
47            4 => Ok(SliceGroupChangeType::RasterScan),
48            5 => Ok(SliceGroupChangeType::WipeOut),
49            _ => Err(PpsError::InvalidSliceGroupChangeType(id)),
50        }
51    }
52}
53
54#[derive(Debug, Clone)]
55pub struct SliceRect {
56    top_left: u32,
57    bottom_right: u32,
58}
59impl SliceRect {
60    fn read<R: BitRead>(r: &mut R, sps: &SeqParameterSet) -> Result<SliceRect, PpsError> {
61        let rect = SliceRect {
62            top_left: r.read_ue("top_left")?,
63            bottom_right: r.read_ue("bottom_right")?,
64        };
65        if rect.top_left > rect.bottom_right {
66            return Err(PpsError::InvalidTopLeft(rect.top_left));
67        }
68        if rect.bottom_right > sps.pic_size_in_map_units() {
69            return Err(PpsError::InvalidBottomRight(rect.bottom_right));
70        }
71        if rect.top_left % sps.pic_width_in_mbs() > rect.bottom_right % sps.pic_width_in_mbs() {
72            return Err(PpsError::InvalidTopLeft(rect.top_left));
73        }
74        Ok(rect)
75    }
76}
77
78#[derive(Debug, Clone)]
79pub enum SliceGroup {
80    Interleaved {
81        run_length_minus1: Vec<u32>,
82    },
83    Dispersed {
84        num_slice_groups_minus1: u32,
85    },
86    ForegroundAndLeftover {
87        rectangles: Vec<SliceRect>,
88    },
89    Changing {
90        change_type: SliceGroupChangeType,
91        num_slice_groups_minus1: u32,
92        slice_group_change_direction_flag: bool,
93        slice_group_change_rate_minus1: u32,
94    },
95    ExplicitAssignment {
96        num_slice_groups_minus1: u32,
97        slice_group_id: Vec<u32>,
98    },
99}
100impl SliceGroup {
101    fn read<R: BitRead>(
102        r: &mut R,
103        num_slice_groups_minus1: u32,
104        sps: &SeqParameterSet,
105    ) -> Result<SliceGroup, PpsError> {
106        let slice_group_map_type = r.read_ue("slice_group_map_type")?;
107        match slice_group_map_type {
108            0 => Ok(SliceGroup::Interleaved {
109                run_length_minus1: Self::read_run_lengths(r, num_slice_groups_minus1, sps)?,
110            }),
111            1 => Ok(SliceGroup::Dispersed {
112                num_slice_groups_minus1,
113            }),
114            2 => Ok(SliceGroup::ForegroundAndLeftover {
115                rectangles: Self::read_rectangles(r, num_slice_groups_minus1, sps)?,
116            }),
117            3 | 4 | 5 => {
118                let slice_group_change_direction_flag =
119                    r.read_bit("slice_group_change_direction_flag")?;
120                let slice_group_change_rate_minus1 = r.read_ue("slice_group_change_rate_minus1")?;
121                if slice_group_change_rate_minus1 > sps.pic_size_in_map_units() - 1 {
122                    return Err(PpsError::InvalidSliceGroupChangeRateMinus1(
123                        slice_group_change_rate_minus1,
124                    ));
125                }
126                Ok(SliceGroup::Changing {
127                    change_type: SliceGroupChangeType::from_id(slice_group_map_type)?,
128                    num_slice_groups_minus1,
129                    slice_group_change_direction_flag,
130                    slice_group_change_rate_minus1,
131                })
132            }
133            6 => Ok(SliceGroup::ExplicitAssignment {
134                num_slice_groups_minus1,
135                slice_group_id: Self::read_group_ids(r, num_slice_groups_minus1, sps)?,
136            }),
137            _ => Err(PpsError::InvalidSliceGroupMapType(slice_group_map_type)),
138        }
139    }
140
141    fn read_run_lengths<R: BitRead>(
142        r: &mut R,
143        num_slice_groups_minus1: u32,
144        sps: &SeqParameterSet,
145    ) -> Result<Vec<u32>, PpsError> {
146        let mut run_lengths = Vec::with_capacity(num_slice_groups_minus1 as usize + 1);
147        for _ in 0..num_slice_groups_minus1 + 1 {
148            let run_length_minus1 = r.read_ue("run_length_minus1")?;
149            if run_length_minus1 > sps.pic_size_in_map_units() - 1 {
150                return Err(PpsError::InvalidRunLengthMinus1(run_length_minus1));
151            }
152            run_lengths.push(run_length_minus1);
153        }
154        Ok(run_lengths)
155    }
156
157    // The spec has:
158    //
159    // else if( slice_group_map_type == 2 )
160    //    for( iGroup = 0; iGroup < num_slice_groups_minus1; iGroup++ ) {
161    //      top_left[ iGroup ]
162    //       bottom_right[ iGroup ]
163    //    }
164    //
165    fn read_rectangles<R: BitRead>(
166        r: &mut R,
167        num_slice_groups_minus1: u32,
168        seq_parameter_set: &SeqParameterSet,
169    ) -> Result<Vec<SliceRect>, PpsError> {
170        let mut run_length_minus1 = Vec::with_capacity(num_slice_groups_minus1 as usize);
171        for _ in 0..num_slice_groups_minus1 {
172            run_length_minus1.push(SliceRect::read(r, seq_parameter_set)?);
173        }
174        Ok(run_length_minus1)
175    }
176
177    fn read_group_ids<R: BitRead>(
178        r: &mut R,
179        num_slice_groups_minus1: u32,
180        sps: &SeqParameterSet,
181    ) -> Result<Vec<u32>, PpsError> {
182        let pic_size_in_map_units_minus1 = r.read_ue("pic_size_in_map_units_minus1")?;
183        let pic_size_in_map_units = sps.pic_size_in_map_units();
184        if pic_size_in_map_units_minus1 >= pic_size_in_map_units {
185            return Err(PpsError::InvalidPicSizeInMapUnitsMinus1 {
186                actual: pic_size_in_map_units_minus1,
187                max: pic_size_in_map_units,
188            });
189        }
190        let size = (1f64 + f64::from(num_slice_groups_minus1)).log2().ceil() as u32;
191        let mut slice_group_id = Vec::with_capacity(pic_size_in_map_units_minus1 as usize + 1);
192        for _ in 0..=pic_size_in_map_units_minus1 {
193            slice_group_id.push(r.read_var(size, "slice_group_id")?);
194        }
195        Ok(slice_group_id)
196    }
197}
198
199#[derive(Debug, Clone)]
200pub struct PicScalingMatrix {
201    pub scaling_lists4x4: ScalingLists4x4,
202    /// `Some` when `transform_8x8_mode_flag` is `true`, `None` otherwise
203    pub scaling_lists8x8: Option<ScalingLists8x8>,
204}
205impl PicScalingMatrix {
206    /// Resolve 4x4 scaling lists using PPS fall-back rules (Table 7-2).
207    ///
208    /// Uses rule set B when the SPS has a scaling matrix
209    /// (`seq_scaling_matrix_present_flag` = 1), otherwise rule set A.
210    pub fn scaling_lists_4x4(&self, sps_chroma: &ChromaInfo) -> [[u8; 16]; 6] {
211        let sps_resolved = sps_chroma
212            .scaling_matrix
213            .as_ref()
214            .map(|m| m.scaling_lists_4x4());
215        sps::resolve_4x4_lists(&self.scaling_lists4x4.0, sps_resolved.as_ref())
216    }
217
218    /// Resolve 8x8 scaling lists using PPS fall-back rules (Table 7-2).
219    ///
220    /// Returns `None` when `transform_8x8_mode_flag` is false (no 8x8 lists in PPS).
221    pub fn scaling_lists_8x8(&self, sps_chroma: &ChromaInfo) -> Option<ScalingLists8x8Resolved> {
222        let pps_lists = self.scaling_lists8x8.as_ref()?;
223        let sps_8x8 = sps_chroma
224            .scaling_matrix
225            .as_ref()
226            .map(|m| m.scaling_lists_8x8());
227        let sps_slice = sps_8x8.as_ref().map(|r| r.as_slice());
228
229        Some(match pps_lists {
230            ScalingLists8x8::Y(lists) => {
231                ScalingLists8x8Resolved::Y(sps::resolve_8x8_lists(lists, sps_slice))
232            }
233            ScalingLists8x8::YCbCr(lists) => {
234                ScalingLists8x8Resolved::YCbCr(sps::resolve_8x8_lists(lists, sps_slice))
235            }
236        })
237    }
238
239    fn read<R: BitRead>(
240        r: &mut R,
241        sps: &sps::SeqParameterSet,
242        transform_8x8_mode_flag: bool,
243    ) -> Result<Option<Box<PicScalingMatrix>>, PpsError> {
244        let pic_scaling_matrix_present_flag = r.read_bit("pic_scaling_matrix_present_flag")?;
245
246        if !pic_scaling_matrix_present_flag {
247            return Ok(None);
248        }
249
250        let scaling_lists4x4 = ScalingLists4x4::read(r).map_err(PpsError::ScalingMatrix)?;
251        let scaling_lists8x8 = transform_8x8_mode_flag
252            .then(|| {
253                ScalingLists8x8::read(r, sps.chroma_info.chroma_format)
254                    .map_err(PpsError::ScalingMatrix)
255            })
256            .transpose()?;
257        Ok(Some(Box::new(PicScalingMatrix {
258            scaling_lists4x4,
259            scaling_lists8x8,
260        })))
261    }
262}
263
264#[derive(Debug, Clone)]
265pub struct PicParameterSetExtra {
266    pub transform_8x8_mode_flag: bool,
267    pub pic_scaling_matrix: Option<Box<PicScalingMatrix>>,
268    pub second_chroma_qp_index_offset: i32,
269}
270impl PicParameterSetExtra {
271    fn read<R: BitRead>(
272        r: &mut R,
273        sps: &sps::SeqParameterSet,
274    ) -> Result<Option<PicParameterSetExtra>, PpsError> {
275        Ok(if r.has_more_rbsp_data("transform_8x8_mode_flag")? {
276            let transform_8x8_mode_flag = r.read_bit("transform_8x8_mode_flag")?;
277            let extra = PicParameterSetExtra {
278                transform_8x8_mode_flag,
279                pic_scaling_matrix: PicScalingMatrix::read(r, sps, transform_8x8_mode_flag)?,
280                second_chroma_qp_index_offset: r.read_se("second_chroma_qp_index_offset")?,
281            };
282            if extra.second_chroma_qp_index_offset < -12 || extra.second_chroma_qp_index_offset > 12
283            {
284                return Err(PpsError::InvalidSecondChromaQpIndexOffset(
285                    extra.second_chroma_qp_index_offset,
286                ));
287            }
288            Some(extra)
289        } else {
290            None
291        })
292    }
293}
294
295#[derive(Debug, PartialEq)]
296pub enum PicParamSetIdError {
297    IdTooLarge(u32),
298}
299
300#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
301pub struct PicParamSetId(u8);
302impl PicParamSetId {
303    pub fn from_u32(id: u32) -> Result<PicParamSetId, PicParamSetIdError> {
304        if id > 255 {
305            Err(PicParamSetIdError::IdTooLarge(id))
306        } else {
307            Ok(PicParamSetId(id as u8))
308        }
309    }
310    pub fn id(self) -> u8 {
311        self.0
312    }
313}
314
315#[derive(Clone, Debug)]
316pub struct PicParameterSet {
317    pub pic_parameter_set_id: PicParamSetId,
318    pub seq_parameter_set_id: SeqParamSetId,
319    pub entropy_coding_mode_flag: bool,
320    pub bottom_field_pic_order_in_frame_present_flag: bool,
321    pub slice_groups: Option<SliceGroup>,
322    pub num_ref_idx_l0_default_active_minus1: u32,
323    pub num_ref_idx_l1_default_active_minus1: u32,
324    pub weighted_pred_flag: bool,
325    pub weighted_bipred_idc: u8,
326    pub pic_init_qp_minus26: i32,
327    pub pic_init_qs_minus26: i32,
328    pub chroma_qp_index_offset: i32,
329    pub deblocking_filter_control_present_flag: bool,
330    pub constrained_intra_pred_flag: bool,
331    pub redundant_pic_cnt_present_flag: bool,
332    pub extension: Option<PicParameterSetExtra>,
333}
334impl PicParameterSet {
335    pub fn from_bits<R: BitRead>(ctx: &Context, mut r: R) -> Result<PicParameterSet, PpsError> {
336        let pic_parameter_set_id = PicParamSetId::from_u32(r.read_ue("pic_parameter_set_id")?)
337            .map_err(PpsError::BadPicParamSetId)?;
338        let seq_parameter_set_id = SeqParamSetId::from_u32(r.read_ue("seq_parameter_set_id")?)
339            .map_err(PpsError::BadSeqParamSetId)?;
340        let seq_parameter_set = ctx
341            .sps_by_id(seq_parameter_set_id)
342            .ok_or_else(|| PpsError::UnknownSeqParamSetId(seq_parameter_set_id))?;
343        let pps = PicParameterSet {
344            pic_parameter_set_id,
345            seq_parameter_set_id,
346            entropy_coding_mode_flag: r.read_bit("entropy_coding_mode_flag")?,
347            bottom_field_pic_order_in_frame_present_flag: r
348                .read_bit("bottom_field_pic_order_in_frame_present_flag")?,
349            slice_groups: Self::read_slice_groups(&mut r, seq_parameter_set)?,
350            num_ref_idx_l0_default_active_minus1: read_num_ref_idx(
351                &mut r,
352                "num_ref_idx_l0_default_active_minus1",
353            )?,
354            num_ref_idx_l1_default_active_minus1: read_num_ref_idx(
355                &mut r,
356                "num_ref_idx_l1_default_active_minus1",
357            )?,
358            weighted_pred_flag: r.read_bit("weighted_pred_flag")?,
359            weighted_bipred_idc: r.read::<2, _>("weighted_bipred_idc")?,
360            pic_init_qp_minus26: r.read_se("pic_init_qp_minus26")?,
361            pic_init_qs_minus26: r.read_se("pic_init_qs_minus26")?,
362            chroma_qp_index_offset: r.read_se("chroma_qp_index_offset")?,
363            deblocking_filter_control_present_flag: r
364                .read_bit("deblocking_filter_control_present_flag")?,
365            constrained_intra_pred_flag: r.read_bit("constrained_intra_pred_flag")?,
366            redundant_pic_cnt_present_flag: r.read_bit("redundant_pic_cnt_present_flag")?,
367            extension: PicParameterSetExtra::read(&mut r, seq_parameter_set)?,
368        };
369        let qp_bd_offset_y = 6 * seq_parameter_set.chroma_info.bit_depth_luma_minus8;
370        if pps.pic_init_qp_minus26 < -(26 + i32::from(qp_bd_offset_y))
371            || pps.pic_init_qp_minus26 > 25
372        {
373            return Err(PpsError::InvalidPicInitQpMinus26(pps.pic_init_qp_minus26));
374        }
375        if pps.pic_init_qs_minus26 < -26 || pps.pic_init_qs_minus26 > 25 {
376            return Err(PpsError::InvalidPicInitQsMinus26(pps.pic_init_qs_minus26));
377        }
378        if pps.chroma_qp_index_offset < -12 || pps.chroma_qp_index_offset > 12 {
379            return Err(PpsError::InvalidChromaQpIndexOffset(
380                pps.chroma_qp_index_offset,
381            ));
382        }
383        r.finish_rbsp()?;
384        Ok(pps)
385    }
386
387    fn read_slice_groups<R: BitRead>(
388        r: &mut R,
389        sps: &SeqParameterSet,
390    ) -> Result<Option<SliceGroup>, PpsError> {
391        let num_slice_groups_minus1 = r.read_ue("num_slice_groups_minus1")?;
392        if num_slice_groups_minus1 > 7 {
393            // 7 is the maximum allowed in any profile; some profiles restrict it to 0.
394            return Err(PpsError::InvalidNumSliceGroupsMinus1(
395                num_slice_groups_minus1,
396            ));
397        }
398        Ok(if num_slice_groups_minus1 > 0 {
399            Some(SliceGroup::read(r, num_slice_groups_minus1, sps)?)
400        } else {
401            None
402        })
403    }
404}
405
406fn read_num_ref_idx<R: BitRead>(r: &mut R, name: &'static str) -> Result<u32, PpsError> {
407    let val = r.read_ue(name)?;
408    if val > 31 {
409        return Err(PpsError::InvalidNumRefIdx(name, val));
410    }
411    Ok(val)
412}
413
414#[cfg(test)]
415mod test {
416    use super::*;
417    use crate::nal::sps::SeqParameterSet;
418    use hex_literal::*;
419
420    #[test]
421    fn test_it() {
422        let data = hex!(
423            "64 00 0A AC 72 84 44 26 84 00 00
424            00 04 00 00 00 CA 3C 48 96 11 80"
425        );
426        let sps = super::sps::SeqParameterSet::from_bits(rbsp::BitReader::new(&data[..]))
427            .expect("unexpected test data");
428        let mut ctx = Context::default();
429        ctx.put_seq_param_set(sps);
430        let data = hex!("E8 43 8F 13 21 30");
431        match PicParameterSet::from_bits(&ctx, rbsp::BitReader::new(&data[..])) {
432            Err(e) => panic!("failed: {:?}", e),
433            Ok(pps) => {
434                println!("pps: {:#?}", pps);
435                assert_eq!(pps.pic_parameter_set_id.id(), 0);
436                assert_eq!(pps.seq_parameter_set_id.id(), 0);
437            }
438        }
439    }
440
441    #[test]
442    fn test_transform_8x8_mode_with_scaling_matrix() {
443        let sps = hex!(
444            "64 00 29 ac 1b 1a 50 1e 00 89 f9 70 11 00 00 03 e9 00 00 bb 80 e2 60 00 04 c3 7a 00 00
445             72 70 e8 c4 b8 c4 c0 00 09 86 f4 00 00 e4 e1 d1 89 70 f8 e1 85 2c"
446        );
447        let pps = hex!(
448            "ea 8d ce 50 94 8d 18 b2 5a 55 28 4a 46 8c 59 2d 2a 50 c9 1a 31 64 b4 aa 85 48 d2 75 d5
449             25 1d 23 49 d2 7a 23 74 93 7a 49 be 95 da ad d5 3d 7a 6b 54 22 9a 4e 93 d6 ea 9f a4 ee
450             aa fd 6e bf f5 f7"
451        );
452        let sps = super::sps::SeqParameterSet::from_bits(rbsp::BitReader::new(&sps[..]))
453            .expect("unexpected test data");
454        let mut ctx = Context::default();
455        ctx.put_seq_param_set(sps);
456
457        let pps = PicParameterSet::from_bits(&ctx, rbsp::BitReader::new(&pps[..]))
458            .expect("we mis-parsed pic_scaling_matrix when transform_8x8_mode_flag is active");
459
460        // if transform_8x8_mode_flag were false or pic_scaling_matrix were None then we wouldn't
461        // be recreating the required conditions for the test
462        assert!(matches!(
463            pps.extension,
464            Some(PicParameterSetExtra {
465                transform_8x8_mode_flag: true,
466                pic_scaling_matrix: Some(m),
467                ..
468            }) if matches!(m.scaling_lists8x8, Some(ScalingLists8x8::Y(_)))
469        ));
470    }
471
472    // Earlier versions of h264-reader incorrectly limited pic_parameter_set_id to at most 32,
473    // while the spec allows up to 255.  Test that a value over 32 is accepted.
474    #[test]
475    fn pps_id_greater32() {
476        // test SPS/PPS values courtesy of @astraw
477        let sps = hex!("42c01643235010020b3cf00f08846a");
478        let pps = hex!("0448e3c8");
479        let sps = sps::SeqParameterSet::from_bits(rbsp::BitReader::new(&sps[..])).unwrap();
480        let mut ctx = Context::default();
481        ctx.put_seq_param_set(sps);
482
483        let pps = PicParameterSet::from_bits(&ctx, rbsp::BitReader::new(&pps[..])).unwrap();
484
485        assert_eq!(pps.pic_parameter_set_id, PicParamSetId(33));
486    }
487
488    #[test]
489    fn invalid_pic_init_qs_minus26() {
490        let mut ctx = Context::default();
491        let sps = SeqParameterSet::from_bits(rbsp::BitReader::new(
492            &hex!("64 00 0b ac d9 42 4d f8 84")[..],
493        ))
494        .expect("sps");
495        println!("{:#?}", sps);
496        ctx.put_seq_param_set(sps);
497        let pps = PicParameterSet::from_bits(
498            &mut ctx,
499            rbsp::BitReader::new(&hex!("eb e8 02 3b 2c 8b")[..]),
500        );
501        // pic_init_qs_minus26 should be in the range [-26, 25]
502        assert!(matches!(pps, Err(PpsError::InvalidPicInitQsMinus26(-285))));
503    }
504}