oxideav_bitstream/
h264.rs

1//! H.264 / AVC minimal IDR header parsing.
2//!
3//! This module parses just enough of an H.264 Annex-B bitstream to
4//! populate any of the slice-data HW backends' parameter buffers
5//! ([`VAPictureParameterBufferH264`], [`VdpPictureInfoH264`],
6//! [`VkVideoDecodeH264PictureInfoKHR`]). It does NOT handle:
7//!
8//! - DCT, entropy decode, inverse transform, motion compensation,
9//!   in-loop filtering — the GPU does all of that.
10//! - Scaling lists (`seq_scaling_matrix_present_flag` / PPS scaling
11//!   lists). The IDR-only backends submit flat 16/16, which is what
12//!   the encoder default emits.
13//! - FMO / ASO (`num_slice_groups_minus1 > 0`).
14//! - Redundant slices.
15//! - B / SP / SI slice header complexity.
16//! - VUI parsing — only the cropping rectangle is captured.
17//!
18//! The boundary is the same the workspace's other "minimal IDR" code
19//! has used in `crates/oxideav-vdpau/src/h264.rs` since Round 3.
20//!
21//! # Spec references
22//!
23//! ITU-T H.264 / ISO/IEC 14496-10. The sections you'll see referenced
24//! in the comments below are 7.3.2.1.1 (SPS), 7.3.2.2 (PPS), 7.3.3
25//! (slice header), 7.3.2.1.1.1 (frame cropping), 7.4.5 (slice header
26//! semantics), 9.1 (Exp-Golomb).
27
28use crate::bit_reader::BitReader;
29use crate::BitstreamError;
30
31// ─────────────────────────── NAL unit types ──────────────────────────────────
32
33/// H.264 NAL unit type 5 — Coded slice of an IDR picture (7.4.1).
34pub const NAL_TYPE_IDR: u8 = 5;
35/// H.264 NAL unit type 7 — Sequence parameter set.
36pub const NAL_TYPE_SPS: u8 = 7;
37/// H.264 NAL unit type 8 — Picture parameter set.
38pub const NAL_TYPE_PPS: u8 = 8;
39/// H.264 NAL unit type 1 — Coded slice of a non-IDR picture.
40pub const NAL_TYPE_NON_IDR_SLICE: u8 = 1;
41/// H.264 NAL unit type 9 — Access unit delimiter.
42pub const NAL_TYPE_AUD: u8 = 9;
43
44// ─────────────────────────── Annex-B framing ─────────────────────────────────
45
46/// Locate every NAL unit in an Annex-B bitstream and return slices
47/// pointing at the NAL body (start-code stripped, NAL header byte at
48/// index 0, emulation bytes still in place — those are stripped only
49/// during bit-level parsing via [`ebsp_to_rbsp`]).
50///
51/// The HW backends keep the original Annex-B bytes around for
52/// submission to the GPU; this function is for *parsing* only.
53pub fn split_annex_b(buf: &[u8]) -> Vec<&[u8]> {
54    let mut out = Vec::new();
55    let mut i = 0;
56    let n = buf.len();
57    let mut last_payload_start: Option<usize> = None;
58    while i < n {
59        let four =
60            i + 3 < n && buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 0 && buf[i + 3] == 1;
61        let three = !four && i + 2 < n && buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 1;
62        if four || three {
63            if let Some(start) = last_payload_start.take() {
64                out.push(&buf[start..i]);
65            }
66            i += if four { 4 } else { 3 };
67            last_payload_start = Some(i);
68            continue;
69        }
70        i += 1;
71    }
72    if let Some(start) = last_payload_start.take() {
73        out.push(&buf[start..n]);
74    }
75    out
76}
77
78/// Strip H.264 emulation-prevention `0x03` bytes from an EBSP to
79/// produce an RBSP (7.4.1.1). Inserted by the encoder after `00 00 0x`
80/// to keep start codes unique inside the payload.
81pub fn ebsp_to_rbsp(ebsp: &[u8]) -> Vec<u8> {
82    let mut out = Vec::with_capacity(ebsp.len());
83    let mut i = 0;
84    let n = ebsp.len();
85    while i < n {
86        if i + 2 < n && ebsp[i] == 0 && ebsp[i + 1] == 0 && ebsp[i + 2] == 3 {
87            out.push(0);
88            out.push(0);
89            i += 3;
90        } else {
91            out.push(ebsp[i]);
92            i += 1;
93        }
94    }
95    out
96}
97
98/// Inspect the NAL header byte and return `(forbidden_zero_bit,
99/// nal_ref_idc, nal_unit_type)`.
100pub fn nal_header(byte: u8) -> (u8, u8, u8) {
101    let forbidden = (byte >> 7) & 1;
102    let nal_ref_idc = (byte >> 5) & 0x3;
103    let nal_type = byte & 0x1f;
104    (forbidden, nal_ref_idc, nal_type)
105}
106
107// ─────────────────────────── Output structs ──────────────────────────────────
108
109/// Frame cropping rectangle in *picture* (output-image) sample units.
110/// The values themselves are the four offsets that
111/// `frame_crop_*_offset` evaluate to after the SubWidth/SubHeight
112/// scaling rules in 6.4 / 7.4.2.1.1.
113#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
114pub struct H264FrameCropping {
115    pub left: u32,
116    pub right: u32,
117    pub top: u32,
118    pub bottom: u32,
119}
120
121/// Sequence parameter set. Carries every field the three slice-data
122/// HW APIs look at, plus a few derived helpers.
123///
124/// See H.264 7.3.2.1.1 for the syntax this is filled from.
125#[derive(Debug, Clone, Default, PartialEq, Eq)]
126pub struct H264Sps {
127    pub profile_idc: u8,
128    /// Bits 7..0 of the constraint_set_flags + reserved_zero block
129    /// in the order the encoder emitted them — useful for round-tripping.
130    pub constraint_set_flags: u8,
131    pub level_idc: u8,
132    pub seq_parameter_set_id: u8,
133
134    pub chroma_format_idc: u8,
135    pub separate_colour_plane_flag: bool,
136    pub bit_depth_luma_minus8: u8,
137    pub bit_depth_chroma_minus8: u8,
138
139    pub log2_max_frame_num_minus4: u8,
140    pub pic_order_cnt_type: u8,
141    pub log2_max_pic_order_cnt_lsb_minus4: u8,
142    pub delta_pic_order_always_zero_flag: bool,
143    pub max_num_ref_frames: u32,
144    pub gaps_in_frame_num_value_allowed_flag: bool,
145
146    pub pic_width_in_mbs_minus1: u32,
147    pub pic_height_in_map_units_minus1: u32,
148    pub frame_mbs_only_flag: bool,
149    pub mb_adaptive_frame_field_flag: bool,
150    pub direct_8x8_inference_flag: bool,
151
152    pub frame_cropping: Option<H264FrameCropping>,
153}
154
155impl H264Sps {
156    /// Whether this SPS requires the High-profile syntax extension
157    /// (chroma_format / bit_depth / scaling_matrix block).
158    pub fn has_high_profile_extension(profile_idc: u8) -> bool {
159        matches!(
160            profile_idc,
161            100 | 110 | 122 | 244 | 44 | 83 | 86 | 118 | 128 | 138 | 139 | 134 | 135
162        )
163    }
164
165    /// Coded width in luma samples (pic_width_in_mbs * 16). 7.4.2.1.1.
166    pub fn coded_width(&self) -> u32 {
167        (self.pic_width_in_mbs_minus1 + 1) * 16
168    }
169
170    /// Coded height in luma samples. Accounts for
171    /// `frame_mbs_only_flag` per 7.4.2.1.1.
172    pub fn coded_height(&self) -> u32 {
173        let mb_only = if self.frame_mbs_only_flag { 1 } else { 0 };
174        (self.pic_height_in_map_units_minus1 + 1) * 16 * (2 - mb_only)
175    }
176
177    /// Per the SubWidthC table in 6.2 (only the values we actually
178    /// hit — 1 for 4:2:0/4:2:2, 1 for 4:4:4 monochrome handled
179    /// separately).
180    fn sub_width_c(&self) -> u32 {
181        match self.chroma_format_idc {
182            1 | 2 => 2, // 4:2:0, 4:2:2
183            3 => 1,     // 4:4:4
184            _ => 1,
185        }
186    }
187
188    fn sub_height_c(&self) -> u32 {
189        match self.chroma_format_idc {
190            1 => 2,     // 4:2:0
191            2 | 3 => 1, // 4:2:2, 4:4:4
192            _ => 1,
193        }
194    }
195
196    /// Display width in luma samples after applying the frame
197    /// cropping rectangle (6.4 / 7.4.2.1.1.1). For monochrome
198    /// (chroma_format_idc==0) the offsets are interpreted in luma
199    /// samples directly.
200    pub fn display_width(&self) -> u32 {
201        let coded = self.coded_width();
202        let Some(c) = self.frame_cropping else {
203            return coded;
204        };
205        let crop_x = if self.chroma_format_idc == 0 || self.separate_colour_plane_flag {
206            1
207        } else {
208            self.sub_width_c()
209        };
210        coded.saturating_sub((c.left + c.right) * crop_x)
211    }
212
213    /// Display height in luma samples after frame cropping.
214    pub fn display_height(&self) -> u32 {
215        let coded = self.coded_height();
216        let Some(c) = self.frame_cropping else {
217            return coded;
218        };
219        let crop_y = if self.chroma_format_idc == 0 || self.separate_colour_plane_flag {
220            2 - if self.frame_mbs_only_flag { 1 } else { 0 }
221        } else {
222            self.sub_height_c() * (2 - if self.frame_mbs_only_flag { 1 } else { 0 })
223        };
224        coded.saturating_sub((c.top + c.bottom) * crop_y)
225    }
226}
227
228/// Picture parameter set (7.3.2.2).
229#[derive(Debug, Clone, Default, PartialEq, Eq)]
230pub struct H264Pps {
231    pub pic_parameter_set_id: u8,
232    pub seq_parameter_set_id: u8,
233    pub entropy_coding_mode_flag: bool,
234    pub bottom_field_pic_order_in_frame_present_flag: bool,
235    pub num_slice_groups_minus1: u32,
236    pub num_ref_idx_l0_default_active_minus1: u8,
237    pub num_ref_idx_l1_default_active_minus1: u8,
238    pub weighted_pred_flag: bool,
239    pub weighted_bipred_idc: u8,
240    pub pic_init_qp_minus26: i32,
241    pub pic_init_qs_minus26: i32,
242    pub chroma_qp_index_offset: i32,
243    pub deblocking_filter_control_present_flag: bool,
244    pub constrained_intra_pred_flag: bool,
245    pub redundant_pic_cnt_present_flag: bool,
246    /// Only meaningful when the High-profile extension block is
247    /// present; otherwise defaulted to false.
248    pub transform_8x8_mode_flag: bool,
249    /// Defaults to a copy of `chroma_qp_index_offset` when the
250    /// High-profile extension block is absent (per 7.4.2.2).
251    pub second_chroma_qp_index_offset: i32,
252}
253
254/// Minimal slice header (7.3.3 + 7.4.3). Only the fields needed to
255/// (a) confirm we have an IDR access unit and (b) populate the HW
256/// backends' picture-info structs.
257#[derive(Debug, Clone, Default, PartialEq, Eq)]
258pub struct H264SliceHeader {
259    pub first_mb_in_slice: u32,
260    /// `slice_type` value as emitted; not modulo-5. 0..9.
261    pub slice_type: u8,
262    pub pic_parameter_set_id: u8,
263    pub frame_num: u32,
264    pub field_pic_flag: bool,
265    pub bottom_field_flag: bool,
266    /// Present only when this is an IDR slice (NAL type 5).
267    pub idr_pic_id: Option<u32>,
268    pub pic_order_cnt_lsb: u32,
269}
270
271impl H264SliceHeader {
272    /// True if `slice_type % 5 == 2` (per 7.4.3 — slice_type values
273    /// 2 and 7 are I-slices).
274    pub fn is_i_slice(&self) -> bool {
275        self.slice_type % 5 == 2
276    }
277}
278
279/// Convenience structure for [`parse_idr_only`].
280#[derive(Debug)]
281pub struct H264IdrParse<'a> {
282    pub sps: H264Sps,
283    pub pps: H264Pps,
284    pub slice_header: H264SliceHeader,
285    /// The raw Annex-B bytes covering the IDR access unit (from the
286    /// start code that precedes the IDR NAL through the end of the
287    /// input). The HW backends submit this slab to the GPU as-is.
288    pub idr_access_unit: &'a [u8],
289}
290
291// ─────────────────────────── Parsers ─────────────────────────────────────────
292
293/// Parse a SPS RBSP (the bytes *after* emulation-prevention stripping
294/// and *after* the NAL header byte).
295///
296/// Note: pass the RBSP payload, not the raw NAL. Use
297/// [`parse_sps_nal`] if you have a NAL with header byte at index 0.
298pub fn parse_sps(rbsp: &[u8]) -> Result<H264Sps, BitstreamError> {
299    if rbsp.len() < 3 {
300        return Err(BitstreamError::unexpected_end("SPS shorter than 3 bytes"));
301    }
302    let mut r = BitReader::new(rbsp);
303
304    let profile_idc = r.u(8) as u8;
305    let constraint_set_flags = r.u(8) as u8;
306    let level_idc = r.u(8) as u8;
307    let mut sps = H264Sps {
308        profile_idc,
309        constraint_set_flags,
310        level_idc,
311        ..H264Sps::default()
312    };
313    sps.seq_parameter_set_id = r.ue()? as u8;
314
315    if H264Sps::has_high_profile_extension(profile_idc) {
316        sps.chroma_format_idc = r.ue()? as u8;
317        if sps.chroma_format_idc == 3 {
318            sps.separate_colour_plane_flag = r.u(1) != 0;
319        }
320        sps.bit_depth_luma_minus8 = r.ue()? as u8;
321        sps.bit_depth_chroma_minus8 = r.ue()? as u8;
322        let _qpprime_y_zero_transform_bypass = r.u(1);
323        let seq_scaling_matrix_present = r.u(1);
324        if seq_scaling_matrix_present != 0 {
325            return Err(BitstreamError::unsupported(
326                "SPS seq_scaling_matrix_present_flag=1 not supported by minimal parser",
327            ));
328        }
329    } else {
330        sps.chroma_format_idc = 1; // implicit 4:2:0
331    }
332
333    sps.log2_max_frame_num_minus4 = r.ue()? as u8;
334    sps.pic_order_cnt_type = r.ue()? as u8;
335    match sps.pic_order_cnt_type {
336        0 => {
337            sps.log2_max_pic_order_cnt_lsb_minus4 = r.ue()? as u8;
338        }
339        1 => {
340            sps.delta_pic_order_always_zero_flag = r.u(1) != 0;
341            let _offset_for_non_ref_pic = r.se()?;
342            let _offset_for_top_to_bottom_field = r.se()?;
343            let num_ref_frames_in_pic_order_cnt_cycle = r.ue()?;
344            for _ in 0..num_ref_frames_in_pic_order_cnt_cycle {
345                let _offset_for_ref_frame = r.se()?;
346            }
347        }
348        2 => { /* nothing further */ }
349        other => {
350            return Err(BitstreamError::invalid(format!(
351                "SPS pic_order_cnt_type={other} (must be 0..2)"
352            )));
353        }
354    }
355    sps.max_num_ref_frames = r.ue()?;
356    sps.gaps_in_frame_num_value_allowed_flag = r.u(1) != 0;
357    sps.pic_width_in_mbs_minus1 = r.ue()?;
358    sps.pic_height_in_map_units_minus1 = r.ue()?;
359    sps.frame_mbs_only_flag = r.u(1) != 0;
360    if !sps.frame_mbs_only_flag {
361        sps.mb_adaptive_frame_field_flag = r.u(1) != 0;
362    }
363    sps.direct_8x8_inference_flag = r.u(1) != 0;
364    let frame_cropping_flag = r.u(1);
365    if frame_cropping_flag != 0 {
366        let left = r.ue()?;
367        let right = r.ue()?;
368        let top = r.ue()?;
369        let bottom = r.ue()?;
370        sps.frame_cropping = Some(H264FrameCropping {
371            left,
372            right,
373            top,
374            bottom,
375        });
376    }
377    // VUI is not parsed.
378    Ok(sps)
379}
380
381/// Parse a SPS NAL — including the NAL header byte at index 0. Errors
382/// out if the NAL type isn't 7.
383pub fn parse_sps_nal(nal: &[u8]) -> Result<H264Sps, BitstreamError> {
384    if nal.is_empty() {
385        return Err(BitstreamError::unexpected_end("empty SPS NAL"));
386    }
387    let (_, _, nal_type) = nal_header(nal[0]);
388    if nal_type != NAL_TYPE_SPS {
389        return Err(BitstreamError::invalid(format!(
390            "expected SPS NAL (type=7), got type={nal_type}"
391        )));
392    }
393    let rbsp = ebsp_to_rbsp(&nal[1..]);
394    parse_sps(&rbsp)
395}
396
397/// Parse a PPS RBSP (after emulation-prevention stripping and after
398/// the NAL header byte).
399pub fn parse_pps(rbsp: &[u8]) -> Result<H264Pps, BitstreamError> {
400    if rbsp.is_empty() {
401        return Err(BitstreamError::unexpected_end("empty PPS"));
402    }
403    let mut r = BitReader::new(rbsp);
404    let mut pps = H264Pps {
405        pic_parameter_set_id: r.ue()? as u8,
406        seq_parameter_set_id: r.ue()? as u8,
407        entropy_coding_mode_flag: r.u(1) != 0,
408        bottom_field_pic_order_in_frame_present_flag: r.u(1) != 0,
409        num_slice_groups_minus1: r.ue()?,
410        ..H264Pps::default()
411    };
412    if pps.num_slice_groups_minus1 != 0 {
413        return Err(BitstreamError::unsupported(
414            "PPS num_slice_groups_minus1>0 (FMO/ASO) not supported by minimal parser",
415        ));
416    }
417    pps.num_ref_idx_l0_default_active_minus1 = r.ue()? as u8;
418    pps.num_ref_idx_l1_default_active_minus1 = r.ue()? as u8;
419    pps.weighted_pred_flag = r.u(1) != 0;
420    pps.weighted_bipred_idc = r.u(2) as u8;
421    pps.pic_init_qp_minus26 = r.se()?;
422    pps.pic_init_qs_minus26 = r.se()?;
423    pps.chroma_qp_index_offset = r.se()?;
424    pps.deblocking_filter_control_present_flag = r.u(1) != 0;
425    pps.constrained_intra_pred_flag = r.u(1) != 0;
426    pps.redundant_pic_cnt_present_flag = r.u(1) != 0;
427
428    if r.more_rbsp_data() {
429        pps.transform_8x8_mode_flag = r.u(1) != 0;
430        let pic_scaling_matrix_present = r.u(1);
431        if pic_scaling_matrix_present != 0 {
432            return Err(BitstreamError::unsupported(
433                "PPS pic_scaling_matrix_present_flag=1 not supported by minimal parser",
434            ));
435        }
436        pps.second_chroma_qp_index_offset = r.se()?;
437    } else {
438        pps.second_chroma_qp_index_offset = pps.chroma_qp_index_offset;
439    }
440    Ok(pps)
441}
442
443/// Parse a PPS NAL — including the NAL header byte at index 0.
444pub fn parse_pps_nal(nal: &[u8]) -> Result<H264Pps, BitstreamError> {
445    if nal.is_empty() {
446        return Err(BitstreamError::unexpected_end("empty PPS NAL"));
447    }
448    let (_, _, nal_type) = nal_header(nal[0]);
449    if nal_type != NAL_TYPE_PPS {
450        return Err(BitstreamError::invalid(format!(
451            "expected PPS NAL (type=8), got type={nal_type}"
452        )));
453    }
454    let rbsp = ebsp_to_rbsp(&nal[1..]);
455    parse_pps(&rbsp)
456}
457
458/// Parse just enough of a slice header to fill the fields the HW
459/// backends need (7.3.3). The `nal_unit_type` is needed because IDR
460/// slices include `idr_pic_id` and non-IDR slices skip it.
461///
462/// `rbsp` must already be the post-emulation-prevention bytes that
463/// follow the NAL header byte. `sps` and `pps` provide the
464/// dependent-fixed-length-code widths (`log2_max_frame_num_minus4`,
465/// `pic_order_cnt_type`, etc.).
466pub fn parse_slice_header_minimal(
467    rbsp: &[u8],
468    nal_unit_type: u8,
469    sps: &H264Sps,
470    pps: &H264Pps,
471) -> Result<H264SliceHeader, BitstreamError> {
472    let _ = pps; // currently consulted only for shape; reserved for future fields
473    let mut r = BitReader::new(rbsp);
474    let mut sh = H264SliceHeader {
475        first_mb_in_slice: r.ue()?,
476        slice_type: r.ue()? as u8,
477        pic_parameter_set_id: r.ue()? as u8,
478        ..H264SliceHeader::default()
479    };
480    if sps.separate_colour_plane_flag {
481        let _colour_plane_id = r.u(2);
482    }
483    sh.frame_num = r.u(sps.log2_max_frame_num_minus4 as u32 + 4);
484    if !sps.frame_mbs_only_flag {
485        sh.field_pic_flag = r.u(1) != 0;
486        if sh.field_pic_flag {
487            sh.bottom_field_flag = r.u(1) != 0;
488        }
489    }
490    if nal_unit_type == NAL_TYPE_IDR {
491        sh.idr_pic_id = Some(r.ue()?);
492    }
493    if sps.pic_order_cnt_type == 0 {
494        sh.pic_order_cnt_lsb = r.u(sps.log2_max_pic_order_cnt_lsb_minus4 as u32 + 4);
495    }
496    // The remaining slice-header fields (ref_pic_list_modification,
497    // pred_weight_table, dec_ref_pic_marking, slice_qp_delta, …) are
498    // not needed to identify an IDR or to populate the IDR-only
499    // parameter buffer, so we stop here.
500    Ok(sh)
501}
502
503/// Walk an Annex-B stream, find the first IDR access unit, parse the
504/// SPS, PPS, and IDR slice header, and return them along with a slice
505/// of the original input covering the IDR access unit.
506///
507/// Errors out if the stream contains no SPS, no PPS, no IDR NAL, or
508/// any of the parsers fails.
509pub fn parse_idr_only(stream: &[u8]) -> Result<H264IdrParse<'_>, BitstreamError> {
510    // Locate the byte ranges of every Annex-B NAL — we need indices
511    // into `stream` (not just borrows of NAL bodies) to slice the
512    // returned `idr_access_unit`.
513    let nals = locate_annex_b(stream);
514    if nals.is_empty() {
515        return Err(BitstreamError::invalid("no NAL units found in stream"));
516    }
517
518    let mut sps: Option<H264Sps> = None;
519    let mut pps: Option<H264Pps> = None;
520    let mut idr_idx: Option<usize> = None;
521    for (idx, n) in nals.iter().enumerate() {
522        let body = &stream[n.body_start..n.body_end];
523        if body.is_empty() {
524            continue;
525        }
526        let (_, _, nal_type) = nal_header(body[0]);
527        match nal_type {
528            NAL_TYPE_SPS if sps.is_none() => {
529                sps = Some(parse_sps_nal(body)?);
530            }
531            NAL_TYPE_PPS if pps.is_none() => {
532                pps = Some(parse_pps_nal(body)?);
533            }
534            NAL_TYPE_IDR if idr_idx.is_none() => {
535                idr_idx = Some(idx);
536            }
537            _ => {}
538        }
539    }
540    let sps = sps.ok_or_else(|| BitstreamError::invalid("stream has no SPS NAL"))?;
541    let pps = pps.ok_or_else(|| BitstreamError::invalid("stream has no PPS NAL"))?;
542    let idr_idx = idr_idx.ok_or_else(|| BitstreamError::invalid("stream has no IDR NAL"))?;
543
544    let idr_nal = &nals[idr_idx];
545    let body = &stream[idr_nal.body_start..idr_nal.body_end];
546    let rbsp = ebsp_to_rbsp(&body[1..]);
547    let slice_header = parse_slice_header_minimal(&rbsp, NAL_TYPE_IDR, &sps, &pps)?;
548
549    let access_unit = &stream[idr_nal.start_code_start..];
550    Ok(H264IdrParse {
551        sps,
552        pps,
553        slice_header,
554        idr_access_unit: access_unit,
555    })
556}
557
558// ─────────────────────────── Annex-B index helper ────────────────────────────
559
560#[derive(Debug, Clone, Copy)]
561struct NalLoc {
562    /// Index of the first byte of the start code (00 00 [00] 01).
563    start_code_start: usize,
564    body_start: usize,
565    body_end: usize,
566}
567
568fn locate_annex_b(buf: &[u8]) -> Vec<NalLoc> {
569    let mut out = Vec::new();
570    let mut i = 0;
571    let n = buf.len();
572    let mut current: Option<(usize, usize)> = None; // (start_code_start, body_start)
573    while i < n {
574        let four =
575            i + 3 < n && buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 0 && buf[i + 3] == 1;
576        let three = !four && i + 2 < n && buf[i] == 0 && buf[i + 1] == 0 && buf[i + 2] == 1;
577        if four || three {
578            if let Some((sc, body_start)) = current.take() {
579                out.push(NalLoc {
580                    start_code_start: sc,
581                    body_start,
582                    body_end: i,
583                });
584            }
585            let sc = i;
586            i += if four { 4 } else { 3 };
587            current = Some((sc, i));
588            continue;
589        }
590        i += 1;
591    }
592    if let Some((sc, body_start)) = current.take() {
593        out.push(NalLoc {
594            start_code_start: sc,
595            body_start,
596            body_end: n,
597        });
598    }
599    out
600}
601
602// ─────────────────────────── Tests ───────────────────────────────────────────
603
604#[cfg(test)]
605mod tests {
606    use super::*;
607
608    #[test]
609    fn split_finds_three_nal_starts_in_synthetic_buffer() {
610        let buf = [
611            0, 0, 0, 1, 0x67, 0xaa, // NAL 1
612            0, 0, 1, 0x68, 0xbb, // NAL 2
613            0, 0, 0, 1, 0x65, 0xcc, 0xdd, // NAL 3
614        ];
615        let nals = split_annex_b(&buf);
616        assert_eq!(nals.len(), 3);
617        assert_eq!(nals[0], &[0x67, 0xaa]);
618        assert_eq!(nals[1], &[0x68, 0xbb]);
619        assert_eq!(nals[2], &[0x65, 0xcc, 0xdd]);
620    }
621
622    #[test]
623    fn ebsp_to_rbsp_strips_03_in_zero_zero_run() {
624        let ebsp = [0x00, 0x00, 0x03, 0x01, 0x02, 0x03];
625        let rbsp = ebsp_to_rbsp(&ebsp);
626        assert_eq!(rbsp, &[0x00, 0x00, 0x01, 0x02, 0x03]);
627    }
628
629    #[test]
630    fn nal_header_decodes_idr_byte() {
631        // 0x65 = 0110_0101 → forbidden=0, nal_ref_idc=3, nal_type=5
632        let (f, nri, t) = nal_header(0x65);
633        assert_eq!(f, 0);
634        assert_eq!(nri, 3);
635        assert_eq!(t, NAL_TYPE_IDR);
636    }
637}
oxideav_bitstream/h264.rs

oxideav_bitstream/
h264.rs
