1#![allow(dead_code)]
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum SliceType {
15 P,
17 B,
19 I,
21 Sp,
23 Si,
25}
26
27impl SliceType {
28 pub fn is_intra(self) -> bool {
30 matches!(self, Self::I | Self::Si)
31 }
32
33 pub fn is_bipredictive(self) -> bool {
35 self == Self::B
36 }
37
38 pub fn is_switching(self) -> bool {
40 matches!(self, Self::Sp | Self::Si)
41 }
42
43 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#[derive(Debug, Clone)]
60pub struct SliceHeader {
61 pub slice_type: SliceType,
63 pub frame_num: u16,
65 pub pic_parameter_set_id: u8,
67 pub field_pic_flag: bool,
69 pub bottom_field_flag: bool,
71 pub idr_pic_id: Option<u16>,
73 pub nal_ref_idc: u8,
75}
76
77impl SliceHeader {
78 pub fn is_reference(&self) -> bool {
80 self.nal_ref_idc > 0
81 }
82
83 pub fn is_idr(&self) -> bool {
85 self.idr_pic_id.is_some()
86 }
87
88 pub fn is_keyframe(&self) -> bool {
90 self.is_idr() || (self.slice_type.is_intra() && self.is_reference())
91 }
92}
93
94#[derive(Debug, Default)]
100pub struct SliceHeaderReader {
101 pub log2_max_frame_num: u8,
103}
104
105impl SliceHeaderReader {
106 pub fn new(log2_max_frame_num: u8) -> Self {
108 Self { log2_max_frame_num }
109 }
110
111 pub fn read_type(&self, data: &[u8]) -> Option<SliceType> {
115 let raw = data.first().copied()?;
116 SliceType::from_raw(raw % 10)
117 }
118
119 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 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 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); 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 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}