From 534eca502c3abf8c63180659eba76c402b452b1b Mon Sep 17 00:00:00 2001 From: Matthew Jackson <1085847+MattJackson@users.noreply.github.com> Date: Thu, 25 Jun 2026 20:38:09 -0700 Subject: [PATCH] mux/h264: measure coding type from slice_type; carry source provenance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Honest PictureInfo population for H.264 — the foundation that lets the video index claim a frame's type and be believed (a faked "P" on a B-frame would make the index confidently wrong). - startcode: shared minimal MSB-first BitReader with Exp-Golomb ue(v), reused by H.264 (and next HEVC). Documents the emulation-prevention caveat: only the leading slice-header fields are read, where 00 00 03 cannot intervene. - h264: decode first_mb_in_slice + slice_type (H.264 §7.3.3) from the first coded slice and map to I/P/B (§7.4.3 Table 7-6; SP→P, SI→I). Set coding = PictureInfo::coding_type_only(...) and source = pes.source. Field order is NOT decoded here, so field_order() stays None — honestly absent, never guessed. - Tests: I/P/B from real ue-encoded slice headers, source carry, field-order absence; BitReader Exp-Golomb table + truncation. --- src/mux/codec/h264.rs | 90 +++++++++++++++++++++++++++++++++-- src/mux/codec/startcode.rs | 96 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 183 insertions(+), 3 deletions(-) diff --git a/src/mux/codec/h264.rs b/src/mux/codec/h264.rs index 167c92f..603a7a5 100644 --- a/src/mux/codec/h264.rs +++ b/src/mux/codec/h264.rs @@ -4,15 +4,32 @@ //! Detects keyframes (IDR slices). //! Each PES packet = one access unit = one frame. -use super::startcode::{find_start_code, skip_start_code}; +use super::coding::{CodingType, PictureInfo}; +use super::startcode::{BitReader, find_start_code, skip_start_code}; use super::{CodecParser, Frame, PesPacket, pts_to_ns}; /// H.264 NAL unit types we care about. +const NAL_SLICE_NON_IDR: u8 = 1; const NAL_SLICE_IDR: u8 = 5; const NAL_SPS: u8 = 7; const NAL_PPS: u8 = 8; const NAL_AUD: u8 = 9; +/// Map an H.264 `slice_type` (Rec. ITU-T H.264 §7.4.3, Table 7-6) to a coding +/// type. Values 5..=9 repeat 0..=4 (the "all slices of this type" forms), so +/// `slice_type % 5`: 0 = P, 1 = B, 2 = I, 3 = SP (predicted → P), 4 = SI +/// (intra → I). Returns `None` for values outside 0..=9 (malformed header). +fn h264_slice_coding_type(slice_type: u32) -> Option { + match slice_type { + 0..=9 => Some(match slice_type % 5 { + 0 | 3 => CodingType::P, // P, SP + 1 => CodingType::B, + _ => CodingType::I, // 2 = I, 4 = SI + }), + _ => None, + } +} + /// H.264 (AVC) Annex B → MKV codec parser: extracts SPS/PPS for the avcC /// codecPrivate, detects IDR keyframes, and converts each PES access unit into /// length-prefixed NAL units. Implements [`CodecParser`]. @@ -141,6 +158,8 @@ impl CodecParser for H264Parser { // Annex B (start-code prefixed) NALUs to length-prefixed NALUs (MKV with // AVCDecoderConfigurationRecord expects a 4-byte length prefix per NAL). let mut keyframe = false; + // Picture coding type, MEASURED from the first coded slice's header. + let mut coding_type: Option = None; // Did this access unit already carry each param-set type in-band? let mut emitted_sps = false; let mut emitted_pps = false; @@ -172,6 +191,21 @@ impl CodecParser for H264Parser { if nal_type == NAL_SLICE_IDR { keyframe = true; } + // Measure the coding type from the FIRST coded slice's header + // (H.264 §7.3.3: first_mb_in_slice ue(v), then slice_type + // ue(v)). Populates PictureInfo so a consumer reads a MEASURED + // I/P/B, never a keyframe-only guess. Only the first slice of + // the access unit is read; `nal[1..]` is the RBSP after the + // 1-byte NAL header (slice_type is too early for an + // emulation-prevention byte to intervene). + if (nal_type == NAL_SLICE_NON_IDR || nal_type == NAL_SLICE_IDR) + && coding_type.is_none() + { + let mut br = BitReader::new(&nal[1..]); + if let (Some(_first_mb), Some(slice_type)) = (br.read_ue(), br.read_ue()) { + coding_type = h264_slice_coding_type(slice_type); + } + } // A NAL longer than u32::MAX can't be length-prefixed in the // 4-byte field; skip it rather than mis-frame the output. // Unreachable in practice (no real AU is >4 GiB). @@ -203,8 +237,10 @@ impl CodecParser for H264Parser { } vec![Frame { - coding: None, - source: None, + // Coding-type only: H.264 field order is not decoded here, so + // `field_order()` stays `None` — honestly absent, never guessed. + coding: coding_type.map(PictureInfo::coding_type_only), + source: pes.source, pts_ns, keyframe, data: frame_data, @@ -573,6 +609,54 @@ mod tests { v } + #[test] + fn h264_populates_measured_coding_type_and_source() { + use super::super::coding::CodingType; + // Slice-header body = first_mb_in_slice=0 ('1') then slice_type ue(v): + // 0x88 = '1 0001000' → slice_type 7 (I) + // 0x98 = '1 00110..' → slice_type 5 (P) + // 0x9C = '1 00111..' → slice_type 6 (B) + let src = crate::pes::SourcePos::at_byte(8192); + let mut parse = |nal_type: u8, body: u8| { + let mut p = H264Parser::new(); + let mut pe = make_pes(h264_nal(nal_type, &[body]), Some(0)); + pe.source = Some(src); + p.parse(&pe) + }; + + // IDR carrying an I-slice → keyframe + MEASURED I; source carried; H.264 + // field order is not decoded, so it is honestly absent (not guessed). + let fi = parse(NAL_SLICE_IDR, 0x88); + assert_eq!(fi.len(), 1); + assert!(fi[0].keyframe, "IDR is a keyframe"); + let ci = fi[0].coding.expect("H.264 frame carries PictureInfo"); + assert_eq!(ci.coding_type(), CodingType::I, "slice_type 7 → I"); + assert!( + ci.field_order().is_none(), + "H.264 field order undecoded → None, never faked" + ); + assert_eq!( + fi[0].source.unwrap().byte, + 8192, + "source provenance carried" + ); + + // Non-IDR P / B slices → MEASURED P / B, not keyframes. + let fp = parse(NAL_SLICE_NON_IDR, 0x98); + assert_eq!( + fp[0].coding.unwrap().coding_type(), + CodingType::P, + "slice_type 5 → P" + ); + assert!(!fp[0].keyframe); + let fb = parse(NAL_SLICE_NON_IDR, 0x9C); + assert_eq!( + fb[0].coding.unwrap().coding_type(), + CodingType::B, + "slice_type 6 → B" + ); + } + /// Regression (Fight Club bug, H.264 variant): PPS id 0 = body A (→ avcC), /// redefined to B, then switched BACK to A. A streaming decoder is on B; the /// revert to A == avcC must still be emitted in-band or the A-segment diff --git a/src/mux/codec/startcode.rs b/src/mux/codec/startcode.rs index cdb801d..af2522b 100644 --- a/src/mux/codec/startcode.rs +++ b/src/mux/codec/startcode.rs @@ -38,10 +38,106 @@ pub fn skip_start_code(data: &[u8], pos: usize) -> Option { None } +/// Minimal MSB-first bit reader over an RBSP, for the leading fields of a coded +/// slice header (H.264 `first_mb_in_slice` + `slice_type`; HEVC +/// `slice_segment_header`). +/// +/// It does NOT remove emulation-prevention bytes (`00 00 03`). Those can only +/// appear after two consecutive `0x00` bytes, which cannot occur within the +/// first Exp-Golomb codes of a slice header (a slice header never begins +/// `00 00`), so the leading fields this reader is used for decode correctly. A +/// caller reading deep enough into a header that `00 00 03` could appear must +/// de-emulate the RBSP first. +pub(crate) struct BitReader<'a> { + data: &'a [u8], + bit: usize, +} + +impl<'a> BitReader<'a> { + /// Reader positioned at the first bit of `data`. + pub fn new(data: &'a [u8]) -> Self { + Self { data, bit: 0 } + } + + /// Read a single bit, MSB-first. `None` once the buffer is exhausted. + pub fn read_bit(&mut self) -> Option { + let byte = self.bit / 8; + if byte >= self.data.len() { + return None; + } + let b = (self.data[byte] >> (7 - (self.bit & 7))) & 1; + self.bit += 1; + Some(b as u32) + } + + /// Skip `n` bits; `None` if that would run past the end. + pub fn skip_bits(&mut self, n: usize) -> Option<()> { + for _ in 0..n { + self.read_bit()?; + } + Some(()) + } + + /// Read an unsigned Exp-Golomb code `ue(v)` (H.264 §9.1 / HEVC §9.2): + /// count leading zeros, read the `1` stop bit, then that many info bits; + /// `code_num = 2^leadingZeros - 1 + info`. `None` on truncation or an + /// absurdly long code (>31 leading zeros — malformed input, not a real + /// slice header). + pub fn read_ue(&mut self) -> Option { + let mut leading_zeros = 0u32; + while self.read_bit()? == 0 { + leading_zeros += 1; + if leading_zeros > 31 { + return None; + } + } + let mut info = 0u32; + for _ in 0..leading_zeros { + info = (info << 1) | self.read_bit()?; + } + Some((1u32 << leading_zeros) - 1 + info) + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn bit_reader_read_ue_exp_golomb_table() { + // ue(v) codes from H.264 Table 9-1: code_num 0='1', 1='010', 2='011', + // 3='00100'. Each crafted byte is left-aligned (MSB-first). + assert_eq!(BitReader::new(&[0x80]).read_ue(), Some(0)); // 1_______ + assert_eq!(BitReader::new(&[0x40]).read_ue(), Some(1)); // 010_____ + assert_eq!(BitReader::new(&[0x60]).read_ue(), Some(2)); // 011_____ + assert_eq!(BitReader::new(&[0x20]).read_ue(), Some(3)); // 00100___ + assert_eq!(BitReader::new(&[0x28]).read_ue(), Some(4)); // 00101___ + } + + #[test] + fn bit_reader_read_ue_sequence_and_bits() { + // '1' '011' '00101' = ue(0), ue(2), ue(4) across the bitstream. + // 1 011 00101 -> 1011 0010 1 -> 0xB2, 0x80. + let mut br = BitReader::new(&[0xB2, 0x80]); + assert_eq!(br.read_ue(), Some(0)); + assert_eq!(br.read_ue(), Some(2)); + assert_eq!(br.read_ue(), Some(4)); + } + + #[test] + fn bit_reader_truncation_and_skip() { + // Empty buffer → None, no panic. + assert_eq!(BitReader::new(&[]).read_ue(), None); + // skip_bits past the end → None. + let mut br = BitReader::new(&[0xFF]); + assert_eq!(br.skip_bits(9), None); + // read_bit MSB-first. + let mut b = BitReader::new(&[0b1010_0000]); + assert_eq!(b.read_bit(), Some(1)); + assert_eq!(b.read_bit(), Some(0)); + assert_eq!(b.read_bit(), Some(1)); + } + #[test] fn find_start_code_3byte() { let data = [0x00, 0x00, 0x01, 0x65];