Skip to main content

arcly_stream/codec/
nal.rs

1//! Shared Annex-B NAL utilities for the NAL-based codecs (H.264, H.265, VVC).
2//!
3//! All three frame their bitstreams identically — `00 00 01` / `00 00 00 01`
4//! start codes with `00 00 03` emulation prevention — and differ only in NAL
5//! header layout and parameter-set field order. Start-code scanning delegates to
6//! the crate-canonical `memchr`-accelerated scanner; this module owns NAL
7//! iteration, RBSP de-emulation, and conformance cropping.
8
9/// Iterate the NAL units of an Annex-B bytestream, yielding each NAL payload
10/// (including its NAL header) **without** the start code.
11pub fn iter_nals(data: &[u8]) -> NalIter<'_> {
12    NalIter { data, pos: 0 }
13}
14
15/// Iterator over Annex-B NAL units. See [`iter_nals`].
16pub struct NalIter<'a> {
17    data: &'a [u8],
18    pos: usize,
19}
20
21impl<'a> Iterator for NalIter<'a> {
22    type Item = &'a [u8];
23
24    fn next(&mut self) -> Option<&'a [u8]> {
25        let start = next_start_code(self.data, self.pos)?;
26        let nal_begin = start.0 + start.1;
27        let nal_end = match next_start_code(self.data, nal_begin) {
28            Some((off, _)) => off,
29            None => self.data.len(),
30        };
31        self.pos = nal_end;
32        if nal_begin >= nal_end {
33            return self.next();
34        }
35        Some(&self.data[nal_begin..nal_end])
36    }
37}
38
39/// Find the next start code at/after `from`, returning `(offset, len)` where
40/// `len` is 3 (`00 00 01`) or 4 (`00 00 00 01`).
41///
42/// The start code must *begin* at or after `from` (its leading zeros are not
43/// sought before the origin) — the contract [`NalIter`] relies on as it resumes
44/// past each consumed NAL.
45pub fn next_start_code(data: &[u8], from: usize) -> Option<(usize, usize)> {
46    crate::bytescan::scan_start_code(data, from, from)
47}
48
49/// Strip emulation-prevention bytes (`00 00 03` → `00 00`) from an RBSP.
50pub fn unescape_rbsp(nal: &[u8]) -> Vec<u8> {
51    let mut out = Vec::with_capacity(nal.len());
52    let mut zeros = 0;
53    for &b in nal {
54        if zeros >= 2 && b == 0x03 {
55            zeros = 0;
56            continue;
57        }
58        if b == 0 {
59            zeros += 1;
60        } else {
61            zeros = 0;
62        }
63        out.push(b);
64    }
65    out
66}
67
68/// Apply HEVC/VVC conformance-window cropping to coded luma dimensions,
69/// returning the displayed `(width, height)`. `chroma_format_idc` selects the
70/// crop unit (4:2:0 → 2×2, 4:2:2 → 2×1, 4:4:4 / mono → 1×1).
71pub fn conformance_dims(
72    width_luma: u32,
73    height_luma: u32,
74    chroma_format_idc: u32,
75    crop_left: u32,
76    crop_right: u32,
77    crop_top: u32,
78    crop_bottom: u32,
79) -> (u32, u32) {
80    let (sub_w, sub_h) = match chroma_format_idc {
81        1 => (2, 2), // 4:2:0
82        2 => (2, 1), // 4:2:2
83        _ => (1, 1), // 4:4:4 or monochrome
84    };
85    let width = width_luma.saturating_sub((crop_left + crop_right) * sub_w);
86    let height = height_luma.saturating_sub((crop_top + crop_bottom) * sub_h);
87    (width, height)
88}
89
90#[cfg(test)]
91mod tests {
92    use super::*;
93
94    #[test]
95    fn iterates_three_and_four_byte_start_codes() {
96        let data = [0, 0, 0, 1, 9, 0xF0, 0, 0, 1, 7, 0x42];
97        let nals: Vec<&[u8]> = iter_nals(&data).collect();
98        assert_eq!(nals, vec![&[9u8, 0xF0][..], &[7u8, 0x42][..]]);
99    }
100
101    #[test]
102    fn unescape_removes_emulation_bytes() {
103        assert_eq!(unescape_rbsp(&[0, 0, 3, 1]), vec![0, 0, 1]);
104        assert_eq!(unescape_rbsp(&[0, 0, 3, 0, 0, 3, 2]), vec![0, 0, 0, 0, 2]);
105    }
106
107    #[test]
108    fn conformance_crop_420() {
109        // 1920x1088 coded, crop 4 rows off the bottom in 4:2:0 → 1920x1080.
110        assert_eq!(conformance_dims(1920, 1088, 1, 0, 0, 0, 4), (1920, 1080));
111    }
112}