less_avc/
nal_unit.rs

1// Copyright 2022-2023 Andrew D. Straw.
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT
5// or http://opensource.org/licenses/MIT>, at your option. This file may not be
6// copied, modified, or distributed except according to those terms.
7
8//! Network Abstraction Layer (NAL) encoding
9
10use super::*;
11
12/// Data to save a NAL unit
13///
14/// The data is in the raw byte sequence payload (RBSP) representation and gets
15/// converted to a NAL unit by the [Self::to_annex_b_data] method.
16pub struct NalUnit {
17    ref_idc: NalRefIdc,
18    unit_type: NalUnitType,
19    rbsp_data: RbspData,
20}
21
22impl NalUnit {
23    /// Create new [NalUnit].
24    pub fn new(ref_idc: NalRefIdc, unit_type: NalUnitType, rbsp_data: RbspData) -> Self {
25        Self {
26            ref_idc,
27            unit_type,
28            rbsp_data,
29        }
30    }
31
32    fn to_buf(&self, with_frame: bool) -> Vec<u8> {
33        #[allow(clippy::identity_op)]
34        // forbidden_zero_bit = 0
35        let nal_byte = 0x00 | (self.ref_idc.nal_ref_idc() << 5) | self.unit_type.nal_unit_type();
36
37        let rbsp_buf = &self.rbsp_data.data;
38        let rbsp_size = rbsp_buf.len();
39        let max_nal_buf_size = calc_max_nal_buf_size(rbsp_size);
40
41        let n_start = if with_frame { 5 } else { 1 };
42        let mut result = vec![0u8; n_start + max_nal_buf_size];
43        if with_frame {
44            result[..4].copy_from_slice(&[0x00, 0x00, 0x00, 0x01]);
45        }
46        result[n_start - 1] = nal_byte;
47
48        let nal_buf_sz = rbsp_to_ebsp(&self.rbsp_data.data, &mut result[n_start..]);
49        let final_sz = n_start + nal_buf_sz;
50        result.truncate(final_sz);
51
52        result
53    }
54
55    /// Return a single "naked" NAL unit.
56    ///
57    /// This is the encapsulated byte sequence payload (EBSP) without NALU
58    /// Header.
59    pub fn to_nal_unit(&self) -> Vec<u8> {
60        self.to_buf(false)
61    }
62    /// Return a single NAL unit encoded for direct saving to `.h264` file.
63    pub fn to_annex_b_data(&self) -> Vec<u8> {
64        self.to_buf(true)
65    }
66}
67
68/// Calculate the maximum possible NAL buffer size for a given RBSP size.
69#[inline]
70fn calc_max_nal_buf_size(rbsp_size: usize) -> usize {
71    (div_ceil(rbsp_size as u32 * 3, 2) * 2).try_into().unwrap()
72}
73
74/// Convert Raw byte sequence payload (RBSP) data to Encapsulated Byte Sequence
75/// Payload (EBSP) bytes.
76pub(crate) fn rbsp_to_ebsp(rbsp_buf: &[u8], nal_buf: &mut [u8]) -> usize {
77    let rbsp_size = rbsp_buf.len();
78    let max_nal_buf_size = calc_max_nal_buf_size(rbsp_size);
79    assert!(nal_buf.len() >= max_nal_buf_size);
80    let mut dest_len = 0;
81
82    let mut input_buf = rbsp_buf;
83
84    while let Some(first_idx) = memchr::memchr(0x00, input_buf) {
85        if first_idx + 1 < input_buf.len() {
86            // more input exists
87            if input_buf[first_idx + 1] == 0x00 {
88                // two nulls in a row
89                if first_idx + 2 < input_buf.len() {
90                    // it is longer
91                    let pos3 = input_buf[first_idx + 2];
92                    if needs_protecting_in_pos3(pos3) {
93                        let src = &input_buf[..first_idx + 2];
94                        nal_buf[dest_len..dest_len + src.len()].copy_from_slice(src);
95                        dest_len += src.len();
96                        nal_buf[dest_len] = 0x03;
97                        dest_len += 1;
98                        input_buf = &input_buf[src.len()..];
99                    } else {
100                        let src = &input_buf[..first_idx + 2];
101                        nal_buf[dest_len..dest_len + src.len()].copy_from_slice(src);
102                        dest_len += src.len();
103                        input_buf = &input_buf[src.len()..];
104                    }
105                } else {
106                    // no more input
107                    break;
108                }
109            } else {
110                // next index is not null, use input up to and including null
111                let src = &input_buf[..first_idx + 1];
112                nal_buf[dest_len..dest_len + src.len()].copy_from_slice(src);
113                dest_len += src.len();
114                input_buf = &input_buf[src.len()..];
115            }
116        } else {
117            // no more input
118            break;
119        }
120    }
121
122    if !input_buf.is_empty() {
123        nal_buf[dest_len..dest_len + input_buf.len()].copy_from_slice(input_buf);
124        dest_len += input_buf.len();
125    }
126
127    dest_len
128}
129
130#[inline]
131/// Returns true if byte is 0x00, 0x01, 0x02 or 0x03.
132fn needs_protecting_in_pos3(byte: u8) -> bool {
133    matches!(byte, 0x00 | 0x01 | 0x02 | 0x03)
134}
135
136#[test]
137fn test_bad_byte() {
138    assert!(needs_protecting_in_pos3(0x00));
139    assert!(needs_protecting_in_pos3(0x01));
140    assert!(needs_protecting_in_pos3(0x02));
141    assert!(needs_protecting_in_pos3(0x03));
142    assert!(!needs_protecting_in_pos3(0x04));
143    for byte in 4..=255 {
144        assert!(!needs_protecting_in_pos3(byte));
145    }
146}
147
148#[test]
149fn test_nal_encoding_roundtrip() {
150    // `h264_reader::rbsp::decode_nal` trims the first byte.
151    let test_vecs = [
152        vec![0x68, 0x00],
153        vec![0x68, 0x01],
154        vec![0x68, 0x02],
155        vec![0x68, 0x03],
156        vec![0x68, 0x04],
157        vec![0x68, 0x00, 0x00],
158        vec![0x68, 0x00, 0x01],
159        vec![0x68, 0x00, 0x02],
160        vec![0x68, 0x00, 0x03],
161        vec![0x68, 0x00, 0x04],
162        vec![0x68, 0x00, 0x00, 0x00],
163        vec![0x68, 0x00, 0x00, 0x01],
164        vec![0x68, 0x00, 0x00, 0x02],
165        vec![0x68, 0x00, 0x00, 0x03],
166        vec![0x68, 0x00, 0x00, 0x04],
167        vec![0x68, 0x00, 0x00, 0x00, 0x00],
168        vec![0x68, 0x00, 0x00, 0x00, 0x01],
169        vec![0x68, 0x00, 0x00, 0x00, 0x02],
170        vec![0x68, 0x00, 0x00, 0x00, 0x03],
171        vec![0x68, 0x00, 0x00, 0x00, 0x04],
172        vec![0x68, 0x00, 0x00, 0x00, 0x05],
173        vec![0x68, 0x03, 0x03, 0x03, 0x03],
174        vec![0x68, 0x00, 0x00, 0x00, 0x00, 0x00],
175        vec![0x68, 0x00, 0x00, 0x00, 0x01, 0x00],
176        vec![0x68, 0x00, 0x00, 0x00, 0x02, 0x00],
177        vec![0x68, 0x00, 0x00, 0x00, 0x03, 0x00],
178        vec![0x68, 0x00, 0x00, 0x00, 0x04, 0x00],
179        vec![0x68, 0x00, 0x00, 0x00, 0x05, 0x00],
180        vec![0x68, 0x03, 0x03, 0x03, 0x03, 0x03],
181        vec![0x68, 0x00, 0x00, 0x03, 0x00, 0x00, 0x00],
182        vec![0x68, 0x00, 0x00, 0x03, 0x00, 0x00, 0x01],
183        vec![0x68, 0x00, 0x00, 0x03, 0x00, 0x00, 0x02],
184        vec![0x68, 0x00, 0x00, 0x03, 0x00, 0x00, 0x03],
185        vec![0x68, 0x00, 0x00, 0x03, 0x00, 0x00, 0x04],
186        vec![0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
187        vec![0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01],
188        vec![0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00],
189        vec![0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01],
190        vec![0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02],
191        vec![0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x03],
192        vec![0x68, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x04],
193    ];
194    for orig in test_vecs.iter() {
195        let mut encoded = vec![0u8; calc_max_nal_buf_size(orig.len())];
196        let sz = rbsp_to_ebsp(orig, &mut encoded);
197        encoded.truncate(sz);
198
199        let decoded = h264_reader::rbsp::decode_nal(&encoded).unwrap();
200        assert_eq!(&orig.as_slice()[1..], decoded.as_ref());
201    }
202}
203
204/// Possible values for the `nal_ref_idc` field in the `nal_unit`.
205///
206/// Encodes to 2 bits.
207pub enum NalRefIdc {
208    // TODO: could these have better names?
209    Zero,
210    One,
211    Two,
212    Three,
213}
214
215impl NalRefIdc {
216    pub(crate) fn nal_ref_idc(&self) -> u8 {
217        match self {
218            Self::Zero => 0,
219            Self::One => 1,
220            Self::Two => 2,
221            Self::Three => 3,
222        }
223    }
224}
225
226/// Possible values for the `nal_unit_type` field in `nal_unit`.
227///
228/// Encodes to 5 bits.
229#[allow(dead_code)]
230#[derive(PartialEq, Eq)]
231#[non_exhaustive]
232pub enum NalUnitType {
233    /// Unspecified
234    Unspecified,
235    /// Coded slice of a non-IDR picture
236    CodedSliceOfANonIDRPicture,
237    /// Coded slice data partition A
238    CodedSliceDataPartitionA,
239    /// Coded slice data partition B
240    CodedSliceDataPartitionB,
241    /// Coded slice data partition C
242    CodedSliceDataPartitionC,
243    /// Coded slice of an IDR picture
244    CodedSliceOfAnIDRPicture,
245    /// Supplemental enhancement information (SEI)
246    SupplementalEnhancementInformation,
247    /// Sequence parameter set
248    SequenceParameterSet,
249    /// Picture parameter set
250    PictureParameterSet,
251    // There are more, which is why this is marked `non_exhaustive`.
252}
253
254impl NalUnitType {
255    pub(crate) fn nal_unit_type(&self) -> u8 {
256        match self {
257            Self::Unspecified => 0,
258            Self::CodedSliceOfANonIDRPicture => 1,
259            Self::CodedSliceDataPartitionA => 2,
260            Self::CodedSliceDataPartitionB => 3,
261            Self::CodedSliceDataPartitionC => 4,
262            Self::CodedSliceOfAnIDRPicture => 5,
263            Self::SupplementalEnhancementInformation => 6,
264            Self::SequenceParameterSet => 7,
265            Self::PictureParameterSet => 8,
266        }
267    }
268}
269
270/// The initial [NalUnit] returned when starting a [LessEncoder].
271pub struct InitialNalUnits {
272    /// sequence parameter set NAL unit
273    pub sps: NalUnit,
274    /// picture parameter set NAL unit
275    pub pps: NalUnit,
276    /// frame NAL unit
277    pub frame: NalUnit,
278}
279
280impl InitialNalUnits {
281    /// Return an [Iterator] over the NAL units generated at the start of encoding.
282    pub fn into_iter(self) -> impl Iterator<Item = NalUnit> {
283        vec![self.sps, self.pps, self.frame].into_iter()
284    }
285}
less_avc/nal_unit.rs

less_avc/
nal_unit.rs
