openh264/
utils.rs

1// How many `0` we have to observe before a `1` means NAL.
2const NAL_MIN_0_COUNT: usize = 2;
3
4/// Given a stream, finds the index of the nth NAL start.
5#[inline]
6fn nth_nal_index(stream: &[u8], nth: usize) -> Option<usize> {
7    let mut count_0 = 0;
8    let mut n = 0;
9
10    for (i, byte) in stream.iter().enumerate() {
11        match byte {
12            0 => count_0 += 1,
13            1 if count_0 >= NAL_MIN_0_COUNT => {
14                if n == nth {
15                    return Some(i - NAL_MIN_0_COUNT);
16                }
17                count_0 = 0;
18                n += 1;
19            }
20            _ => count_0 = 0,
21        }
22    }
23
24    None
25}
26
27/// Splits a bitstream into NAL units.
28///
29/// This function is useful if you happen to have a H.264 bitstream and want to decode it frame by frame: You
30/// apply this function to the underlying stream and run your decoder on each returned slice, preferably
31/// ignoring isolated decoding errors.
32///
33/// In detail, given a bitstream like so (`001` being the NAL start prefix code):
34///
35/// ```text
36/// ......001.........001......001.....
37/// ```
38///
39/// This function will return an iterator returning packets:
40/// ```text
41///      [001.......][001....][001.....]
42/// ```
43///
44/// In other words, any incomplete data at the beginning of the buffer is skipped,
45/// NAL units in the middle are split at their boundaries, the last packet is returned
46/// as-is.
47///
48pub fn nal_units(mut stream: &[u8]) -> impl Iterator<Item = &[u8]> {
49    std::iter::from_fn(move || {
50        let first = nth_nal_index(stream, 0);
51        let next = nth_nal_index(stream, 1);
52
53        match (first, next) {
54            (Some(f), Some(n)) => {
55                let rval = &stream[f..n];
56                stream = &stream[n..];
57                Some(rval)
58            }
59            (Some(f), None) => {
60                let rval = &stream[f..];
61                stream = &stream[f + NAL_MIN_0_COUNT..];
62                Some(rval)
63            }
64            _ => None,
65        }
66    })
67}
68
69/// Splits an incrementally arriving bitstream into NAL units.
70///
71/// This searches for `001` marks in a byte stream, and deals with cross-boundary checks when
72/// a frame is partially read.
73#[derive(Default)]
74pub struct NalParser {
75    leftover_buffer: Vec<u8>,
76    curr_offset: usize,
77    last_nal: Option<usize>,
78}
79
80impl NalParser {
81    /// Creates a new NAL parser.
82    #[must_use]
83    pub fn new() -> Self {
84        Self::default()
85    }
86
87    /// Tries to retrieve the next NAL unit, if present.
88    ///
89    /// After feeding new data you should keep calling this method until it returns `None`.
90    #[allow(clippy::should_implement_trait)]
91    pub fn next(&mut self) -> Option<Vec<u8>> {
92        if self.leftover_buffer.is_empty() {
93            return None;
94        }
95
96        if let Some(idx) = self.get_nal_mark() {
97            if let Some(last_offset) = self.last_nal {
98                // Last mark and current mark found, process packet
99                let packet = self.leftover_buffer[last_offset..idx].to_vec();
100                self.leftover_buffer = self.leftover_buffer[idx..].to_vec();
101                self.last_nal = Some(0);
102                self.curr_offset = 2;
103                Some(packet)
104            } else {
105                // Try your luck searching for 0, 0, 1
106                // In case there is no 0, 0, 1 in the next try, you get ReadMore
107                self.curr_offset = idx + 2;
108                self.last_nal = Some(idx);
109                None
110            }
111        } else {
112            // No 0, 0, 1 mark here, read more data
113            None
114        }
115    }
116
117    /// Feeds more data to the processor.
118    ///
119    /// After calling this method, there may be between 0 to M new NAL units present, which you can query with [`Self::next()`].
120    pub fn feed(&mut self, buffer: impl AsRef<[u8]>) {
121        self.leftover_buffer.extend_from_slice(buffer.as_ref());
122    }
123
124    fn get_nal_mark(&self) -> Option<usize> {
125        (self.curr_offset..self.leftover_buffer.len() - 2)
126            .find(|&i| self.leftover_buffer[i] == 0 && self.leftover_buffer[i + 1] == 0 && self.leftover_buffer[i + 2] == 1)
127    }
128}
129
130#[cfg(test)]
131mod test {
132    use super::{NalParser, nal_units};
133
134    #[test]
135    fn splits_at_nal() {
136        let stream = [];
137        assert!(nal_units(&stream).next().is_none());
138
139        let stream = [2, 3];
140        assert!(nal_units(&stream).next().is_none());
141
142        let stream = [0, 0, 1];
143        assert_eq!(nal_units(&stream).next().unwrap(), &[0, 0, 1]);
144
145        let stream = [0, 0, 1, 2];
146        assert_eq!(nal_units(&stream).next().unwrap(), &[0, 0, 1, 2]);
147
148        let stream = [0, 0, 1, 2, 0, 0, 1];
149        let mut split = nal_units(&stream);
150        assert_eq!(split.next().unwrap(), &[0, 0, 1, 2]);
151        assert_eq!(split.next().unwrap(), &[0, 0, 1]);
152        assert!(split.next().is_none());
153
154        let stream = [0, 0, 0, 0, 0, 1, 2, 0, 0, 1];
155        let mut split = nal_units(&stream);
156        assert_eq!(split.next().unwrap(), &[0, 0, 1, 2]);
157        assert_eq!(split.next().unwrap(), &[0, 0, 1]);
158        assert!(split.next().is_none());
159
160        let stream = [0, 0, 0, 0, 0, 1, 2, 0, 0];
161        let mut split = nal_units(&stream);
162        assert_eq!(split.next().unwrap(), &[0, 0, 1, 2, 0, 0]);
163        assert!(split.next().is_none());
164
165        let stream = [0, 0, 0, 0, 0, 1, 2, 0, 0, 1, 2, 3, 0, 0, 1];
166        let mut split = nal_units(&stream);
167        assert_eq!(split.next().unwrap(), &[0, 0, 1, 2]);
168        assert_eq!(split.next().unwrap(), &[0, 0, 1, 2, 3]);
169        assert_eq!(split.next().unwrap(), &[0, 0, 1]);
170        assert!(split.next().is_none());
171    }
172
173    #[test]
174    fn nal_mark_stream_boundary() {
175        let v1 = [1, 2, 3, 0];
176        let v2 = [0, 1, 104, 238, 56, 127, 0];
177        let v3 = [0, 0, 1, 104, 238, 56, 128, 0];
178
179        let mut np = NalParser::new();
180
181        // nothing read, read some data
182        assert_eq!(None, np.next());
183
184        np.feed(v1);
185        assert_eq!(None, np.next());
186
187        np.feed(v2);
188        assert_eq!(None, np.next());
189
190        np.feed(v3);
191        assert_eq!(Some(vec![0, 0, 1, 104, 238, 56, 127, 0]), np.next());
192        assert_eq!(None, np.next());
193    }
194
195    #[test]
196    fn nal_mark_empty() {
197        let mut np = NalParser::new();
198        assert_eq!(None, np.next());
199    }
200
201    #[test]
202    fn nal_mark_no_mark() {
203        let mut np = NalParser::new();
204        np.feed([2, 3]);
205        assert_eq!(None, np.next());
206    }
207
208    #[test]
209    fn nal_mark_single_mark() {
210        let mut np = NalParser::new();
211        np.feed([0, 0, 1]);
212        assert_eq!(None, np.next());
213    }
214
215    #[test]
216    fn nal_mark_multiple_marks_same_vec() {
217        let mut np = NalParser::new();
218        np.feed([1, 2, 3, 4, 5, 0, 0, 1, 22, 33, 44, 0, 0, 0, 1, 0, 5, 6, 7, 0, 0, 1, 7, 8, 9]);
219        assert_eq!(None, np.next());
220        assert_eq!(Some(vec![0, 0, 1, 22, 33, 44, 0]), np.next());
221        assert_eq!(Some(vec![0, 0, 1, 0, 5, 6, 7]), np.next());
222        assert_eq!(None, np.next());
223    }
224
225    #[test]
226    fn nal_mark_multiple_marks() {
227        let mut np = NalParser::new();
228
229        np.feed([0, 0, 1, 2, 3, 4, 0, 0, 1]);
230        assert_eq!(None, np.next());
231        assert_eq!(Some(vec![0, 0, 1, 2, 3, 4]), np.next());
232        assert_eq!(None, np.next());
233
234        np.feed([2, 2, 2]);
235        assert_eq!(None, np.next());
236
237        np.feed([3, 3, 3, 0, 0, 1, 5, 6, 7]);
238        assert_eq!(Some(vec![0, 0, 1, 2, 2, 2, 3, 3, 3]), np.next());
239        assert_eq!(None, np.next());
240    }
241}
openh264/utils.rs

openh264/
utils.rs
