oximedia_codec/
nal_unit.rs

1//! NAL unit handling for H.264/H.265.
2//!
3//! Provides NAL unit type classification, start code detection, RBSP
4//! trailing-bit stripping, and a simple Annex B byte-stream parser.
5
6#![allow(dead_code)]
7
8/// H.264 (AVC) NAL unit type.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum H264NalType {
11    /// Unspecified.
12    Unspecified,
13    /// Non-IDR slice (P- or B-frame).
14    NonIdrSlice,
15    /// Slice data partition A.
16    SlicePartitionA,
17    /// Slice data partition B.
18    SlicePartitionB,
19    /// Slice data partition C.
20    SlicePartitionC,
21    /// IDR slice (keyframe).
22    IdrSlice,
23    /// Supplemental enhancement information.
24    Sei,
25    /// Sequence parameter set.
26    Sps,
27    /// Picture parameter set.
28    Pps,
29    /// Access unit delimiter.
30    AccessUnitDelimiter,
31    /// End of sequence.
32    EndOfSequence,
33    /// End of stream.
34    EndOfStream,
35    /// Filler data.
36    FillerData,
37    /// Reserved or unknown type.
38    Reserved(u8),
39}
40
41impl H264NalType {
42    /// Parses the first byte of a NAL unit header (ignoring forbidden zero bit
43    /// and NRI bits, extracting the 5-bit nal_unit_type).
44    #[must_use]
45    pub fn from_header_byte(byte: u8) -> Self {
46        match byte & 0x1F {
47            0 => Self::Unspecified,
48            1 => Self::NonIdrSlice,
49            2 => Self::SlicePartitionA,
50            3 => Self::SlicePartitionB,
51            4 => Self::SlicePartitionC,
52            5 => Self::IdrSlice,
53            6 => Self::Sei,
54            7 => Self::Sps,
55            8 => Self::Pps,
56            9 => Self::AccessUnitDelimiter,
57            10 => Self::EndOfSequence,
58            11 => Self::EndOfStream,
59            12 => Self::FillerData,
60            n => Self::Reserved(n),
61        }
62    }
63
64    /// Returns `true` if this NAL unit is a slice (IDR or non-IDR).
65    #[must_use]
66    pub fn is_slice(self) -> bool {
67        matches!(self, Self::NonIdrSlice | Self::IdrSlice)
68    }
69
70    /// Returns `true` if this NAL unit is a parameter set (SPS or PPS).
71    #[must_use]
72    pub fn is_parameter_set(self) -> bool {
73        matches!(self, Self::Sps | Self::Pps)
74    }
75}
76
77/// H.265 (HEVC) NAL unit type.
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub enum H265NalType {
80    /// Trailing non-reference picture (B-frame).
81    TrailN,
82    /// Trailing reference picture (P-frame).
83    TrailR,
84    /// Instantaneous decoder refresh picture (IDR).
85    IdrWRadl,
86    /// IDR without leading pictures.
87    IdrNLp,
88    /// Video parameter set.
89    Vps,
90    /// Sequence parameter set.
91    Sps,
92    /// Picture parameter set.
93    Pps,
94    /// Access unit delimiter.
95    Aud,
96    /// Supplemental enhancement information (prefix).
97    PrefixSei,
98    /// Reserved or unknown type.
99    Reserved(u8),
100}
101
102impl H265NalType {
103    /// Parses from the 6-bit nal_unit_type field of an HEVC NAL header.
104    #[must_use]
105    pub fn from_nal_type(nal_type: u8) -> Self {
106        match nal_type & 0x3F {
107            0 => Self::TrailN,
108            1 => Self::TrailR,
109            19 => Self::IdrWRadl,
110            20 => Self::IdrNLp,
111            32 => Self::Vps,
112            33 => Self::Sps,
113            34 => Self::Pps,
114            35 => Self::Aud,
115            39 => Self::PrefixSei,
116            n => Self::Reserved(n),
117        }
118    }
119
120    /// Returns `true` if this is an IDR picture.
121    #[must_use]
122    pub fn is_idr(self) -> bool {
123        matches!(self, Self::IdrWRadl | Self::IdrNLp)
124    }
125}
126
127/// The 3-byte Annex B start code prefix `0x00 0x00 0x01`.
128pub const START_CODE_3: [u8; 3] = [0x00, 0x00, 0x01];
129/// The 4-byte Annex B start code prefix `0x00 0x00 0x00 0x01`.
130pub const START_CODE_4: [u8; 4] = [0x00, 0x00, 0x00, 0x01];
131
132/// Returns `true` if `data` begins with a 3-byte or 4-byte Annex B start code.
133#[must_use]
134pub fn has_start_code(data: &[u8]) -> bool {
135    data.starts_with(&START_CODE_4) || data.starts_with(&START_CODE_3)
136}
137
138/// Removes trailing RBSP stop bit and zero-padding bytes from a raw RBSP
139/// payload, returning the trimmed slice.
140#[must_use]
141pub fn strip_rbsp_trailing(data: &[u8]) -> &[u8] {
142    // Walk backwards, skipping 0x00 bytes until we hit the stop bit (0x80).
143    let mut end = data.len();
144    while end > 0 {
145        end -= 1;
146        let byte = data[end];
147        if byte == 0x80 {
148            return &data[..end];
149        }
150        if byte != 0x00 {
151            // No trailing stop bit found; return original.
152            return data;
153        }
154    }
155    data
156}
157
158/// A parsed NAL unit extracted from an Annex B byte stream.
159#[derive(Debug, Clone)]
160pub struct NalUnit<'a> {
161    /// Raw bytes of this NAL unit (without the leading start code).
162    pub data: &'a [u8],
163    /// Byte offset of the first byte of this NAL unit within the original stream.
164    pub offset: usize,
165}
166
167impl NalUnit<'_> {
168    /// Returns the NAL unit header byte (first byte of `data`).
169    #[must_use]
170    pub fn header_byte(&self) -> Option<u8> {
171        self.data.first().copied()
172    }
173
174    /// Parses the H.264 NAL type from the header byte.
175    #[must_use]
176    pub fn h264_type(&self) -> Option<H264NalType> {
177        self.header_byte().map(H264NalType::from_header_byte)
178    }
179}
180
181/// Parses all NAL units from an Annex B byte stream.
182///
183/// Returns a `Vec` of [`NalUnit`] references into the original `data` slice.
184#[must_use]
185pub fn parse_annex_b(data: &[u8]) -> Vec<NalUnit<'_>> {
186    let mut units = Vec::new();
187    let mut i = 0usize;
188    let len = data.len();
189
190    while i < len {
191        // Detect start code: try 4-byte first, then 3-byte.
192        let start_code_len = if i + 4 <= len && data[i..i + 4] == START_CODE_4 {
193            4
194        } else if i + 3 <= len && data[i..i + 3] == START_CODE_3 {
195            3
196        } else {
197            i += 1;
198            continue;
199        };
200
201        let nal_start = i + start_code_len;
202        // Find the end of this NAL unit (next start code or end of data).
203        let mut j = nal_start;
204        let mut found_next = false;
205        while j + 3 <= len {
206            if data[j..j + 3] == START_CODE_3 || (j + 4 <= len && data[j..j + 4] == START_CODE_4) {
207                found_next = true;
208                break;
209            }
210            j += 1;
211        }
212        if !found_next {
213            j = len;
214        }
215        // Strip trailing zero bytes before the next start code.
216        let mut nal_end = j;
217        while nal_end > nal_start && data[nal_end - 1] == 0x00 {
218            nal_end -= 1;
219        }
220        if nal_end > nal_start {
221            units.push(NalUnit {
222                data: &data[nal_start..nal_end],
223                offset: nal_start,
224            });
225        }
226        i = j;
227    }
228    units
229}
230
231#[cfg(test)]
232mod tests {
233    use super::*;
234
235    #[test]
236    fn test_h264_idr_type() {
237        assert_eq!(H264NalType::from_header_byte(0x65), H264NalType::IdrSlice);
238    }
239
240    #[test]
241    fn test_h264_sps_type() {
242        assert_eq!(H264NalType::from_header_byte(0x67), H264NalType::Sps);
243    }
244
245    #[test]
246    fn test_h264_pps_type() {
247        assert_eq!(H264NalType::from_header_byte(0x68), H264NalType::Pps);
248    }
249
250    #[test]
251    fn test_h264_non_idr_slice() {
252        assert_eq!(
253            H264NalType::from_header_byte(0x41),
254            H264NalType::NonIdrSlice
255        );
256    }
257
258    #[test]
259    fn test_h264_is_slice() {
260        assert!(H264NalType::IdrSlice.is_slice());
261        assert!(H264NalType::NonIdrSlice.is_slice());
262        assert!(!H264NalType::Sps.is_slice());
263    }
264
265    #[test]
266    fn test_h264_is_parameter_set() {
267        assert!(H264NalType::Sps.is_parameter_set());
268        assert!(H264NalType::Pps.is_parameter_set());
269        assert!(!H264NalType::IdrSlice.is_parameter_set());
270    }
271
272    #[test]
273    fn test_h265_idr_type() {
274        assert_eq!(H265NalType::from_nal_type(19), H265NalType::IdrWRadl);
275        assert!(H265NalType::IdrWRadl.is_idr());
276    }
277
278    #[test]
279    fn test_h265_non_idr_not_idr() {
280        assert!(!H265NalType::TrailR.is_idr());
281    }
282
283    #[test]
284    fn test_has_start_code_4byte() {
285        assert!(has_start_code(&[0x00, 0x00, 0x00, 0x01, 0x67]));
286    }
287
288    #[test]
289    fn test_has_start_code_3byte() {
290        assert!(has_start_code(&[0x00, 0x00, 0x01, 0x67]));
291    }
292
293    #[test]
294    fn test_no_start_code() {
295        assert!(!has_start_code(&[0x01, 0x02, 0x03]));
296    }
297
298    #[test]
299    fn test_strip_rbsp_trailing_removes_stop_bit() {
300        let data = &[0xAB, 0xCD, 0x80, 0x00, 0x00];
301        let stripped = strip_rbsp_trailing(data);
302        assert_eq!(stripped, &[0xAB, 0xCD]);
303    }
304
305    #[test]
306    fn test_strip_rbsp_no_stop_bit_returns_original() {
307        let data = &[0x01, 0x02, 0x03];
308        assert_eq!(strip_rbsp_trailing(data), data);
309    }
310
311    #[test]
312    fn test_parse_annex_b_single_nal() {
313        let stream = [0x00, 0x00, 0x00, 0x01, 0x67, 0xAB, 0xCD];
314        let units = parse_annex_b(&stream);
315        assert_eq!(units.len(), 1);
316        assert_eq!(units[0].data[0], 0x67);
317    }
318
319    #[test]
320    fn test_parse_annex_b_multiple_nals() {
321        let stream = [
322            0x00, 0x00, 0x00, 0x01, 0x67, 0x11, 0x00, 0x00, 0x01, 0x68, 0x22,
323        ];
324        let units = parse_annex_b(&stream);
325        assert_eq!(units.len(), 2);
326        assert_eq!(units[0].h264_type(), Some(H264NalType::Sps));
327        assert_eq!(units[1].h264_type(), Some(H264NalType::Pps));
328    }
329}
oximedia_codec/nal_unit.rs

oximedia_codec/
nal_unit.rs
