edgefirst_codec/video/
h264.rs

1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! H.264 Annex B byte-stream framing (ITU-T H.264 §7.3.1 and Annex B).
5//!
6//! Encoders hand back Annex B access units: NAL units each preceded by a
7//! `00 00 01` or `00 00 00 01` start code. This module finds the NAL units
8//! and classifies the ones the encoder cares about, without parsing slice
9//! data.
10
11/// `nal_unit_type` of a coded slice of an IDR picture.
12pub(crate) const NAL_IDR: u8 = 5;
13/// `nal_unit_type` of a sequence parameter set.
14pub(crate) const NAL_SPS: u8 = 7;
15/// `nal_unit_type` of a picture parameter set.
16pub(crate) const NAL_PPS: u8 = 8;
17/// `nal_unit_type` of an end of sequence.
18const NAL_END_OF_SEQUENCE: u8 = 10;
19/// `nal_unit_type` of an end of stream.
20const NAL_END_OF_STREAM: u8 = 11;
21/// `nal_unit_type` of filler data.
22const NAL_FILLER: u8 = 12;
23
24/// The `nal_unit_type` values of every NAL unit in an Annex B buffer, in
25/// order. Bytes before the first start code are ignored.
26pub(crate) fn nal_types(data: &[u8]) -> Vec<u8> {
27    let mut types = Vec::new();
28    let mut i = 0;
29    while let Some(start) = next_start_code(data, i) {
30        if let Some(&header) = data.get(start) {
31            types.push(header & 0x1f);
32        }
33        i = start;
34    }
35    types
36}
37
38/// The SPS and PPS NAL units of an Annex B buffer, each with a 4-byte
39/// start code, in stream order.
40pub(crate) fn parameter_sets(data: &[u8]) -> Vec<u8> {
41    let mut starts = Vec::new();
42    let mut i = 0;
43    while let Some(start) = next_start_code(data, i) {
44        starts.push(start);
45        i = start;
46    }
47    let mut out = Vec::new();
48    for (k, &start) in starts.iter().enumerate() {
49        let Some(&header) = data.get(start) else {
50            continue;
51        };
52        if !matches!(header & 0x1f, NAL_SPS | NAL_PPS) {
53            continue;
54        }
55        // The NAL unit ends where the next start code (and any zero byte of a
56        // 4-byte code) begins.
57        let mut end = starts.get(k + 1).map_or(data.len(), |&next| next - 3);
58        while end > start && data[end - 1] == 0 && k + 1 < starts.len() {
59            end -= 1;
60        }
61        out.extend_from_slice(&[0, 0, 0, 1]);
62        out.extend_from_slice(&data[start..end]);
63    }
64    out
65}
66
67/// Index of the first byte after the next `00 00 01` start code at or
68/// after `from` (a 4-byte `00 00 00 01` code ends with the same 3 bytes).
69fn next_start_code(data: &[u8], from: usize) -> Option<usize> {
70    data.get(from..)?
71        .windows(3)
72        .position(|w| w == [0, 0, 1])
73        .map(|p| from + p + 3)
74}
75
76/// Splits an H.264 Annex B byte stream into access units, the input
77/// [`VideoDecoder::decode`](super::VideoDecoder::decode) takes.
78///
79/// An access unit starts at the first NAL unit after a picture that may only
80/// precede one (access unit delimiter, SEI, SPS, PPS; §7.4.1.2.3), or at a
81/// slice that starts a new picture (`first_mb_in_slice` is 0). Streams that
82/// send a picture's slices out of order (arbitrary slice order) are split at
83/// each slice that starts at macroblock 0, which such streams do not
84/// guarantee. Bytes before the first start code are skipped.
85pub fn access_units(stream: &[u8]) -> AccessUnits<'_> {
86    AccessUnits {
87        data: stream,
88        next: next_start_code(stream, 0),
89    }
90}
91
92/// Iterator over the access units of a byte stream, from [`access_units`].
93#[derive(Debug, Clone)]
94pub struct AccessUnits<'a> {
95    data: &'a [u8],
96    /// First byte of the next access unit's first NAL unit header.
97    next: Option<usize>,
98}
99
100impl<'a> Iterator for AccessUnits<'a> {
101    type Item = &'a [u8];
102
103    fn next(&mut self) -> Option<&'a [u8]> {
104        let data = self.data;
105        let first = self.next?;
106        let start = start_code_begin(data, first);
107        let mut picture = false;
108        let mut nal = first;
109        loop {
110            let Some(&header) = data.get(nal) else {
111                self.next = None;
112                return Some(&data[start..]);
113            };
114            let kind = header & 0x1f;
115            let slice = matches!(kind, 1 | 5);
116            // ue(v) first_mb_in_slice is 0 exactly when its first bit is 1.
117            let new_picture = slice && data.get(nal + 1).is_some_and(|b| b & 0x80 != 0);
118            let leads = matches!(kind, 6..=9 | 14..=18);
119            if picture && (leads || new_picture) {
120                self.next = Some(nal);
121                return Some(&data[start..start_code_begin(data, nal)]);
122            }
123            picture |= slice;
124            match next_start_code(data, nal) {
125                Some(n) => nal = n,
126                None => {
127                    self.next = None;
128                    return Some(&data[start..]);
129                }
130            }
131        }
132    }
133}
134
135/// Where the start code before the NAL unit header at `nal` begins,
136/// including the zero byte of a 4-byte code.
137fn start_code_begin(data: &[u8], nal: usize) -> usize {
138    let begin = nal - 3;
139    if begin > 0 && data[begin - 1] == 0 {
140        begin - 1
141    } else {
142        begin
143    }
144}
145
146/// What an encoded buffer contains, by its NAL units.
147#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
148pub(crate) struct AccessUnit {
149    /// Contains a slice of an IDR picture.
150    pub idr: bool,
151    /// Contains an SPS and a PPS.
152    pub parameter_sets: bool,
153    /// Contains at least one coded slice (`nal_unit_type` 1 to 5).
154    pub has_picture: bool,
155    /// Has NAL units, and every one may only follow a picture (end of
156    /// sequence, end of stream, filler data; §7.4.1.2.3).
157    pub trailing_only: bool,
158}
159
160impl AccessUnit {
161    /// Classifies an Annex B buffer.
162    pub(crate) fn parse(data: &[u8]) -> Self {
163        let types = nal_types(data);
164        Self {
165            idr: types.contains(&NAL_IDR),
166            parameter_sets: types.contains(&NAL_SPS) && types.contains(&NAL_PPS),
167            has_picture: types.iter().any(|t| (1..=5).contains(t)),
168            trailing_only: !types.is_empty()
169                && types
170                    .iter()
171                    .all(|t| matches!(*t, NAL_END_OF_SEQUENCE | NAL_END_OF_STREAM | NAL_FILLER)),
172        }
173    }
174
175    /// Belongs to the access unit before it: no picture, only NAL units that
176    /// end one (an encoder's end of sequence after a drain).
177    pub(crate) fn ends_previous(&self) -> bool {
178        !self.has_picture && self.trailing_only
179    }
180}
181
182#[cfg(test)]
183mod tests {
184    use super::*;
185
186    const SPS: &[u8] = &[0, 0, 0, 1, 0x67, 0x64, 0x00, 0x28];
187    const PPS: &[u8] = &[0, 0, 0, 1, 0x68, 0xee, 0x3c, 0x80];
188    const IDR: &[u8] = &[0, 0, 1, 0x65, 0x88, 0x84, 0x00];
189    const P_SLICE: &[u8] = &[0, 0, 0, 1, 0x41, 0x9a, 0x02];
190    const SEI: &[u8] = &[0, 0, 1, 0x06, 0x05, 0xff];
191
192    fn cat(parts: &[&[u8]]) -> Vec<u8> {
193        parts.concat()
194    }
195
196    #[test]
197    fn finds_nal_types_after_3_and_4_byte_start_codes() {
198        let au = cat(&[SPS, PPS, SEI, IDR]);
199        assert_eq!(nal_types(&au), vec![NAL_SPS, NAL_PPS, 6, NAL_IDR]);
200    }
201
202    #[test]
203    fn ignores_leading_garbage_and_empty_input() {
204        assert!(nal_types(&[]).is_empty());
205        assert!(nal_types(&[0, 0]).is_empty());
206        assert_eq!(nal_types(&cat(&[&[0xaa, 0xbb], P_SLICE])), vec![1]);
207    }
208
209    #[test]
210    fn start_code_at_the_very_end_has_no_nal() {
211        assert_eq!(nal_types(&cat(&[P_SLICE, &[0, 0, 1]])), vec![1]);
212    }
213
214    #[test]
215    fn keyframe_with_headers() {
216        let au = AccessUnit::parse(&cat(&[SPS, PPS, IDR]));
217        assert!(au.idr && au.parameter_sets && au.has_picture);
218        assert!(!au.ends_previous());
219    }
220
221    #[test]
222    fn idr_without_headers_is_still_idr() {
223        let au = AccessUnit::parse(IDR);
224        assert!(au.idr && !au.parameter_sets);
225    }
226
227    #[test]
228    fn predicted_frame() {
229        let au = AccessUnit::parse(P_SLICE);
230        assert!(!au.idr && !au.parameter_sets && au.has_picture);
231    }
232
233    #[test]
234    fn end_of_sequence_belongs_to_the_previous_access_unit() {
235        let eos: &[u8] = &[0, 0, 0, 1, 0x0a];
236        let au = AccessUnit::parse(eos);
237        assert!(au.ends_previous() && !au.parameter_sets);
238        assert!(AccessUnit::parse(&cat(&[eos, &[0, 0, 1, 0x0b]])).ends_previous());
239        assert!(!AccessUnit::parse(&cat(&[P_SLICE, eos])).ends_previous());
240        assert!(
241            !AccessUnit::parse(SEI).ends_previous(),
242            "SEI precedes a picture"
243        );
244        assert!(!AccessUnit::parse(&[]).ends_previous());
245    }
246
247    #[test]
248    fn extracts_parameter_sets_with_4_byte_start_codes() {
249        let au = cat(&[SPS, PPS, SEI, IDR]);
250        assert_eq!(parameter_sets(&au), cat(&[SPS, PPS]));
251        // 3-byte start codes come back as 4-byte ones.
252        let short = cat(&[&[0, 0, 1, 0x67, 0x42], &[0, 0, 1, 0x68, 0xce], IDR]);
253        assert_eq!(
254            parameter_sets(&short),
255            vec![0, 0, 0, 1, 0x67, 0x42, 0, 0, 0, 1, 0x68, 0xce]
256        );
257        assert!(parameter_sets(&cat(&[P_SLICE, SEI])).is_empty());
258        assert_eq!(parameter_sets(&cat(&[IDR, PPS])), PPS.to_vec(), "last NAL");
259    }
260
261    #[test]
262    fn splits_access_units_at_headers_and_new_pictures() {
263        // A second slice of the same picture: first_mb_in_slice = 1 (ue "010").
264        let p_slice2: &[u8] = &[0, 0, 1, 0x41, 0x40, 0x11];
265        let aud: &[u8] = &[0, 0, 0, 1, 0x09, 0xf0];
266        let stream = cat(&[SPS, PPS, IDR, P_SLICE, p_slice2, aud, SEI, P_SLICE]);
267        let aus: Vec<&[u8]> = access_units(&stream).collect();
268        assert_eq!(aus.len(), 3);
269        assert_eq!(aus[0], cat(&[SPS, PPS, IDR]).as_slice());
270        assert_eq!(aus[1], cat(&[P_SLICE, p_slice2]).as_slice());
271        assert_eq!(aus[2], cat(&[aud, SEI, P_SLICE]).as_slice());
272        assert_eq!(aus.concat(), stream, "nothing is lost or duplicated");
273    }
274
275    #[test]
276    fn end_of_sequence_stays_with_its_picture() {
277        let eos: &[u8] = &[0, 0, 0, 1, 0x0a];
278        let stream = cat(&[IDR, eos, SPS, PPS, IDR]);
279        let aus: Vec<&[u8]> = access_units(&stream).collect();
280        assert_eq!(aus, vec![cat(&[IDR, eos]), cat(&[SPS, PPS, IDR])]);
281    }
282
283    #[test]
284    fn access_units_of_degenerate_input() {
285        assert_eq!(access_units(&[]).count(), 0);
286        assert_eq!(access_units(&[1, 2, 3]).count(), 0);
287        assert_eq!(
288            access_units(&cat(&[&[0xaa], SPS])).collect::<Vec<_>>(),
289            vec![SPS]
290        );
291        // A start code with no NAL header after it ends the last unit.
292        let tail = cat(&[IDR, &[0, 0, 1]]);
293        assert_eq!(
294            access_units(&tail).collect::<Vec<_>>(),
295            vec![tail.as_slice()]
296        );
297    }
298
299    #[test]
300    fn header_only_buffer() {
301        let au = AccessUnit::parse(&cat(&[SPS, PPS]));
302        assert!(au.parameter_sets && !au.has_picture && !au.ends_previous());
303        assert!(!AccessUnit::parse(SPS).parameter_sets);
304        assert!(!AccessUnit::parse(SEI).has_picture);
305    }
306}
edgefirst_codec/video/h264.rs

edgefirst_codec/video/
h264.rs
