h264_reader/nal/
mod.rs

1//! Types for reading H264 _Network Abstraction Layer_ Units (NAL Units).
2//!
3//! The data presented must already be in _RBSP_ form (i.e. have been passed through
4//! [`RbspDecoder`](../rbsp/struct.RbspDecoder.html)), where it has been encoded with
5//! 'emulation prevention bytes'.
6
7pub mod pps;
8pub mod sei;
9pub mod slice;
10pub mod sps;
11
12use crate::rbsp;
13use hex_slice::AsHex;
14use std::fmt;
15
16#[derive(PartialEq, Hash, Debug, Copy, Clone)]
17pub enum UnitType {
18    /// The values `0` and `24`-`31` are unspecified in the H264 spec
19    Unspecified(u8),
20    SliceLayerWithoutPartitioningNonIdr,
21    SliceDataPartitionALayer,
22    SliceDataPartitionBLayer,
23    SliceDataPartitionCLayer,
24    SliceLayerWithoutPartitioningIdr,
25    /// Supplemental enhancement information
26    SEI,
27    SeqParameterSet,
28    PicParameterSet,
29    AccessUnitDelimiter,
30    EndOfSeq,
31    EndOfStream,
32    FillerData,
33    SeqParameterSetExtension,
34    PrefixNALUnit,
35    SubsetSeqParameterSet,
36    DepthParameterSet,
37    SliceLayerWithoutPartitioningAux,
38    SliceExtension,
39    SliceExtensionViewComponent,
40    /// The values `17`, `18`, `22` and `23` are reserved for future use by the H264 spec
41    Reserved(u8),
42}
43impl UnitType {
44    pub fn for_id(id: u8) -> Result<UnitType, UnitTypeError> {
45        if id > 31 {
46            Err(UnitTypeError::ValueOutOfRange(id))
47        } else {
48            let t = match id {
49                0 => UnitType::Unspecified(0),
50                1 => UnitType::SliceLayerWithoutPartitioningNonIdr,
51                2 => UnitType::SliceDataPartitionALayer,
52                3 => UnitType::SliceDataPartitionBLayer,
53                4 => UnitType::SliceDataPartitionCLayer,
54                5 => UnitType::SliceLayerWithoutPartitioningIdr,
55                6 => UnitType::SEI,
56                7 => UnitType::SeqParameterSet,
57                8 => UnitType::PicParameterSet,
58                9 => UnitType::AccessUnitDelimiter,
59                10 => UnitType::EndOfSeq,
60                11 => UnitType::EndOfStream,
61                12 => UnitType::FillerData,
62                13 => UnitType::SeqParameterSetExtension,
63                14 => UnitType::PrefixNALUnit,
64                15 => UnitType::SubsetSeqParameterSet,
65                16 => UnitType::DepthParameterSet,
66                17..=18 => UnitType::Reserved(id),
67                19 => UnitType::SliceLayerWithoutPartitioningAux,
68                20 => UnitType::SliceExtension,
69                21 => UnitType::SliceExtensionViewComponent,
70                22..=23 => UnitType::Reserved(id),
71                24..=31 => UnitType::Unspecified(id),
72                _ => panic!("unexpected {}", id), // shouldn't happen
73            };
74            Ok(t)
75        }
76    }
77
78    pub fn id(self) -> u8 {
79        match self {
80            UnitType::Unspecified(v) => v,
81            UnitType::SliceLayerWithoutPartitioningNonIdr => 1,
82            UnitType::SliceDataPartitionALayer => 2,
83            UnitType::SliceDataPartitionBLayer => 3,
84            UnitType::SliceDataPartitionCLayer => 4,
85            UnitType::SliceLayerWithoutPartitioningIdr => 5,
86            UnitType::SEI => 6,
87            UnitType::SeqParameterSet => 7,
88            UnitType::PicParameterSet => 8,
89            UnitType::AccessUnitDelimiter => 9,
90            UnitType::EndOfSeq => 10,
91            UnitType::EndOfStream => 11,
92            UnitType::FillerData => 12,
93            UnitType::SeqParameterSetExtension => 13,
94            UnitType::PrefixNALUnit => 14,
95            UnitType::SubsetSeqParameterSet => 15,
96            UnitType::DepthParameterSet => 16,
97            UnitType::SliceLayerWithoutPartitioningAux => 19,
98            UnitType::SliceExtension => 20,
99            UnitType::SliceExtensionViewComponent => 21,
100            UnitType::Reserved(v) => v,
101        }
102    }
103}
104
105#[derive(Debug)]
106pub enum UnitTypeError {
107    /// if the value was outside the range `0` - `31`.
108    ValueOutOfRange(u8),
109}
110
111#[derive(Copy, Clone, PartialEq, Eq)]
112pub struct NalHeader(u8);
113
114#[derive(Debug)]
115pub enum NalHeaderError {
116    /// The most significant bit of the header, called `forbidden_zero_bit`, was set to 1.
117    ForbiddenZeroBit,
118}
119impl NalHeader {
120    pub fn new(header_value: u8) -> Result<NalHeader, NalHeaderError> {
121        if header_value & 0b1000_0000 != 0 {
122            Err(NalHeaderError::ForbiddenZeroBit)
123        } else {
124            Ok(NalHeader(header_value))
125        }
126    }
127
128    pub fn nal_ref_idc(self) -> u8 {
129        (self.0 & 0b0110_0000) >> 5
130    }
131
132    pub fn nal_unit_type(self) -> UnitType {
133        UnitType::for_id(self.0 & 0b0001_1111).unwrap()
134    }
135}
136impl From<NalHeader> for u8 {
137    fn from(v: NalHeader) -> Self {
138        v.0
139    }
140}
141
142impl fmt::Debug for NalHeader {
143    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
144        f.debug_struct("NalHeader")
145            .field("nal_ref_idc", &self.nal_ref_idc())
146            .field("nal_unit_type", &self.nal_unit_type())
147            .finish()
148    }
149}
150
151/// A partially- or completely-buffered encoded NAL.
152
153/// Must have at least one byte (the header). Partially-encoded NALs are *prefixes*
154/// of a complete NAL. They can always be parsed from the beginning.
155///
156///
157/// ```
158/// use h264_reader::nal::{Nal, RefNal, UnitType};
159/// use h264_reader::rbsp::BitRead;
160/// use std::io::{ErrorKind, Read};
161/// let nal_bytes = &b"\x68\x12\x34\x00\x00\x03\x00\x86"[..];
162/// let nal = RefNal::new(nal_bytes, &[], true);
163///
164/// // Basic inspection:
165/// assert!(nal.is_complete());
166/// assert_eq!(nal.header().unwrap().nal_unit_type(), UnitType::PicParameterSet);
167///
168/// // Reading NAL bytes:
169/// let mut buf = Vec::new();
170/// nal.reader().read_to_end(&mut buf);
171/// assert_eq!(buf, nal_bytes);
172///
173/// // Reading from a partial NAL:
174/// let partial_nal = RefNal::new(&nal_bytes[0..2], &[], false);
175/// assert!(!partial_nal.is_complete());
176/// let mut r = partial_nal.reader();
177/// buf.resize(2, 0u8);
178/// r.read_exact(&mut buf).unwrap(); // reading buffered bytes works.
179/// assert_eq!(&buf[..], &b"\x68\x12"[..]);
180/// buf.resize(1, 0u8);
181/// let e = r.read_exact(&mut buf).unwrap_err(); // beyond returns WouldBlock.
182/// assert_eq!(e.kind(), ErrorKind::WouldBlock);
183///
184/// // Reading RBSP bytes (no header byte, `03` removed from `00 00 03` sequences):
185/// buf.clear();
186/// nal.rbsp_bytes().read_to_end(&mut buf);
187/// assert_eq!(buf, &b"\x12\x34\x00\x00\x00\x86"[..]);
188///
189/// // Reading RBSP bytes of invalid NALs:
190/// let invalid_nal = RefNal::new(&b"\x68\x12\x34\x00\x00\x00\x86"[..], &[], true);
191/// buf.clear();
192/// assert_eq!(invalid_nal.rbsp_bytes().read_to_end(&mut buf).unwrap_err().kind(),
193///            ErrorKind::InvalidData);
194///
195/// // Reading RBSP as a bit sequence:
196/// let mut r = nal.rbsp_bits();
197/// assert_eq!(r.read::<u8>(4, "first nibble").unwrap(), 0x1);
198/// assert_eq!(r.read::<u8>(4, "second nibble").unwrap(), 0x2);
199/// assert_eq!(r.read::<u32>(23, "23 bits at a time").unwrap(), 0x1a_00_00);
200/// assert!(r.has_more_rbsp_data("more left").unwrap());
201/// ```
202pub trait Nal {
203    type BufRead: std::io::BufRead + Clone;
204
205    /// Returns whether the NAL is completely buffered.
206    fn is_complete(&self) -> bool;
207
208    /// Returns the NAL header or error if corrupt.
209    fn header(&self) -> Result<NalHeader, NalHeaderError>;
210
211    /// Reads the bytes in NAL form (including the header byte and
212    /// any emulation-prevention-three-bytes) as a [`std::io::BufRead`].
213    /// If the NAL is incomplete, reads may fail with [`std::io::ErrorKind::WouldBlock`].
214    fn reader(&self) -> Self::BufRead;
215
216    /// Reads the bytes in RBSP form (skipping header byte and
217    /// emulation-prevention-three-bytes).
218    #[inline]
219    fn rbsp_bytes(&self) -> rbsp::ByteReader<Self::BufRead> {
220        rbsp::ByteReader::skipping_h264_header(self.reader())
221    }
222
223    /// Reads bits within the RBSP form.
224    #[inline]
225    fn rbsp_bits(&self) -> rbsp::BitReader<rbsp::ByteReader<Self::BufRead>> {
226        rbsp::BitReader::new(self.rbsp_bytes())
227    }
228}
229
230/// A partially- or completely-buffered [`Nal`] backed by borrowed `&[u8]`s. See [`Nal`] docs.
231#[derive(Clone, Eq, PartialEq)]
232pub struct RefNal<'a> {
233    header: u8,
234    complete: bool,
235
236    // Non-empty chunks.
237    head: &'a [u8],
238    tail: &'a [&'a [u8]],
239}
240impl<'a> RefNal<'a> {
241    /// The caller must ensure that each provided chunk is non-empty.
242    #[inline]
243    pub fn new(head: &'a [u8], tail: &'a [&'a [u8]], complete: bool) -> Self {
244        for buf in tail {
245            debug_assert!(!buf.is_empty());
246        }
247        Self {
248            header: *head.first().expect("RefNal must be non-empty"),
249            head,
250            tail,
251            complete,
252        }
253    }
254}
255impl<'a> Nal for RefNal<'a> {
256    type BufRead = RefNalReader<'a>;
257
258    #[inline]
259    fn is_complete(&self) -> bool {
260        self.complete
261    }
262
263    #[inline]
264    fn header(&self) -> Result<NalHeader, NalHeaderError> {
265        NalHeader::new(self.header)
266    }
267
268    #[inline]
269    fn reader(&self) -> Self::BufRead {
270        RefNalReader {
271            cur: self.head,
272            tail: self.tail,
273            complete: self.complete,
274        }
275    }
276}
277impl<'a> std::fmt::Debug for RefNal<'a> {
278    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
279        // Interpret the NAL header and display the data as a hex string.
280        f.debug_struct("RefNal")
281            .field("header", &self.header())
282            .field(
283                "data",
284                &RefNalReader {
285                    cur: self.head,
286                    tail: self.tail,
287                    complete: self.complete,
288                },
289            )
290            .finish()
291    }
292}
293
294/// A reader through the bytes of a partially- or fully-buffered [`RefNal`]
295/// that implements [`std::io::BufRead`].
296///
297/// Returns [`std::io::ErrorKind::WouldBlock`] on reaching the end of partially-buffered NAL.
298/// Construct via [`Nal::reader`].
299#[derive(Clone)]
300pub struct RefNalReader<'a> {
301    /// Empty only if at end.
302    cur: &'a [u8],
303    tail: &'a [&'a [u8]],
304    complete: bool,
305}
306impl<'a> RefNalReader<'a> {
307    fn next_chunk(&mut self) {
308        match self.tail {
309            [first, tail @ ..] => {
310                self.cur = first;
311                self.tail = tail;
312            }
313            _ => self.cur = &[], // EOF.
314        }
315    }
316}
317impl<'a> std::io::Read for RefNalReader<'a> {
318    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
319        let len;
320        if buf.is_empty() {
321            len = 0;
322        } else if self.cur.is_empty() && !self.complete {
323            return Err(std::io::Error::new(
324                std::io::ErrorKind::WouldBlock,
325                "reached end of partially-buffered NAL",
326            ));
327        } else if buf.len() < self.cur.len() {
328            len = buf.len();
329            let (copy, keep) = self.cur.split_at(len);
330            buf.copy_from_slice(copy);
331            self.cur = keep;
332        } else {
333            len = self.cur.len();
334            buf[..len].copy_from_slice(self.cur);
335            self.next_chunk();
336        }
337        Ok(len)
338    }
339}
340impl<'a> std::io::BufRead for RefNalReader<'a> {
341    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
342        if self.cur.is_empty() && !self.complete {
343            return Err(std::io::Error::new(
344                std::io::ErrorKind::WouldBlock,
345                "reached end of partially-buffered NAL",
346            ));
347        }
348        Ok(self.cur)
349    }
350    fn consume(&mut self, amt: usize) {
351        self.cur = &self.cur[amt..];
352        if self.cur.is_empty() {
353            self.next_chunk();
354        }
355    }
356}
357impl<'a> std::fmt::Debug for RefNalReader<'a> {
358    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
359        write!(f, "{:02x}", self.cur.plain_hex(true))?;
360        for buf in self.tail {
361            write!(f, " {:02x}", buf.plain_hex(true))?;
362        }
363        if !self.complete {
364            f.write_str(" ...")?;
365        }
366        Ok(())
367    }
368}
369
370#[cfg(test)]
371mod test {
372    use std::io::{BufRead, Read};
373
374    use super::*;
375
376    #[test]
377    fn header() {
378        let h = NalHeader::new(0b0101_0001).unwrap();
379        assert_eq!(0b10, h.nal_ref_idc());
380        assert_eq!(UnitType::Reserved(17), h.nal_unit_type());
381    }
382
383    #[test]
384    fn ref_nal() {
385        fn common<'a>(head: &'a [u8], tail: &'a [&'a [u8]], complete: bool) -> RefNal<'a> {
386            let nal = RefNal::new(head, tail, complete);
387            assert_eq!(NalHeader::new(0b0101_0001).unwrap(), nal.header().unwrap());
388
389            // Try the Read impl.
390            let mut r = nal.reader();
391            let mut buf = [0u8; 5];
392            r.read_exact(&mut buf).unwrap();
393            assert_eq!(&buf[..], &[0b0101_0001, 1, 2, 3, 4]);
394            if complete {
395                assert_eq!(r.read(&mut buf[..]).unwrap(), 0);
396
397                // Also try read_to_end.
398                let mut buf = Vec::new();
399                nal.reader().read_to_end(&mut buf).unwrap();
400                assert_eq!(buf, &[0b0101_0001, 1, 2, 3, 4]);
401            } else {
402                assert_eq!(
403                    r.read(&mut buf[..]).unwrap_err().kind(),
404                    std::io::ErrorKind::WouldBlock
405                );
406            }
407
408            // Let the caller try the BufRead impl.
409            nal
410        }
411
412        // Incomplete NAL with a first chunk only.
413        let nal = common(&[0b0101_0001, 1, 2, 3, 4], &[], false);
414        let mut r = nal.reader();
415        assert_eq!(r.fill_buf().unwrap(), &[0b0101_0001, 1, 2, 3, 4]);
416        r.consume(1);
417        assert_eq!(r.fill_buf().unwrap(), &[1, 2, 3, 4]);
418        r.consume(4);
419        assert_eq!(
420            r.fill_buf().unwrap_err().kind(),
421            std::io::ErrorKind::WouldBlock
422        );
423
424        // Incomplete NAL with multiple chunks.
425        let nal = common(&[0b0101_0001], &[&[1, 2], &[3, 4]], false);
426        let mut r = nal.reader();
427        assert_eq!(r.fill_buf().unwrap(), &[0b0101_0001]);
428        r.consume(1);
429        assert_eq!(r.fill_buf().unwrap(), &[1, 2]);
430        r.consume(2);
431        assert_eq!(r.fill_buf().unwrap(), &[3, 4]);
432        r.consume(1);
433        assert_eq!(r.fill_buf().unwrap(), &[4]);
434        r.consume(1);
435        assert_eq!(
436            r.fill_buf().unwrap_err().kind(),
437            std::io::ErrorKind::WouldBlock
438        );
439
440        // Complete NAL with first chunk only.
441        let nal = common(&[0b0101_0001, 1, 2, 3, 4], &[], true);
442        let mut r = nal.reader();
443        assert_eq!(r.fill_buf().unwrap(), &[0b0101_0001, 1, 2, 3, 4]);
444        r.consume(1);
445        assert_eq!(r.fill_buf().unwrap(), &[1, 2, 3, 4]);
446        r.consume(4);
447        assert!(r.fill_buf().unwrap().is_empty());
448    }
449
450    #[test]
451    fn reader_debug() {
452        assert_eq!(
453            format!(
454                "{:?}",
455                RefNalReader {
456                    cur: &b"\x00"[..],
457                    tail: &[&b"\x01"[..], &b"\x02\x03"[..]],
458                    complete: false,
459                }
460            ),
461            "00 01 02 03 ..."
462        );
463    }
464}
h264_reader/nal/mod.rs

h264_reader/nal/
mod.rs
