mux/h264: measure coding type from slice_type; carry source provenance

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.
This commit is contained in:
Matthew Jackson
2026-06-25 20:38:09 -07:00
parent e3dbafcebd
commit 534eca502c
2 changed files with 183 additions and 3 deletions
+87 -3
View File
@@ -4,15 +4,32 @@
//! Detects keyframes (IDR slices). //! Detects keyframes (IDR slices).
//! Each PES packet = one access unit = one frame. //! 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}; use super::{CodecParser, Frame, PesPacket, pts_to_ns};
/// H.264 NAL unit types we care about. /// H.264 NAL unit types we care about.
const NAL_SLICE_NON_IDR: u8 = 1;
const NAL_SLICE_IDR: u8 = 5; const NAL_SLICE_IDR: u8 = 5;
const NAL_SPS: u8 = 7; const NAL_SPS: u8 = 7;
const NAL_PPS: u8 = 8; const NAL_PPS: u8 = 8;
const NAL_AUD: u8 = 9; 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<CodingType> {
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 /// 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 /// codecPrivate, detects IDR keyframes, and converts each PES access unit into
/// length-prefixed NAL units. Implements [`CodecParser`]. /// 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 // Annex B (start-code prefixed) NALUs to length-prefixed NALUs (MKV with
// AVCDecoderConfigurationRecord expects a 4-byte length prefix per NAL). // AVCDecoderConfigurationRecord expects a 4-byte length prefix per NAL).
let mut keyframe = false; let mut keyframe = false;
// Picture coding type, MEASURED from the first coded slice's header.
let mut coding_type: Option<CodingType> = None;
// Did this access unit already carry each param-set type in-band? // Did this access unit already carry each param-set type in-band?
let mut emitted_sps = false; let mut emitted_sps = false;
let mut emitted_pps = false; let mut emitted_pps = false;
@@ -172,6 +191,21 @@ impl CodecParser for H264Parser {
if nal_type == NAL_SLICE_IDR { if nal_type == NAL_SLICE_IDR {
keyframe = true; 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 // A NAL longer than u32::MAX can't be length-prefixed in the
// 4-byte field; skip it rather than mis-frame the output. // 4-byte field; skip it rather than mis-frame the output.
// Unreachable in practice (no real AU is >4 GiB). // Unreachable in practice (no real AU is >4 GiB).
@@ -203,8 +237,10 @@ impl CodecParser for H264Parser {
} }
vec![Frame { vec![Frame {
coding: None, // Coding-type only: H.264 field order is not decoded here, so
source: None, // `field_order()` stays `None` — honestly absent, never guessed.
coding: coding_type.map(PictureInfo::coding_type_only),
source: pes.source,
pts_ns, pts_ns,
keyframe, keyframe,
data: frame_data, data: frame_data,
@@ -573,6 +609,54 @@ mod tests {
v 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), /// 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 /// 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 /// revert to A == avcC must still be emitted in-band or the A-segment
+96
View File
@@ -38,10 +38,106 @@ pub fn skip_start_code(data: &[u8], pos: usize) -> Option<usize> {
None 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<u32> {
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<u32> {
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)] #[cfg(test)]
mod tests { mod tests {
use super::*; 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] #[test]
fn find_start_code_3byte() { fn find_start_code_3byte() {
let data = [0x00, 0x00, 0x01, 0x65]; let data = [0x00, 0x00, 0x01, 0x65];