MattJackson
1b95193517
- MkvTrack::audio emits A_DTS/MA, A_DTS/HR, A_DTS per the DTS family instead of mislabelling everything as A_DTS. Plex transcoder and strict hardware decoders reject DTS-HD MA payload under a plain A_DTS track. - PgsParser is now stateful: pairs display PCS with the following empty PCS to compute a duration. Frame::duration_ns + PesFrame::duration_ns carry it through; MkvMuxer::write_frame gains a final Option<u64> parameter that emits BlockGroup + BlockDuration when set. Fixes subtitle bitmaps lingering past their intended end-time.
477 lines
14 KiB
Rust
477 lines
14 KiB
Rust
//! H.264 (AVC) elementary stream parser.
|
||
//!
|
||
//! Extracts SPS and PPS NAL units for MKV codecPrivate.
|
||
//! Detects keyframes (IDR slices).
|
||
//! Each PES packet = one access unit = one frame.
|
||
|
||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||
|
||
/// H.264 NAL unit types we care about.
|
||
const NAL_SLICE_IDR: u8 = 5;
|
||
const NAL_SPS: u8 = 7;
|
||
const NAL_PPS: u8 = 8;
|
||
const NAL_AUD: u8 = 9;
|
||
|
||
pub struct H264Parser {
|
||
sps: Option<Vec<u8>>,
|
||
pps: Option<Vec<u8>>,
|
||
}
|
||
|
||
impl Default for H264Parser {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|
||
|
||
impl H264Parser {
|
||
pub fn new() -> Self {
|
||
Self {
|
||
sps: None,
|
||
pps: None,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl CodecParser for H264Parser {
|
||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||
if pes.data.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
// Use DTS when available (monotonic for B-frame content), fall back to PTS
|
||
let pts_ns = pes.dts.or(pes.pts).map(pts_to_ns).unwrap_or(0);
|
||
|
||
// Scan NAL units for SPS, PPS, and IDR detection
|
||
let mut keyframe = false;
|
||
let mut frame_data = Vec::new();
|
||
|
||
for nal in NalIterator::new(&pes.data) {
|
||
let nal_type = nal[0] & 0x1F;
|
||
|
||
match nal_type {
|
||
NAL_SPS => {
|
||
self.sps = Some(nal.to_vec());
|
||
}
|
||
NAL_PPS => {
|
||
self.pps = Some(nal.to_vec());
|
||
}
|
||
NAL_SLICE_IDR => {
|
||
keyframe = true;
|
||
}
|
||
_ => {}
|
||
}
|
||
}
|
||
|
||
// Convert Annex B (start code prefixed) to length-prefixed NALUs.
|
||
// MKV with AVCDecoderConfigurationRecord expects 4-byte length prefix per NALU.
|
||
// Skip SPS/PPS/AUD NALUs — they're in codecPrivate, not in frame data.
|
||
for nal in NalIterator::new(&pes.data) {
|
||
let nal_type = nal[0] & 0x1F;
|
||
// Skip parameter sets and access unit delimiters
|
||
if nal_type == NAL_SPS || nal_type == NAL_PPS || nal_type == NAL_AUD {
|
||
continue;
|
||
}
|
||
// 4-byte big-endian length prefix
|
||
let len = nal.len() as u32;
|
||
frame_data.extend_from_slice(&len.to_be_bytes());
|
||
frame_data.extend_from_slice(nal);
|
||
}
|
||
|
||
if frame_data.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
|
||
vec![Frame {
|
||
pts_ns,
|
||
keyframe,
|
||
data: frame_data,
|
||
duration_ns: None,
|
||
}]
|
||
}
|
||
|
||
fn codec_private(&self) -> Option<Vec<u8>> {
|
||
// Build AVCDecoderConfigurationRecord from SPS + PPS
|
||
let sps = self.sps.as_ref()?;
|
||
let pps = self.pps.as_ref()?;
|
||
|
||
if sps.len() < 4 {
|
||
return None;
|
||
}
|
||
|
||
// AVCDecoderConfigurationRecord (ISO 14496-15):
|
||
// configurationVersion = 1
|
||
// AVCProfileIndication = SPS[1]
|
||
// profile_compatibility = SPS[2]
|
||
// AVCLevelIndication = SPS[3]
|
||
// lengthSizeMinusOne = 3 (4-byte length prefix)
|
||
// numOfSequenceParameterSets = 1
|
||
// sequenceParameterSetLength = sps.len()
|
||
// sequenceParameterSetNALUnit = sps
|
||
// numOfPictureParameterSets = 1
|
||
// pictureParameterSetLength = pps.len()
|
||
// pictureParameterSetNALUnit = pps
|
||
|
||
let mut record = vec![
|
||
1, // configurationVersion
|
||
sps[1], // profile
|
||
sps[2], // compatibility
|
||
sps[3], // level
|
||
0xFF, // 6 bits reserved (111111) + 2 bits lengthSizeMinusOne (11 = 3)
|
||
0xE1, // 3 bits reserved (111) + 5 bits numSPS (1)
|
||
(sps.len() >> 8) as u8,
|
||
sps.len() as u8,
|
||
];
|
||
record.extend_from_slice(sps);
|
||
record.push(1); // numPPS
|
||
record.push((pps.len() >> 8) as u8);
|
||
record.push(pps.len() as u8);
|
||
record.extend_from_slice(pps);
|
||
|
||
Some(record)
|
||
}
|
||
}
|
||
|
||
/// Iterator over NAL units in Annex B byte stream.
|
||
/// Finds start codes (00 00 01 or 00 00 00 01) and yields the data between them.
|
||
struct NalIterator<'a> {
|
||
data: &'a [u8],
|
||
pos: usize,
|
||
}
|
||
|
||
impl<'a> NalIterator<'a> {
|
||
fn new(data: &'a [u8]) -> Self {
|
||
// Skip to first start code
|
||
let pos = find_start_code(data, 0).unwrap_or(data.len());
|
||
Self { data, pos }
|
||
}
|
||
}
|
||
|
||
impl<'a> Iterator for NalIterator<'a> {
|
||
type Item = &'a [u8];
|
||
|
||
fn next(&mut self) -> Option<&'a [u8]> {
|
||
if self.pos >= self.data.len() {
|
||
return None;
|
||
}
|
||
|
||
// Skip the start code at current position
|
||
let nal_start = skip_start_code(self.data, self.pos)?;
|
||
|
||
// Find next start code (or end of data)
|
||
let nal_end = find_start_code(self.data, nal_start).unwrap_or(self.data.len());
|
||
|
||
// Remove trailing zeros (part of next start code's zero prefix)
|
||
let mut end = nal_end;
|
||
while end > nal_start && self.data[end - 1] == 0x00 {
|
||
end -= 1;
|
||
}
|
||
|
||
self.pos = nal_end;
|
||
|
||
if end > nal_start {
|
||
Some(&self.data[nal_start..end])
|
||
} else {
|
||
self.next()
|
||
}
|
||
}
|
||
}
|
||
|
||
/// Find the position of the next start code (00 00 01) at or after `from`.
|
||
///
|
||
/// Backed by `memchr::memmem::find` for SIMD-accelerated bytestring
|
||
/// search. On AVX2-capable x86_64 this runs ~5–10× the byte-by-byte
|
||
/// scan that preceded it; on a 200 KB UHD HEVC frame the saving is
|
||
/// in the hundreds of microseconds per call.
|
||
pub fn find_start_code(data: &[u8], from: usize) -> Option<usize> {
|
||
if data.len() < from + 3 {
|
||
return None;
|
||
}
|
||
memchr::memmem::find(&data[from..], b"\x00\x00\x01").map(|rel| from + rel)
|
||
}
|
||
|
||
/// Skip past the start code at position `pos`, returning the first byte after it.
|
||
pub fn skip_start_code(data: &[u8], pos: usize) -> Option<usize> {
|
||
if pos + 2 >= data.len() {
|
||
return None;
|
||
}
|
||
if data[pos] == 0x00 && data[pos + 1] == 0x00 {
|
||
if pos + 3 < data.len() && data[pos + 2] == 0x00 && data[pos + 3] == 0x01 {
|
||
return Some(pos + 4); // 4-byte start code
|
||
}
|
||
if data[pos + 2] == 0x01 {
|
||
return Some(pos + 3); // 3-byte start code
|
||
}
|
||
}
|
||
None
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use crate::mux::ts::PesPacket;
|
||
|
||
fn make_pes(data: Vec<u8>, pts: Option<i64>) -> PesPacket {
|
||
PesPacket {
|
||
pid: 0x1011,
|
||
pts,
|
||
dts: None,
|
||
data,
|
||
}
|
||
}
|
||
|
||
// --- find_start_code tests ---
|
||
|
||
#[test]
|
||
fn find_start_code_3byte() {
|
||
let data = [0x00, 0x00, 0x01, 0x65];
|
||
assert_eq!(find_start_code(&data, 0), Some(0));
|
||
}
|
||
|
||
#[test]
|
||
fn find_start_code_4byte() {
|
||
let data = [0x00, 0x00, 0x00, 0x01, 0x65];
|
||
// find_start_code looks for 00 00 01 pattern, which starts at offset 1 in a 4-byte start code
|
||
assert_eq!(find_start_code(&data, 0), Some(1));
|
||
}
|
||
|
||
#[test]
|
||
fn find_start_code_offset() {
|
||
let data = [0xFF, 0xFF, 0x00, 0x00, 0x01, 0x09];
|
||
assert_eq!(find_start_code(&data, 0), Some(2));
|
||
}
|
||
|
||
#[test]
|
||
fn find_start_code_none() {
|
||
let data = [0x00, 0x00, 0x00, 0x00];
|
||
assert_eq!(find_start_code(&data, 0), None);
|
||
}
|
||
|
||
#[test]
|
||
fn find_start_code_too_short() {
|
||
let data = [0x00, 0x00];
|
||
assert_eq!(find_start_code(&data, 0), None);
|
||
}
|
||
|
||
// --- parse SPS+PPS → codec_private ---
|
||
|
||
#[test]
|
||
fn parse_sps_pps() {
|
||
let mut parser = H264Parser::new();
|
||
|
||
// Build PES with SPS (type 7) + PPS (type 8) + IDR slice (type 5)
|
||
// SPS NAL: 0x67 = 0_11_00111 (nal_type = 7), followed by profile/compat/level + payload
|
||
// PPS NAL: 0x68 = 0_11_01000 (nal_type = 8)
|
||
let mut data = Vec::new();
|
||
// SPS: 00 00 01 [67 42 00 1E <payload>]
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x67); // SPS
|
||
data.extend_from_slice(&[0x42, 0x00, 0x1E, 0xAB, 0xCD]); // profile=0x42, compat=0x00, level=0x1E
|
||
// PPS: 00 00 01 [68 <payload>]
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x68); // PPS
|
||
data.extend_from_slice(&[0xCE, 0x01]);
|
||
// IDR slice: 00 00 01 [65 <payload>]
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x65); // IDR
|
||
data.extend_from_slice(&[0x88, 0x00, 0x10]);
|
||
|
||
let pes = make_pes(data, Some(90000));
|
||
let frames = parser.parse(&pes);
|
||
|
||
// codec_private should now be available
|
||
let cp = parser.codec_private();
|
||
assert!(
|
||
cp.is_some(),
|
||
"codec_private should be Some after seeing SPS+PPS"
|
||
);
|
||
let cp = cp.unwrap();
|
||
|
||
// AVCDecoderConfigurationRecord checks
|
||
assert_eq!(cp[0], 1, "configurationVersion");
|
||
assert_eq!(cp[1], 0x42, "profile from SPS[1]");
|
||
assert_eq!(cp[2], 0x00, "compatibility from SPS[2]");
|
||
assert_eq!(cp[3], 0x1E, "level from SPS[3]");
|
||
assert_eq!(cp[4], 0xFF, "reserved + lengthSizeMinusOne=3");
|
||
assert_eq!(cp[5], 0xE1, "reserved + numSPS=1");
|
||
|
||
// Frames should have been produced
|
||
assert_eq!(frames.len(), 1);
|
||
}
|
||
|
||
#[test]
|
||
fn codec_private_none_before_sps_pps() {
|
||
let parser = H264Parser::new();
|
||
assert!(parser.codec_private().is_none());
|
||
}
|
||
|
||
// --- IDR keyframe detection ---
|
||
|
||
#[test]
|
||
fn parse_idr_keyframe() {
|
||
let mut parser = H264Parser::new();
|
||
|
||
// PES with IDR NAL (type 5 = 0x65)
|
||
let mut data = Vec::new();
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x65); // IDR slice (nal_type = 5)
|
||
data.extend_from_slice(&[0x88, 0x00, 0x10, 0x20]);
|
||
|
||
let pes = make_pes(data, Some(90000));
|
||
let frames = parser.parse(&pes);
|
||
|
||
assert_eq!(frames.len(), 1);
|
||
assert!(
|
||
frames[0].keyframe,
|
||
"IDR slice should be detected as keyframe"
|
||
);
|
||
}
|
||
|
||
// --- non-IDR → not keyframe ---
|
||
|
||
#[test]
|
||
fn parse_non_idr() {
|
||
let mut parser = H264Parser::new();
|
||
|
||
// PES with non-IDR slice (type 1 = 0x61 or 0x41)
|
||
let mut data = Vec::new();
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x41); // non-IDR coded slice (nal_type = 1)
|
||
data.extend_from_slice(&[0x9A, 0x00, 0x10]);
|
||
|
||
let pes = make_pes(data, Some(180000));
|
||
let frames = parser.parse(&pes);
|
||
|
||
assert_eq!(frames.len(), 1);
|
||
assert!(!frames[0].keyframe, "non-IDR slice should not be keyframe");
|
||
}
|
||
|
||
// --- length prefix conversion ---
|
||
|
||
#[test]
|
||
fn length_prefix_conversion() {
|
||
let mut parser = H264Parser::new();
|
||
|
||
// PES with a single non-IDR NAL
|
||
let nal_payload = [0x41, 0xAA, 0xBB, 0xCC, 0xDD]; // type 1, 5 bytes
|
||
let mut data = Vec::new();
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.extend_from_slice(&nal_payload);
|
||
|
||
let pes = make_pes(data, Some(0));
|
||
let frames = parser.parse(&pes);
|
||
|
||
assert_eq!(frames.len(), 1);
|
||
let frame_data = &frames[0].data;
|
||
|
||
// Should start with 4-byte big-endian length prefix
|
||
assert!(
|
||
frame_data.len() >= 4,
|
||
"frame data should have length prefix"
|
||
);
|
||
let length =
|
||
u32::from_be_bytes([frame_data[0], frame_data[1], frame_data[2], frame_data[3]]);
|
||
assert_eq!(
|
||
length as usize,
|
||
nal_payload.len(),
|
||
"length prefix should match NAL size"
|
||
);
|
||
|
||
// Followed by the NAL data itself
|
||
assert_eq!(&frame_data[4..], &nal_payload);
|
||
|
||
// No start code (00 00 01) should appear in the output
|
||
for i in 0..frame_data.len().saturating_sub(2) {
|
||
let is_sc =
|
||
frame_data[i] == 0x00 && frame_data[i + 1] == 0x00 && frame_data[i + 2] == 0x01;
|
||
assert!(!is_sc, "output should not contain Annex B start codes");
|
||
}
|
||
}
|
||
|
||
// --- SPS/PPS/AUD are stripped from frame data ---
|
||
|
||
#[test]
|
||
fn sps_pps_aud_stripped_from_frame_data() {
|
||
let mut parser = H264Parser::new();
|
||
|
||
let mut data = Vec::new();
|
||
// AUD (type 9)
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x09);
|
||
data.push(0xF0);
|
||
// SPS (type 7)
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x67);
|
||
data.extend_from_slice(&[0x42, 0x00, 0x1E, 0xAB]);
|
||
// PPS (type 8)
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x68);
|
||
data.extend_from_slice(&[0xCE, 0x01]);
|
||
// IDR (type 5) - only this should appear in frame data
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x65);
|
||
data.extend_from_slice(&[0x88, 0x00]);
|
||
|
||
let pes = make_pes(data, Some(0));
|
||
let frames = parser.parse(&pes);
|
||
assert_eq!(frames.len(), 1);
|
||
|
||
// Frame data should only contain the IDR NAL (length-prefixed)
|
||
let fd = &frames[0].data;
|
||
let length = u32::from_be_bytes([fd[0], fd[1], fd[2], fd[3]]);
|
||
// IDR NAL is 0x65, 0x88 (trailing 0x00 is stripped as potential start code prefix)
|
||
assert_eq!(length, 2);
|
||
assert_eq!(fd[4], 0x65); // IDR NAL type byte
|
||
}
|
||
|
||
// --- PTS conversion ---
|
||
|
||
#[test]
|
||
fn pts_conversion() {
|
||
let mut parser = H264Parser::new();
|
||
|
||
let mut data = Vec::new();
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x41);
|
||
data.extend_from_slice(&[0x00, 0x10]);
|
||
|
||
// PTS = 90000 (1 second at 90kHz) → 1_000_000_000 ns
|
||
let pes = make_pes(data, Some(90000));
|
||
let frames = parser.parse(&pes);
|
||
assert_eq!(frames.len(), 1);
|
||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||
}
|
||
|
||
// --- empty PES ---
|
||
|
||
#[test]
|
||
fn parse_empty_pes() {
|
||
let mut parser = H264Parser::new();
|
||
let pes = make_pes(Vec::new(), Some(0));
|
||
let frames = parser.parse(&pes);
|
||
assert!(frames.is_empty());
|
||
}
|
||
|
||
// --- DTS preferred over PTS when present ---
|
||
|
||
#[test]
|
||
fn dts_preferred_over_pts() {
|
||
let mut parser = H264Parser::new();
|
||
|
||
let mut data = Vec::new();
|
||
data.extend_from_slice(&[0x00, 0x00, 0x01]);
|
||
data.push(0x41);
|
||
data.extend_from_slice(&[0x00, 0x10]);
|
||
|
||
let pes = PesPacket {
|
||
pid: 0x1011,
|
||
pts: Some(180000), // 2 seconds
|
||
dts: Some(90000), // 1 second
|
||
data,
|
||
};
|
||
let frames = parser.parse(&pes);
|
||
assert_eq!(frames.len(), 1);
|
||
// DTS should be used, not PTS
|
||
assert_eq!(frames[0].pts_ns, 1_000_000_000);
|
||
}
|
||
}
|