arcly_stream/codec/
h264.rs

1//! H.264 (AVC) bitstream helpers: NAL iteration, Annex-B ↔ AVCC conversion,
2//! and SPS resolution extraction.
3
4use super::bitreader::BitReader;
5use super::nal::{iter_nals, unescape_rbsp};
6use super::parser::{CodecParser, VideoParams};
7use crate::CodecId;
8use bytes::{BufMut, Bytes, BytesMut};
9
10/// NAL unit type for a sequence parameter set.
11pub const NAL_SPS: u8 = 7;
12/// NAL unit type for a picture parameter set.
13pub const NAL_PPS: u8 = 8;
14/// NAL unit type for a coded slice of an IDR picture (keyframe).
15pub const NAL_IDR: u8 = 5;
16
17/// Iterate the NAL units of an Annex-B (`00 00 01` / `00 00 00 01`) bytestream,
18/// yielding each NAL payload **without** its start code.
19pub fn iter_nals_annexb(data: &[u8]) -> impl Iterator<Item = &[u8]> {
20    iter_nals(data)
21}
22
23/// Convert an Annex-B bytestream to AVCC (length-prefixed) with a 4-byte length.
24pub fn annexb_to_avcc(data: &[u8]) -> Bytes {
25    let mut out = BytesMut::with_capacity(data.len());
26    for nal in iter_nals(data) {
27        out.put_u32(nal.len() as u32);
28        out.put_slice(nal);
29    }
30    out.freeze()
31}
32
33/// Convert an AVCC (length-prefixed) bytestream to Annex-B (`00 00 00 01`).
34///
35/// `nal_length_size` is the prefix width in bytes (1–4), from the
36/// AVCDecoderConfigurationRecord (`lengthSizeMinusOne + 1`).
37pub fn avcc_to_annexb(data: &[u8], nal_length_size: usize) -> Option<Bytes> {
38    if !(1..=4).contains(&nal_length_size) {
39        return None;
40    }
41    let mut out = BytesMut::with_capacity(data.len() + 16);
42    let mut i = 0;
43    while i + nal_length_size <= data.len() {
44        let mut len = 0usize;
45        for _ in 0..nal_length_size {
46            len = (len << 8) | data[i] as usize;
47            i += 1;
48        }
49        if i + len > data.len() {
50            return None; // truncated NAL
51        }
52        out.put_slice(&[0, 0, 0, 1]);
53        out.put_slice(&data[i..i + len]);
54        i += len;
55    }
56    Some(out.freeze())
57}
58
59/// Decoded video dimensions and profile from a sequence parameter set.
60#[derive(Debug, Clone, Copy, PartialEq, Eq)]
61pub struct SpsInfo {
62    /// Luma width in pixels (after cropping).
63    pub width: u32,
64    /// Luma height in pixels (after cropping).
65    pub height: u32,
66    /// `profile_idc` (66 = baseline, 77 = main, 100 = high, …).
67    pub profile_idc: u8,
68    /// `level_idc` × 10 (e.g. 31 = level 3.1).
69    pub level_idc: u8,
70}
71
72/// Parse a sequence parameter set NAL (including its 1-byte NAL header) and
73/// extract the coded resolution. Returns `None` on a malformed SPS or one using
74/// a scaling matrix (unsupported here).
75pub fn parse_sps(nal: &[u8]) -> Option<SpsInfo> {
76    if nal.is_empty() || nal[0] & 0x1f != NAL_SPS {
77        return None;
78    }
79    let rbsp = unescape_rbsp(&nal[1..]);
80    let mut r = BitReader::new(&rbsp);
81    if r.remaining() < 24 {
82        return None; // too short to hold profile/constraints/level
83    }
84
85    let profile_idc = r.read_bits(8)? as u8;
86    let _constraints = r.read_bits(8)?;
87    let level_idc = r.read_bits(8)? as u8;
88    let _sps_id = r.read_ue()?;
89
90    let mut chroma_format_idc = 1u32; // 4:2:0 default
91    if matches!(
92        profile_idc,
93        100 | 110 | 122 | 244 | 44 | 83 | 86 | 118 | 128 | 138 | 139 | 134 | 135
94    ) {
95        chroma_format_idc = r.read_ue()?;
96        if chroma_format_idc == 3 {
97            let _separate_colour_plane = r.read_bit()?;
98        }
99        let _bit_depth_luma = r.read_ue()?;
100        let _bit_depth_chroma = r.read_ue()?;
101        let _qpprime = r.read_bit()?;
102        let scaling_matrix_present = r.read_bit()?;
103        if scaling_matrix_present == 1 {
104            return None; // scaling lists not decoded
105        }
106    }
107
108    let _log2_max_frame_num = r.read_ue()?;
109    let pic_order_cnt_type = r.read_ue()?;
110    if pic_order_cnt_type == 0 {
111        let _log2_max_poc_lsb = r.read_ue()?;
112    } else if pic_order_cnt_type == 1 {
113        let _delta_pic_order_always_zero = r.read_bit()?;
114        let _offset_for_non_ref = r.read_se()?;
115        let _offset_for_top_to_bottom = r.read_se()?;
116        let cycle = r.read_ue()?;
117        for _ in 0..cycle {
118            let _ = r.read_se()?;
119        }
120    }
121
122    let _max_num_ref_frames = r.read_ue()?;
123    let _gaps_allowed = r.read_bit()?;
124    let pic_width_in_mbs_minus1 = r.read_ue()?;
125    let pic_height_in_map_units_minus1 = r.read_ue()?;
126    let frame_mbs_only_flag = r.read_bit()?;
127    if frame_mbs_only_flag == 0 {
128        let _mb_adaptive = r.read_bit()?;
129    }
130    let _direct_8x8 = r.read_bit()?;
131
132    let frame_cropping_flag = r.read_bit()?;
133    let (mut crop_left, mut crop_right, mut crop_top, mut crop_bottom) = (0u32, 0u32, 0u32, 0u32);
134    if frame_cropping_flag == 1 {
135        crop_left = r.read_ue()?;
136        crop_right = r.read_ue()?;
137        crop_top = r.read_ue()?;
138        crop_bottom = r.read_ue()?;
139    }
140
141    let width_mbs = pic_width_in_mbs_minus1 + 1;
142    let height_map = pic_height_in_map_units_minus1 + 1;
143    let frame_height_mbs = (2 - frame_mbs_only_flag) * height_map;
144
145    // Chroma sub-sampling factors for the crop unit.
146    let (sub_w, sub_h) = match chroma_format_idc {
147        0 => (1, 1), // monochrome
148        1 => (2, 2), // 4:2:0
149        2 => (2, 1), // 4:2:2
150        _ => (1, 1), // 4:4:4
151    };
152    let crop_unit_x = if chroma_format_idc == 0 { 1 } else { sub_w };
153    let crop_unit_y = (if chroma_format_idc == 0 { 1 } else { sub_h }) * (2 - frame_mbs_only_flag);
154
155    let width = width_mbs * 16 - (crop_left + crop_right) * crop_unit_x;
156    let height = frame_height_mbs * 16 - (crop_top + crop_bottom) * crop_unit_y;
157
158    Some(SpsInfo {
159        width,
160        height,
161        profile_idc,
162        level_idc,
163    })
164}
165
166/// The HLS `CODECS` attribute for H.264, e.g. `"avc1.4d401f"` (profile/level
167/// encoded as `avc1.PPCCLL`). The middle byte is the constraint-set flags (0).
168pub fn hls_codec_string(p: &VideoParams) -> String {
169    format!("avc1.{:02x}00{:02x}", p.profile, p.level)
170}
171
172/// [`CodecParser`] implementation for H.264 / AVC.
173pub struct H264;
174
175impl CodecParser for H264 {
176    const CODEC: CodecId = CodecId::H264;
177
178    fn parse_config(data: &[u8]) -> Option<VideoParams> {
179        for nal in iter_nals(data) {
180            if !nal.is_empty() && nal[0] & 0x1f == NAL_SPS {
181                let sps = parse_sps(nal)?;
182                return Some(VideoParams {
183                    width: sps.width,
184                    height: sps.height,
185                    profile: sps.profile_idc,
186                    level: sps.level_idc,
187                    tier: 0,
188                    bit_depth: 8,
189                });
190            }
191        }
192        None
193    }
194
195    fn is_random_access_point(data: &[u8]) -> bool {
196        iter_nals(data).any(|n| !n.is_empty() && n[0] & 0x1f == NAL_IDR)
197    }
198
199    fn carries_config(data: &[u8]) -> bool {
200        iter_nals(data).any(|n| !n.is_empty() && matches!(n[0] & 0x1f, NAL_SPS | NAL_PPS))
201    }
202
203    fn hls_codec_string(params: &VideoParams) -> String {
204        hls_codec_string(params)
205    }
206}
207
208#[cfg(test)]
209mod tests {
210    use super::*;
211    use crate::codec::testutil::BitWriter;
212
213    #[test]
214    fn iterates_and_converts_nals() {
215        // Two NALs: [9 0xF0] and [7 0x42], with mixed 4- and 3-byte start codes.
216        let annexb = [0, 0, 0, 1, 9, 0xF0, 0, 0, 1, 7, 0x42];
217        let nals: Vec<&[u8]> = iter_nals_annexb(&annexb).collect();
218        assert_eq!(nals, vec![&[9u8, 0xF0][..], &[7u8, 0x42][..]]);
219
220        // Round-trip Annex-B → AVCC → Annex-B (normalizes to 4-byte codes).
221        let avcc = annexb_to_avcc(&annexb);
222        assert_eq!(&avcc[..], &[0, 0, 0, 2, 9, 0xF0, 0, 0, 0, 2, 7, 0x42]);
223        let back = avcc_to_annexb(&avcc, 4).unwrap();
224        assert_eq!(&back[..], &[0, 0, 0, 1, 9, 0xF0, 0, 0, 0, 1, 7, 0x42]);
225    }
226
227    /// Synthesize a baseline (profile 66) SPS for 1280x720, no cropping.
228    fn sps_1280x720() -> Vec<u8> {
229        let mut w = BitWriter::default();
230        w.bits(66, 8); // profile_idc (baseline → no chroma block)
231        w.bits(0, 8); // constraint flags
232        w.bits(31, 8); // level 3.1
233        w.ue(0); // seq_parameter_set_id
234        w.ue(0); // log2_max_frame_num_minus4
235        w.ue(0); // pic_order_cnt_type = 0
236        w.ue(0); // log2_max_pic_order_cnt_lsb_minus4
237        w.ue(1); // max_num_ref_frames
238        w.bit(0); // gaps_in_frame_num_value_allowed
239        w.ue(79); // pic_width_in_mbs_minus1 → 80*16 = 1280
240        w.ue(44); // pic_height_in_map_units_minus1 → 45*16 = 720
241        w.bit(1); // frame_mbs_only_flag
242        w.bit(0); // direct_8x8_inference_flag
243        w.bit(0); // frame_cropping_flag
244        w.bit(0); // vui_parameters_present_flag
245        let mut nal = vec![0x67u8]; // NAL header: SPS (type 7), ref_idc 3
246        nal.extend_from_slice(&w.bytes());
247        nal
248    }
249
250    #[test]
251    fn parse_sps_extracts_resolution() {
252        let sps = parse_sps(&sps_1280x720()).expect("parse");
253        assert_eq!((sps.width, sps.height), (1280, 720));
254        assert_eq!(sps.profile_idc, 66);
255        assert_eq!(sps.level_idc, 31);
256    }
257
258    #[test]
259    fn codec_parser_classifies_and_extracts() {
260        // Annex-B AU: SPS + an IDR slice (type 5).
261        let mut au = vec![0, 0, 0, 1];
262        au.extend_from_slice(&sps_1280x720());
263        au.extend_from_slice(&[0, 0, 0, 1, 0x65, 0xAA]); // IDR NAL (type 5)
264
265        assert!(H264::is_random_access_point(&au));
266        assert!(H264::carries_config(&au));
267        let p = H264::parse_config(&au).expect("params");
268        assert_eq!((p.width, p.height, p.profile, p.level), (1280, 720, 66, 31));
269        assert_eq!(H264::hls_codec_string(&p), "avc1.42001f");
270
271        // A lone non-IDR slice (type 1) is not a RAP.
272        let delta = [0, 0, 0, 1, 0x41, 0xAA];
273        assert!(!H264::is_random_access_point(&delta));
274    }
275}
arcly_stream/codec/h264.rs

arcly_stream/codec/
h264.rs
