Files
libfreemkv/src/mux/codec/hevc.rs
T

575 lines
19 KiB
Rust

//! HEVC (H.265) elementary stream parser.
//!
//! Extracts VPS, SPS, PPS NAL units for MKV codecPrivate.
//! Detects keyframes (IRAP pictures: IDR, CRA, BLA).
//! Each PES packet = one access unit = one frame.
use super::h264::{find_start_code, skip_start_code};
use super::{pts_to_ns, CodecParser, Frame, PesPacket};
// HEVC NAL unit types
const NAL_VPS: u8 = 32;
const NAL_SPS: u8 = 33;
const NAL_PPS: u8 = 34;
const NAL_AUD: u8 = 35;
// Dolby Vision RPU (Reference Processing Unit) — NAL type 62 (UNSPEC62).
// This is NOT filtered: all NAL types except VPS/SPS/PPS/AUD pass through
// to frame data, so DV enhancement layer RPU NALs are preserved automatically.
const _NAL_UNSPEC62_DV_RPU: u8 = 62;
// IRAP types (keyframes): BLA, IDR, CRA
const NAL_BLA_W_LP: u8 = 16;
const NAL_RSV_IRAP_VCL23: u8 = 23;
pub struct HevcParser {
vps: Option<Vec<u8>>,
sps: Option<Vec<u8>>,
pps: Option<Vec<u8>>,
}
impl Default for HevcParser {
fn default() -> Self {
Self::new()
}
}
impl HevcParser {
pub fn new() -> Self {
Self {
vps: None,
sps: None,
pps: None,
}
}
}
impl CodecParser for HevcParser {
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);
let data = &pes.data;
let mut keyframe = false;
// Scan NAL units
let mut pos = 0;
while let Some(sc_pos) = find_start_code(data, pos) {
if let Some(nal_start) = skip_start_code(data, sc_pos) {
let next = find_start_code(data, nal_start).unwrap_or(data.len());
let mut end = next;
while end > nal_start && data[end - 1] == 0x00 {
end -= 1;
}
if nal_start < data.len() {
// HEVC NAL header: 2 bytes. Type is bits 1-6 of first byte.
let nal_type = (data[nal_start] >> 1) & 0x3F;
match nal_type {
NAL_VPS => self.vps = Some(data[nal_start..end].to_vec()),
NAL_SPS => self.sps = Some(data[nal_start..end].to_vec()),
NAL_PPS => self.pps = Some(data[nal_start..end].to_vec()),
t if (NAL_BLA_W_LP..=NAL_RSV_IRAP_VCL23).contains(&t) => {
keyframe = true;
}
_ => {}
}
}
pos = next;
} else {
break;
}
}
// Convert Annex B to length-prefixed NALUs.
// Skip VPS/SPS/PPS/AUD — they're in codecPrivate.
let mut frame_data = Vec::new();
let mut pos = 0;
while let Some(sc_pos) = find_start_code(&pes.data, pos) {
if let Some(nal_start) = skip_start_code(&pes.data, sc_pos) {
let next = find_start_code(&pes.data, nal_start).unwrap_or(pes.data.len());
let mut end = next;
while end > nal_start && pes.data[end - 1] == 0x00 {
end -= 1;
}
if nal_start < pes.data.len() {
let nal_type = (pes.data[nal_start] >> 1) & 0x3F;
// Skip parameter sets and AUD
if nal_type != NAL_VPS
&& nal_type != NAL_SPS
&& nal_type != NAL_PPS
&& nal_type != NAL_AUD
{
let nal = &pes.data[nal_start..end];
let len = nal.len() as u32;
frame_data.extend_from_slice(&len.to_be_bytes());
frame_data.extend_from_slice(nal);
}
}
pos = next;
} else {
break;
}
}
if frame_data.is_empty() {
return Vec::new();
}
vec![Frame {
pts_ns,
keyframe,
data: frame_data,
}]
}
fn codec_private(&self) -> Option<Vec<u8>> {
// HEVCDecoderConfigurationRecord (ISO 14496-15)
let vps = self.vps.as_ref()?;
let sps = self.sps.as_ref()?;
let pps = self.pps.as_ref()?;
// Simplified: store as arrays in Annex B format
// Full HEVCDecoderConfigurationRecord is complex — for now, concatenate
let mut record = Vec::new();
// Minimal HEVCDecoderConfigurationRecord header
record.push(1); // configurationVersion
// General profile space, tier flag, profile IDC from SPS
if sps.len() > 3 {
record.push(sps[1]); // general_profile_space + general_tier_flag + general_profile_idc
} else {
record.push(0);
}
// general_profile_compatibility_flags (4 bytes)
record.extend_from_slice(&[0, 0, 0, 0]);
// general_constraint_indicator_flags (6 bytes)
record.extend_from_slice(&[0, 0, 0, 0, 0, 0]);
// general_level_idc
record.push(if sps.len() > 12 { sps[12] } else { 0 });
// min_spatial_segmentation_idc (4 + 12 bits)
record.extend_from_slice(&[0xF0, 0x00]);
// parallelismType (6 + 2 bits)
record.push(0xFC);
// chromaFormat (6 + 2 bits)
record.push(0xFC | 1); // 4:2:0
// bitDepthLumaMinus8 (5 + 3 bits)
record.push(0xF8);
// bitDepthChromaMinus8 (5 + 3 bits)
record.push(0xF8);
// avgFrameRate
record.extend_from_slice(&[0, 0]);
// constantFrameRate + numTemporalLayers + temporalIdNested + lengthSizeMinusOne
record.push(0x03); // lengthSizeMinusOne = 3 (4 bytes)
// numOfArrays
record.push(3); // VPS, SPS, PPS
// VPS array
record.push(0x20 | (NAL_VPS & 0x3F)); // array_completeness + NAL type
record.extend_from_slice(&[0, 1]); // numNalus = 1
record.push((vps.len() >> 8) as u8);
record.push(vps.len() as u8);
record.extend_from_slice(vps);
// SPS array
record.push(0x20 | (NAL_SPS & 0x3F));
record.extend_from_slice(&[0, 1]);
record.push((sps.len() >> 8) as u8);
record.push(sps.len() as u8);
record.extend_from_slice(sps);
// PPS array
record.push(0x20 | (NAL_PPS & 0x3F));
record.extend_from_slice(&[0, 1]);
record.push((pps.len() >> 8) as u8);
record.push(pps.len() as u8);
record.extend_from_slice(pps);
Some(record)
}
}
#[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,
}
}
/// Build an HEVC NAL header (2 bytes). Type is bits 1-6 of first byte.
/// Format: forbidden(1) | type(6) | layer_id_high(1) || layer_id_low(5) | tid(3)
fn hevc_nal_header(nal_type: u8) -> [u8; 2] {
[(nal_type & 0x3F) << 1, 0x01] // tid=1
}
// --- VPS+SPS+PPS → codec_private ---
#[test]
fn parse_vps_sps_pps() {
let mut parser = HevcParser::new();
let mut data = Vec::new();
// VPS (type 32)
data.extend_from_slice(&[0x00, 0x00, 0x01]);
let vps_hdr = hevc_nal_header(32);
data.extend_from_slice(&vps_hdr);
data.extend_from_slice(&[0xAA, 0xBB, 0xCC]); // VPS payload
// SPS (type 33)
data.extend_from_slice(&[0x00, 0x00, 0x01]);
let sps_hdr = hevc_nal_header(33);
data.extend_from_slice(&sps_hdr);
data.extend_from_slice(&[
0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D,
]); // SPS payload (>12 bytes for level)
// PPS (type 34)
data.extend_from_slice(&[0x00, 0x00, 0x01]);
let pps_hdr = hevc_nal_header(34);
data.extend_from_slice(&pps_hdr);
data.extend_from_slice(&[0xDD, 0xEE]); // PPS payload
// IRAP slice (type 19 = IDR_W_RADL) so a frame is emitted
data.extend_from_slice(&[0x00, 0x00, 0x01]);
let idr_hdr = hevc_nal_header(19);
data.extend_from_slice(&idr_hdr);
data.extend_from_slice(&[0x10, 0x20, 0x30]);
let pes = make_pes(data, Some(90000));
let _frames = parser.parse(&pes);
let cp = parser.codec_private();
assert!(
cp.is_some(),
"codec_private should be Some after VPS+SPS+PPS"
);
let cp = cp.unwrap();
// configurationVersion = 1
assert_eq!(cp[0], 1);
// numOfArrays = 3 (VPS, SPS, PPS)
assert_eq!(cp[22], 3);
// Should be longer than the minimal header (23 bytes) + array entries
assert!(
cp.len() > 23,
"codec_private should contain VPS+SPS+PPS data"
);
}
#[test]
fn codec_private_none_before_params() {
let parser = HevcParser::new();
assert!(parser.codec_private().is_none());
}
#[test]
fn codec_private_none_missing_pps() {
let mut parser = HevcParser::new();
// Only VPS + SPS, no PPS
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(32));
data.extend_from_slice(&[0xAA, 0xBB]);
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(33));
data.extend_from_slice(&[0x01, 0x02, 0x03, 0x04]);
// Add a slice so parse doesn't return empty
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(1)); // TRAIL_R
data.extend_from_slice(&[0x10, 0x20]);
let pes = make_pes(data, Some(0));
parser.parse(&pes);
assert!(
parser.codec_private().is_none(),
"should be None without PPS"
);
}
// --- IRAP keyframe detection ---
#[test]
fn parse_irap_keyframe_idr_w_radl() {
let mut parser = HevcParser::new();
let mut data = Vec::new();
// IDR_W_RADL = type 19
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(19));
data.extend_from_slice(&[0x10, 0x20, 0x30]);
let pes = make_pes(data, Some(90000));
let frames = parser.parse(&pes);
assert_eq!(frames.len(), 1);
assert!(
frames[0].keyframe,
"IDR_W_RADL (type 19) should be keyframe"
);
}
#[test]
fn parse_irap_keyframe_bla() {
let mut parser = HevcParser::new();
// BLA_W_LP = type 16
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(16));
data.extend_from_slice(&[0x10, 0x20]);
let pes = make_pes(data, Some(0));
let frames = parser.parse(&pes);
assert_eq!(frames.len(), 1);
assert!(frames[0].keyframe, "BLA_W_LP (type 16) should be keyframe");
}
#[test]
fn parse_irap_keyframe_cra() {
let mut parser = HevcParser::new();
// CRA_NUT = type 21
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(21));
data.extend_from_slice(&[0x10, 0x20]);
let pes = make_pes(data, Some(0));
let frames = parser.parse(&pes);
assert_eq!(frames.len(), 1);
assert!(frames[0].keyframe, "CRA (type 21) should be keyframe");
}
#[test]
fn parse_irap_type_23() {
let mut parser = HevcParser::new();
// RSV_IRAP_VCL23 = type 23 (upper boundary)
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(23));
data.extend_from_slice(&[0x10, 0x20]);
let pes = make_pes(data, Some(0));
let frames = parser.parse(&pes);
assert_eq!(frames.len(), 1);
assert!(frames[0].keyframe, "type 23 should be keyframe");
}
// --- non-IRAP (trailing) → not keyframe ---
#[test]
fn parse_trailing_not_keyframe() {
let mut parser = HevcParser::new();
// TRAIL_R = type 1
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(1));
data.extend_from_slice(&[0x10, 0x20, 0x30]);
let pes = make_pes(data, Some(180000));
let frames = parser.parse(&pes);
assert_eq!(frames.len(), 1);
assert!(
!frames[0].keyframe,
"TRAIL_R (type 1) should not be keyframe"
);
}
#[test]
fn parse_tsa_not_keyframe() {
let mut parser = HevcParser::new();
// TSA_N = type 2
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(2));
data.extend_from_slice(&[0x10, 0x20]);
let pes = make_pes(data, Some(0));
let frames = parser.parse(&pes);
assert_eq!(frames.len(), 1);
assert!(!frames[0].keyframe, "TSA_N (type 2) should not be keyframe");
}
// --- VPS/SPS/PPS stripped from frame data ---
#[test]
fn param_sets_stripped_from_frame() {
let mut parser = HevcParser::new();
let mut data = Vec::new();
// VPS
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(32));
data.extend_from_slice(&[0xAA]);
// SPS
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(33));
data.extend_from_slice(&[0xBB]);
// PPS
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(34));
data.extend_from_slice(&[0xCC]);
// IDR slice
data.extend_from_slice(&[0x00, 0x00, 0x01]);
let idr_hdr = hevc_nal_header(19);
data.extend_from_slice(&idr_hdr);
data.extend_from_slice(&[0x10, 0x20]);
let pes = make_pes(data, Some(0));
let frames = parser.parse(&pes);
assert_eq!(frames.len(), 1);
// Frame data should only have 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 = 2 bytes header + 2 bytes payload = 4 bytes
assert_eq!(
length as usize + 4,
fd.len(),
"frame should contain exactly one length-prefixed NAL"
);
}
// --- empty PES ---
#[test]
fn parse_empty_pes() {
let mut parser = HevcParser::new();
let pes = make_pes(Vec::new(), Some(0));
let frames = parser.parse(&pes);
assert!(frames.is_empty());
}
// --- PTS conversion ---
#[test]
fn pts_conversion() {
let mut parser = HevcParser::new();
let mut data = Vec::new();
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(1));
data.extend_from_slice(&[0x10, 0x20]);
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);
}
// --- Dolby Vision enhancement layer ---
#[test]
fn dv_rpu_nal_preserved() {
// Dolby Vision enhancement layer streams contain RPU (Reference Processing
// Unit) metadata as NAL type 62 (UNSPEC62). The HEVC parser must pass these
// through to the frame data — only VPS/SPS/PPS/AUD are stripped.
let mut parser = HevcParser::new();
let mut data = Vec::new();
// VPS (type 32) — should be stripped from frame data
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(32));
data.extend_from_slice(&[0xAA, 0xBB]);
// SPS (type 33) — should be stripped from frame data
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(33));
data.extend_from_slice(&[0x01, 0x02, 0x03, 0x04]);
// PPS (type 34) — should be stripped from frame data
data.extend_from_slice(&[0x00, 0x00, 0x01]);
data.extend_from_slice(&hevc_nal_header(34));
data.extend_from_slice(&[0xDD, 0xEE]);
// IDR_W_RADL slice (type 19) — should appear in frame data
data.extend_from_slice(&[0x00, 0x00, 0x01]);
let idr_hdr = hevc_nal_header(19);
data.extend_from_slice(&idr_hdr);
data.extend_from_slice(&[0x10, 0x20, 0x30]);
// Dolby Vision RPU (type 62 = UNSPEC62) — MUST appear in frame data
data.extend_from_slice(&[0x00, 0x00, 0x01]);
let rpu_hdr = hevc_nal_header(62);
data.extend_from_slice(&rpu_hdr);
let rpu_payload = [0xF0, 0xF1, 0xF2, 0xF3, 0xF4];
data.extend_from_slice(&rpu_payload);
let pes = make_pes(data, Some(90000));
let frames = parser.parse(&pes);
assert_eq!(frames.len(), 1, "should produce one frame");
assert!(frames[0].keyframe, "IDR should mark keyframe");
// Verify the frame data contains both the IDR NAL and the RPU NAL.
// Frame data is length-prefixed NALUs (4-byte big-endian length + NAL bytes).
let fd = &frames[0].data;
// Walk the length-prefixed NALUs and collect their types
let mut nal_types = Vec::new();
let mut offset = 0;
while offset + 4 <= fd.len() {
let length =
u32::from_be_bytes([fd[offset], fd[offset + 1], fd[offset + 2], fd[offset + 3]])
as usize;
offset += 4;
assert!(offset + length <= fd.len(), "NAL length exceeds frame data");
let nal_type = (fd[offset] >> 1) & 0x3F;
nal_types.push(nal_type);
offset += length;
}
assert!(
nal_types.contains(&19),
"frame data must contain IDR NAL (type 19), got: {:?}",
nal_types
);
assert!(
nal_types.contains(&62),
"frame data must contain Dolby Vision RPU NAL (type 62), got: {:?}",
nal_types
);
assert_eq!(
nal_types.len(),
2,
"frame data should have exactly 2 NALs (IDR + RPU), got: {:?}",
nal_types
);
// Verify RPU payload is intact
let mut offset = 0;
while offset + 4 <= fd.len() {
let length =
u32::from_be_bytes([fd[offset], fd[offset + 1], fd[offset + 2], fd[offset + 3]])
as usize;
offset += 4;
let nal_type = (fd[offset] >> 1) & 0x3F;
if nal_type == 62 {
// NAL = 2-byte header + payload
let nal_payload = &fd[offset + 2..offset + length];
assert_eq!(
nal_payload, &rpu_payload,
"RPU payload must be preserved verbatim"
);
}
offset += length;
}
}
}