1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use anyhow::anyhow;
use bytes::Buf;
use std::fmt::Debug;
use std::io::Cursor;

#[allow(clippy::len_without_is_empty)]
pub trait Header: Sized {
    /// Parse the NALU header, returning it.
    fn parse<T: AsRef<[u8]>>(cursor: &Cursor<T>) -> anyhow::Result<Self>;
    /// Whether this header type indicates EOS.
    fn is_end(&self) -> bool;
    /// The length of the header.
    fn len(&self) -> usize;
}

#[derive(Debug)]
pub struct Nalu<T, U> {
    header: U,
    /// The mapping that backs this NALU. Possibly shared with the other NALUs
    /// in the Access Unit.
    data: T,

    size: usize,
    offset: usize,
    sc_offset: usize,
}

impl<T, U> Nalu<T, U>
where
    T: AsRef<[u8]> + Clone,
    U: Debug + Header,
{
    /// Find the next Annex B encoded NAL unit.
    pub fn next(cursor: &mut Cursor<T>) -> anyhow::Result<Option<Nalu<T, U>>> {
        let bitstream = cursor.clone().into_inner();
        let pos = usize::try_from(cursor.position())?;

        // Find the start code for this NALU
        let current_nalu_offset = match Nalu::<T, U>::find_start_code(cursor, pos) {
            Some(offset) => offset,
            None => return Err(anyhow!("No NAL found")),
        };

        let mut start_code_offset = pos + current_nalu_offset;

        // If the preceding byte is 00, then we actually have a four byte SC,
        // i.e. 00 00 00 01 Where the first 00 is the "zero_byte()"
        if start_code_offset > 0 && cursor.get_ref().as_ref()[start_code_offset - 1] == 00 {
            start_code_offset -= 1;
        }

        // The NALU offset is its offset + 3 bytes to skip the start code.
        let nalu_offset = pos + current_nalu_offset + 3;

        // Set the bitstream position to the start of the current NALU
        cursor.set_position(u64::try_from(nalu_offset)?);

        let hdr = U::parse(cursor)?;

        // Find the start of the subsequent NALU.
        let mut next_nalu_offset = match Nalu::<T, U>::find_start_code(cursor, nalu_offset) {
            Some(offset) => offset,
            None => cursor.chunk().len(), // Whatever data is left must be part of the current NALU
        };

        while next_nalu_offset > 0
            && cursor.get_ref().as_ref()[nalu_offset + next_nalu_offset - 1] == 00
        {
            // Discard trailing_zero_8bits
            next_nalu_offset -= 1;
        }

        let nal_size = if hdr.is_end() { 1 } else { next_nalu_offset };

        Ok(Some(Nalu {
            header: hdr,
            data: bitstream,
            size: nal_size,
            offset: nalu_offset,
            sc_offset: start_code_offset,
        }))
    }
}

impl<T, U> Nalu<T, U>
where
    T: AsRef<[u8]>,
    U: Debug,
{
    fn find_start_code(data: &mut Cursor<T>, offset: usize) -> Option<usize> {
        // discard all zeroes until the start code pattern is found
        data.get_ref().as_ref()[offset..]
            .windows(3)
            .position(|window| window == [0x00, 0x00, 0x01])
    }

    /// Get a reference to the nalu's header.
    pub fn header(&self) -> &U {
        &self.header
    }

    /// Get a reference to the nalu's data.
    pub fn data(&self) -> &T {
        &self.data
    }

    /// Get a reference to the nalu's size.
    pub fn size(&self) -> usize {
        self.size
    }

    /// Get a reference to the nalu's offset.
    pub fn offset(&self) -> usize {
        self.offset
    }

    /// Get a reference to the nalu's sc offset.
    pub fn sc_offset(&self) -> usize {
        self.sc_offset
    }
}

impl<T: AsRef<[u8]>, U> AsRef<[u8]> for Nalu<T, U> {
    fn as_ref(&self) -> &[u8] {
        let data = self.data.as_ref();
        &data[self.offset..self.offset + self.size]
    }
}