Skip to main content

oximedia_codec/
slice_header.rs

1//! H.264/H.265 slice header types and parsing utilities.
2//!
3//! Provides the [`SliceType`] enum, [`SliceHeader`] struct, and
4//! [`SliceHeaderReader`] helper that extracts slice type and frame number
5//! from a raw bitstream prefix.
6
7#![allow(dead_code)]
8
9/// H.264 slice types as defined in Table 7-6.
10///
11/// The spec defines values 0–9; values 5–9 are equivalent to 0–4 but
12/// signal that all slices in the picture are the same type.
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum SliceType {
15    /// Predictive slice (inter-predicted from reference frames).
16    P,
17    /// Bi-predictive slice (two reference frames).
18    B,
19    /// Intra-only slice.
20    I,
21    /// Switching P slice.
22    Sp,
23    /// Switching I slice.
24    Si,
25}
26
27impl SliceType {
28    /// Returns `true` when this slice type uses only intra prediction.
29    pub fn is_intra(self) -> bool {
30        matches!(self, Self::I | Self::Si)
31    }
32
33    /// Returns `true` when this is a bi-predictive slice.
34    pub fn is_bipredictive(self) -> bool {
35        self == Self::B
36    }
37
38    /// Returns `true` when this is a switching slice (SP or SI).
39    pub fn is_switching(self) -> bool {
40        matches!(self, Self::Sp | Self::Si)
41    }
42
43    /// Decodes a raw ue(v) slice_type value (0–9) into a [`SliceType`].
44    ///
45    /// Values 5–9 are mapped to their equivalents 0–4.
46    pub fn from_raw(raw: u8) -> Option<Self> {
47        match raw % 5 {
48            0 => Some(Self::P),
49            1 => Some(Self::B),
50            2 => Some(Self::I),
51            3 => Some(Self::Sp),
52            4 => Some(Self::Si),
53            _ => None,
54        }
55    }
56}
57
58/// Parsed fields from a slice header.
59#[derive(Debug, Clone)]
60pub struct SliceHeader {
61    /// Slice type.
62    pub slice_type: SliceType,
63    /// `frame_num` syntax element (wraps at `MaxFrameNum`).
64    pub frame_num: u16,
65    /// PPS id referenced by this slice.
66    pub pic_parameter_set_id: u8,
67    /// `field_pic_flag`: `true` when this slice belongs to a field (not a frame).
68    pub field_pic_flag: bool,
69    /// `bottom_field_flag`: relevant only when `field_pic_flag` is `true`.
70    pub bottom_field_flag: bool,
71    /// IDR picture id (only meaningful for IDR slices).
72    pub idr_pic_id: Option<u16>,
73    /// `nal_ref_idc` from the enclosing NAL unit (0 = non-reference).
74    pub nal_ref_idc: u8,
75}
76
77impl SliceHeader {
78    /// Returns `true` if this slice is a reference slice (nal_ref_idc > 0).
79    pub fn is_reference(&self) -> bool {
80        self.nal_ref_idc > 0
81    }
82
83    /// Returns `true` if this slice is an IDR slice.
84    pub fn is_idr(&self) -> bool {
85        self.idr_pic_id.is_some()
86    }
87
88    /// Returns `true` if this is a key-frame (IDR or I-slice that is reference).
89    pub fn is_keyframe(&self) -> bool {
90        self.is_idr() || (self.slice_type.is_intra() && self.is_reference())
91    }
92}
93
94/// Reads slice header fields from a raw byte buffer.
95///
96/// This is a simplified reader that expects the buffer to begin immediately
97/// after the NAL unit start code and header byte. Full Exp-Golomb parsing
98/// is approximated with a fixed layout suitable for testing.
99#[derive(Debug, Default)]
100pub struct SliceHeaderReader {
101    /// Log2 of `MaxFrameNum` (from SPS); used to mask `frame_num`.
102    pub log2_max_frame_num: u8,
103}
104
105impl SliceHeaderReader {
106    /// Creates a new reader with the given `log2_max_frame_num`.
107    pub fn new(log2_max_frame_num: u8) -> Self {
108        Self { log2_max_frame_num }
109    }
110
111    /// Reads the slice type from the first byte of `data`.
112    ///
113    /// Returns `None` if `data` is empty or the value is out of range.
114    pub fn read_type(&self, data: &[u8]) -> Option<SliceType> {
115        let raw = data.first().copied()?;
116        SliceType::from_raw(raw % 10)
117    }
118
119    /// Reads the `frame_num` from bytes 1–2 of `data` (big-endian u16),
120    /// masked to `(1 << log2_max_frame_num) - 1`.
121    ///
122    /// Returns `None` if `data` has fewer than 3 bytes.
123    pub fn frame_num(&self, data: &[u8]) -> Option<u16> {
124        if data.len() < 3 {
125            return None;
126        }
127        let raw = u16::from_be_bytes([data[1], data[2]]);
128        let mask = if self.log2_max_frame_num >= 16 {
129            u16::MAX
130        } else {
131            ((1u32 << self.log2_max_frame_num) - 1) as u16
132        };
133        Some(raw & mask)
134    }
135
136    /// Full parse: reads slice type, frame_num, and other fixed fields.
137    ///
138    /// Returns `None` if `data` is too short (requires at least 5 bytes).
139    pub fn parse(&self, nal_ref_idc: u8, data: &[u8]) -> Option<SliceHeader> {
140        if data.len() < 5 {
141            return None;
142        }
143        let slice_type = self.read_type(data)?;
144        let frame_num = self.frame_num(data)?;
145        let pic_parameter_set_id = data[3];
146        let flags = data[4];
147        let field_pic_flag = (flags & 0x80) != 0;
148        let bottom_field_flag = (flags & 0x40) != 0;
149        let is_idr = (flags & 0x20) != 0;
150        let idr_pic_id = if is_idr {
151            Some(u16::from_be_bytes([data[3], data[4]]) & 0x0FFF)
152        } else {
153            None
154        };
155
156        Some(SliceHeader {
157            slice_type,
158            frame_num,
159            pic_parameter_set_id,
160            field_pic_flag,
161            bottom_field_flag,
162            idr_pic_id,
163            nal_ref_idc,
164        })
165    }
166}
167
168#[cfg(test)]
169mod tests {
170    use super::*;
171
172    #[test]
173    fn test_slice_type_is_intra_i() {
174        assert!(SliceType::I.is_intra());
175    }
176
177    #[test]
178    fn test_slice_type_is_intra_si() {
179        assert!(SliceType::Si.is_intra());
180    }
181
182    #[test]
183    fn test_slice_type_p_is_not_intra() {
184        assert!(!SliceType::P.is_intra());
185    }
186
187    #[test]
188    fn test_slice_type_b_is_bipredictive() {
189        assert!(SliceType::B.is_bipredictive());
190    }
191
192    #[test]
193    fn test_slice_type_p_is_not_bipredictive() {
194        assert!(!SliceType::P.is_bipredictive());
195    }
196
197    #[test]
198    fn test_slice_type_is_switching() {
199        assert!(SliceType::Sp.is_switching());
200        assert!(SliceType::Si.is_switching());
201        assert!(!SliceType::I.is_switching());
202    }
203
204    #[test]
205    fn test_slice_type_from_raw_0_to_4() {
206        assert_eq!(SliceType::from_raw(0), Some(SliceType::P));
207        assert_eq!(SliceType::from_raw(1), Some(SliceType::B));
208        assert_eq!(SliceType::from_raw(2), Some(SliceType::I));
209        assert_eq!(SliceType::from_raw(3), Some(SliceType::Sp));
210        assert_eq!(SliceType::from_raw(4), Some(SliceType::Si));
211    }
212
213    #[test]
214    fn test_slice_type_from_raw_5_to_9_maps_same() {
215        assert_eq!(SliceType::from_raw(5), Some(SliceType::P));
216        assert_eq!(SliceType::from_raw(7), Some(SliceType::I));
217        assert_eq!(SliceType::from_raw(9), Some(SliceType::Si));
218    }
219
220    #[test]
221    fn test_slice_header_is_reference() {
222        let hdr = SliceHeader {
223            slice_type: SliceType::I,
224            frame_num: 0,
225            pic_parameter_set_id: 0,
226            field_pic_flag: false,
227            bottom_field_flag: false,
228            idr_pic_id: None,
229            nal_ref_idc: 1,
230        };
231        assert!(hdr.is_reference());
232    }
233
234    #[test]
235    fn test_slice_header_non_reference() {
236        let hdr = SliceHeader {
237            slice_type: SliceType::B,
238            frame_num: 3,
239            pic_parameter_set_id: 0,
240            field_pic_flag: false,
241            bottom_field_flag: false,
242            idr_pic_id: None,
243            nal_ref_idc: 0,
244        };
245        assert!(!hdr.is_reference());
246    }
247
248    #[test]
249    fn test_slice_header_is_idr() {
250        let hdr = SliceHeader {
251            slice_type: SliceType::I,
252            frame_num: 0,
253            pic_parameter_set_id: 0,
254            field_pic_flag: false,
255            bottom_field_flag: false,
256            idr_pic_id: Some(0),
257            nal_ref_idc: 3,
258        };
259        assert!(hdr.is_idr());
260        assert!(hdr.is_keyframe());
261    }
262
263    #[test]
264    fn test_reader_read_type_i_slice() {
265        let reader = SliceHeaderReader::new(4);
266        // raw = 2 → I
267        let ty = reader.read_type(&[2, 0, 0, 0, 0]);
268        assert_eq!(ty, Some(SliceType::I));
269    }
270
271    #[test]
272    fn test_reader_read_type_empty_returns_none() {
273        let reader = SliceHeaderReader::new(4);
274        assert!(reader.read_type(&[]).is_none());
275    }
276
277    #[test]
278    fn test_reader_frame_num_masked() {
279        let reader = SliceHeaderReader::new(4); // mask = 0x000F
280                                                // raw frame_num bytes: 0x01, 0xFF → 0x01FF & 0x000F = 0x000F
281        let fn_ = reader.frame_num(&[0, 0x01, 0xFF, 0, 0]);
282        assert_eq!(fn_, Some(0x000F));
283    }
284
285    #[test]
286    fn test_reader_frame_num_short_returns_none() {
287        let reader = SliceHeaderReader::new(4);
288        assert!(reader.frame_num(&[0, 1]).is_none());
289    }
290
291    #[test]
292    fn test_reader_parse_full() {
293        let reader = SliceHeaderReader::new(8);
294        // slice_type=2 (I), frame_num=0x0005, pps_id=0, flags=0x00
295        let data = [2u8, 0x00, 0x05, 0x00, 0x00];
296        let hdr = reader.parse(1, &data).expect("parse should succeed");
297        assert_eq!(hdr.slice_type, SliceType::I);
298        assert_eq!(hdr.frame_num, 5);
299        assert!(hdr.is_reference());
300    }
301
302    #[test]
303    fn test_reader_parse_too_short_returns_none() {
304        let reader = SliceHeaderReader::new(4);
305        assert!(reader.parse(1, &[2, 0, 5]).is_none());
306    }
307}