Add MKV muxer and event system
- Add event.rs: structured event system for progress reporting - Add mux/: MKV muxer pipeline (EBML writer, TS demuxer, stream assembly) - Codec parsers: H.264, HEVC, AC-3, DTS, TrueHD, PGS, VC-1 - Lookahead buffer for codec private data extraction - Direct m2ts-to-MKV streaming without intermediate files
This commit is contained in:
MattJackson
@@ -0,0 +1,76 @@
|
||||
//! Event system for progress and status reporting.
|
||||
//!
|
||||
//! The lib fires events during operations like rip().
|
||||
//! The app registers a callback to receive them.
|
||||
//! No display logic, no English text — just data.
|
||||
//!
|
||||
//! ```rust
|
||||
//! disc.rip(&mut session, 0, output, |event| {
|
||||
//! match event.kind {
|
||||
//! EventKind::BytesRead { bytes, total } => update_progress(bytes, total),
|
||||
//! EventKind::ReadError { sector, .. } => log_error(sector),
|
||||
//! _ => {}
|
||||
//! }
|
||||
//! });
|
||||
//! ```
|
||||
|
||||
use crate::error::Error;
|
||||
|
||||
/// An event fired by the lib during operations.
|
||||
#[derive(Debug)]
|
||||
pub struct Event {
|
||||
pub kind: EventKind,
|
||||
}
|
||||
|
||||
/// Types of events the lib can fire.
|
||||
#[derive(Debug)]
|
||||
pub enum EventKind {
|
||||
/// Bytes successfully read and written to output.
|
||||
BytesRead {
|
||||
/// Bytes written so far.
|
||||
bytes: u64,
|
||||
/// Total bytes expected (0 if unknown).
|
||||
total: u64,
|
||||
},
|
||||
|
||||
/// A read error occurred. The lib will retry automatically.
|
||||
ReadError {
|
||||
/// Sector that failed.
|
||||
sector: u64,
|
||||
/// Error code.
|
||||
error: Error,
|
||||
},
|
||||
|
||||
/// Retrying a failed read.
|
||||
Retry {
|
||||
/// Current attempt number (1-based).
|
||||
attempt: u32,
|
||||
},
|
||||
|
||||
/// Drive speed changed (error recovery or restoration).
|
||||
SpeedChange {
|
||||
/// New speed in KB/s (0xFFFF = max).
|
||||
speed_kbs: u16,
|
||||
},
|
||||
|
||||
/// Starting a new disc extent.
|
||||
ExtentStart {
|
||||
/// Extent index (0-based).
|
||||
index: usize,
|
||||
/// First sector of extent.
|
||||
start_sector: u64,
|
||||
/// Number of sectors in extent.
|
||||
sector_count: u64,
|
||||
},
|
||||
|
||||
/// Operation complete.
|
||||
Complete {
|
||||
/// Total bytes written.
|
||||
bytes: u64,
|
||||
/// Total read errors encountered.
|
||||
errors: u32,
|
||||
},
|
||||
}
|
||||
|
||||
/// A no-op event handler. Ignores all events.
|
||||
pub fn ignore(_event: Event) {}
|
||||
@@ -0,0 +1,37 @@
|
||||
//! AC3 (Dolby Digital) / EAC3 (Dolby Digital Plus) frame parser.
|
||||
//!
|
||||
//! AC3 frames are self-contained and always start with syncword 0x0B77.
|
||||
//! Each PES packet typically contains exactly one AC3 frame.
|
||||
//! All AC3 frames are effectively keyframes (no inter-frame dependencies).
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
pub struct Ac3Parser;
|
||||
|
||||
impl Ac3Parser {
|
||||
pub fn new() -> Self {
|
||||
Self
|
||||
}
|
||||
}
|
||||
|
||||
impl CodecParser for Ac3Parser {
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||||
if pes.data.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
|
||||
// AC3: each PES = one frame, always a keyframe
|
||||
vec![Frame {
|
||||
pts_ns,
|
||||
keyframe: true,
|
||||
data: pes.data.clone(),
|
||||
}]
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> {
|
||||
// AC3 doesn't need codecPrivate in MKV
|
||||
None
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
//! DTS / DTS-HD elementary stream parser.
|
||||
//!
|
||||
//! DTS core syncword: 0x7FFE8001 (32 bits).
|
||||
//! DTS-HD MA/HRA extension follows the core frame.
|
||||
//! All frames are keyframes (no inter-frame dependencies).
|
||||
//! Each PES packet = one frame.
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
pub struct DtsParser;
|
||||
|
||||
impl DtsParser {
|
||||
pub fn new() -> Self { Self }
|
||||
}
|
||||
|
||||
impl CodecParser for DtsParser {
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||||
if pes.data.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
vec![Frame { pts_ns, keyframe: true, data: pes.data.clone() }]
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> { None }
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
//! 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: u8 = 1;
|
||||
const NAL_SLICE_IDR: u8 = 5;
|
||||
const NAL_SEI: u8 = 6;
|
||||
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 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();
|
||||
}
|
||||
|
||||
let pts_ns = 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,
|
||||
}]
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> {
|
||||
// Build AVCDecoderConfigurationRecord from SPS + PPS
|
||||
let sps = self.sps.as_ref()?;
|
||||
let pps = self.pps.as_ref()?;
|
||||
|
||||
// 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::new();
|
||||
record.push(1); // configurationVersion
|
||||
record.push(sps[1]); // profile
|
||||
record.push(sps[2]); // compatibility
|
||||
record.push(sps[3]); // level
|
||||
record.push(0xFF); // 6 bits reserved (111111) + 2 bits lengthSizeMinusOne (11 = 3)
|
||||
record.push(0xE1); // 3 bits reserved (111) + 5 bits numSPS (1)
|
||||
record.push((sps.len() >> 8) as u8);
|
||||
record.push(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`.
|
||||
pub fn find_start_code(data: &[u8], from: usize) -> Option<usize> {
|
||||
if data.len() < from + 3 {
|
||||
return None;
|
||||
}
|
||||
for i in from..data.len() - 2 {
|
||||
if data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01 {
|
||||
return Some(i);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// 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
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
//! 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::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
use super::h264::{find_start_code, skip_start_code};
|
||||
|
||||
// HEVC NAL unit types
|
||||
const NAL_VPS: u8 = 32;
|
||||
const NAL_SPS: u8 = 33;
|
||||
const NAL_PPS: u8 = 34;
|
||||
const NAL_AUD: u8 = 35;
|
||||
// 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 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();
|
||||
}
|
||||
|
||||
let pts_ns = 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 t >= NAL_BLA_W_LP && t <= NAL_RSV_IRAP_VCL23 => {
|
||||
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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
//! Elementary stream codec parsers.
|
||||
//!
|
||||
//! Each parser takes PES packets and produces frames suitable for MKV muxing.
|
||||
//! Responsibilities:
|
||||
//! - Find frame boundaries
|
||||
//! - Extract codec initialization data (SPS/PPS, etc.)
|
||||
//! - Determine keyframe status
|
||||
//! - Convert PTS from 90kHz to nanoseconds
|
||||
|
||||
pub mod ac3;
|
||||
pub mod h264;
|
||||
pub mod hevc;
|
||||
pub mod vc1;
|
||||
pub mod dts;
|
||||
pub mod truehd;
|
||||
pub mod pgs;
|
||||
|
||||
use crate::disc::Codec;
|
||||
use super::ts::PesPacket;
|
||||
|
||||
/// A single frame ready for MKV muxing.
|
||||
pub struct Frame {
|
||||
/// Presentation timestamp in nanoseconds.
|
||||
pub pts_ns: i64,
|
||||
/// Whether this is a keyframe (used for cue points).
|
||||
pub keyframe: bool,
|
||||
/// Frame data (elementary stream bytes).
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Convert 90kHz PTS to nanoseconds.
|
||||
pub fn pts_to_ns(pts: i64) -> i64 {
|
||||
pts * 100_000 / 9
|
||||
}
|
||||
|
||||
/// Trait for codec-specific elementary stream parsers.
|
||||
pub trait CodecParser: Send {
|
||||
/// Parse a PES packet into zero or more frames.
|
||||
/// Most codecs: one PES = one frame.
|
||||
/// Some (TrueHD): multiple access units per PES.
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame>;
|
||||
|
||||
/// Get codec initialization data (e.g., SPS+PPS for H.264).
|
||||
/// Returns None until enough data has been seen.
|
||||
fn codec_private(&self) -> Option<Vec<u8>>;
|
||||
}
|
||||
|
||||
/// Passthrough parser — treats each PES as one frame, no parsing.
|
||||
/// Used for codecs where PES = frame (AC3, DTS, PGS).
|
||||
pub struct PassthroughParser {
|
||||
keyframe: bool,
|
||||
}
|
||||
|
||||
impl PassthroughParser {
|
||||
pub fn new(always_keyframe: bool) -> Self {
|
||||
Self { keyframe: always_keyframe }
|
||||
}
|
||||
}
|
||||
|
||||
impl CodecParser for PassthroughParser {
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
vec![Frame {
|
||||
pts_ns,
|
||||
keyframe: self.keyframe,
|
||||
data: pes.data.clone(),
|
||||
}]
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Create the appropriate parser for a codec.
|
||||
pub fn parser_for_codec(codec: Codec) -> Box<dyn CodecParser> {
|
||||
match codec {
|
||||
Codec::H264 => Box::new(h264::H264Parser::new()),
|
||||
Codec::Hevc => Box::new(hevc::HevcParser::new()),
|
||||
Codec::Vc1 => Box::new(vc1::Vc1Parser::new()),
|
||||
Codec::Ac3 | Codec::Ac3Plus => Box::new(ac3::Ac3Parser::new()),
|
||||
Codec::DtsHdMa | Codec::DtsHdHr | Codec::Dts => Box::new(dts::DtsParser::new()),
|
||||
Codec::TrueHd => Box::new(truehd::TrueHdParser::new()),
|
||||
Codec::Pgs => Box::new(pgs::PgsParser::new()),
|
||||
Codec::Lpcm => Box::new(PassthroughParser::new(true)),
|
||||
_ => Box::new(PassthroughParser::new(true)),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
//! HDMV PGS (Presentation Graphics Stream) subtitle parser.
|
||||
//!
|
||||
//! PGS segments: PCS, WDS, PDS, ODS, END.
|
||||
//! Each PES packet contains one or more segments.
|
||||
//! All segments are keyframes (no inter-segment dependencies).
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
pub struct PgsParser;
|
||||
|
||||
impl PgsParser {
|
||||
pub fn new() -> Self { Self }
|
||||
}
|
||||
|
||||
impl CodecParser for PgsParser {
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||||
if pes.data.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
vec![Frame { pts_ns, keyframe: true, data: pes.data.clone() }]
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> { None }
|
||||
}
|
||||
@@ -0,0 +1,27 @@
|
||||
//! Dolby TrueHD / Atmos elementary stream parser.
|
||||
//!
|
||||
//! TrueHD major sync: 0xF8726FBA at a 4-byte aligned position.
|
||||
//! Access units consist of a major sync followed by minor syncs.
|
||||
//! An embedded AC3 core is in substream 0 for backward compatibility.
|
||||
//! All access units are keyframes.
|
||||
//! Each PES packet = one access unit.
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
pub struct TrueHdParser;
|
||||
|
||||
impl TrueHdParser {
|
||||
pub fn new() -> Self { Self }
|
||||
}
|
||||
|
||||
impl CodecParser for TrueHdParser {
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||||
if pes.data.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
vec![Frame { pts_ns, keyframe: true, data: pes.data.clone() }]
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> { None }
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
//! VC-1 (SMPTE 421M) elementary stream parser.
|
||||
//!
|
||||
//! VC-1 uses start codes similar to MPEG-2.
|
||||
//! Sequence header (0x0F) contains codec initialization data.
|
||||
//! Frame start = Frame header start code (0x0D).
|
||||
//! I-frames (keyframes) are identified from the frame header.
|
||||
|
||||
use super::{CodecParser, Frame, PesPacket, pts_to_ns};
|
||||
|
||||
const SC_SEQUENCE_HEADER: u8 = 0x0F;
|
||||
const SC_ENTRY_POINT: u8 = 0x0E;
|
||||
const SC_FRAME: u8 = 0x0D;
|
||||
|
||||
pub struct Vc1Parser {
|
||||
seq_header: Option<Vec<u8>>,
|
||||
entry_point: Option<Vec<u8>>,
|
||||
}
|
||||
|
||||
impl Vc1Parser {
|
||||
pub fn new() -> Self {
|
||||
Self { seq_header: None, entry_point: None }
|
||||
}
|
||||
}
|
||||
|
||||
impl CodecParser for Vc1Parser {
|
||||
fn parse(&mut self, pes: &PesPacket) -> Vec<Frame> {
|
||||
if pes.data.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
|
||||
let pts_ns = pes.pts.map(pts_to_ns).unwrap_or(0);
|
||||
let mut keyframe = false;
|
||||
|
||||
// Scan for start codes (00 00 01 XX)
|
||||
let data = &pes.data;
|
||||
let mut i = 0;
|
||||
while i + 3 < data.len() {
|
||||
if data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01 {
|
||||
let sc_type = data[i + 3];
|
||||
match sc_type {
|
||||
SC_SEQUENCE_HEADER => {
|
||||
// Capture everything from here to next start code
|
||||
let end = find_next_sc(data, i + 4).unwrap_or(data.len());
|
||||
self.seq_header = Some(data[i..end].to_vec());
|
||||
}
|
||||
SC_ENTRY_POINT => {
|
||||
let end = find_next_sc(data, i + 4).unwrap_or(data.len());
|
||||
self.entry_point = Some(data[i..end].to_vec());
|
||||
}
|
||||
SC_FRAME => {
|
||||
// First frame after sequence header + entry point is a keyframe
|
||||
if self.seq_header.is_some() && self.entry_point.is_some() {
|
||||
keyframe = true;
|
||||
}
|
||||
// Also check frame type from bitstream (bit after start code)
|
||||
if i + 4 < data.len() {
|
||||
// For Advanced profile: first 2 bits of frame data indicate type
|
||||
// But simpler: any frame preceded by seq+entry is I-frame
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
i += 4;
|
||||
} else {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
|
||||
vec![Frame {
|
||||
pts_ns,
|
||||
keyframe,
|
||||
data: pes.data.clone(),
|
||||
}]
|
||||
}
|
||||
|
||||
fn codec_private(&self) -> Option<Vec<u8>> {
|
||||
// MKV V_MS/VFW/FOURCC requires BITMAPINFOHEADER (40 bytes) + extra codec data.
|
||||
// The sequence header + entry point go as extra data after the header.
|
||||
let sh = self.seq_header.as_ref()?;
|
||||
let ep = self.entry_point.as_ref()?;
|
||||
|
||||
let extra_len = sh.len() + ep.len();
|
||||
let header_size: u32 = 40 + extra_len as u32;
|
||||
|
||||
let mut cp = Vec::with_capacity(header_size as usize);
|
||||
|
||||
// BITMAPINFOHEADER (40 bytes, little-endian)
|
||||
cp.extend_from_slice(&header_size.to_le_bytes()); // biSize
|
||||
cp.extend_from_slice(&1920u32.to_le_bytes()); // biWidth (updated by player)
|
||||
cp.extend_from_slice(&1080u32.to_le_bytes()); // biHeight
|
||||
cp.extend_from_slice(&1u16.to_le_bytes()); // biPlanes
|
||||
cp.extend_from_slice(&24u16.to_le_bytes()); // biBitCount
|
||||
cp.extend_from_slice(b"WVC1"); // biCompression = "WVC1" FOURCC
|
||||
cp.extend_from_slice(&0u32.to_le_bytes()); // biSizeImage
|
||||
cp.extend_from_slice(&0u32.to_le_bytes()); // biXPelsPerMeter
|
||||
cp.extend_from_slice(&0u32.to_le_bytes()); // biYPelsPerMeter
|
||||
cp.extend_from_slice(&0u32.to_le_bytes()); // biClrUsed
|
||||
cp.extend_from_slice(&0u32.to_le_bytes()); // biClrImportant
|
||||
|
||||
// Extra codec data: sequence header + entry point (Annex B)
|
||||
cp.extend_from_slice(sh);
|
||||
cp.extend_from_slice(ep);
|
||||
|
||||
Some(cp)
|
||||
}
|
||||
}
|
||||
|
||||
fn find_next_sc(data: &[u8], from: usize) -> Option<usize> {
|
||||
for i in from..data.len().saturating_sub(2) {
|
||||
if data[i] == 0x00 && data[i + 1] == 0x00 && data[i + 2] == 0x01 {
|
||||
return Some(i);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
+272
@@ -0,0 +1,272 @@
|
||||
//! EBML (Extensible Binary Meta Language) write primitives for Matroska.
|
||||
//!
|
||||
//! EBML uses variable-length integers for element IDs and sizes.
|
||||
//! This module provides low-level writers for constructing MKV files.
|
||||
|
||||
use std::io::{self, Write, Seek, SeekFrom};
|
||||
|
||||
/// Write an EBML element ID (1-4 bytes, already encoded).
|
||||
/// Element IDs are predefined constants — we write them verbatim.
|
||||
pub fn write_id(w: &mut impl Write, id: u32) -> io::Result<()> {
|
||||
if id <= 0x7F {
|
||||
w.write_all(&[id as u8])
|
||||
} else if id <= 0x7FFF {
|
||||
w.write_all(&[(id >> 8) as u8, id as u8])
|
||||
} else if id <= 0x7F_FFFF {
|
||||
w.write_all(&[(id >> 16) as u8, (id >> 8) as u8, id as u8])
|
||||
} else {
|
||||
w.write_all(&[(id >> 24) as u8, (id >> 16) as u8, (id >> 8) as u8, id as u8])
|
||||
}
|
||||
}
|
||||
|
||||
/// Write an EBML variable-length size (1-8 bytes).
|
||||
/// Uses the EBML VINT encoding: leading bits indicate width.
|
||||
pub fn write_size(w: &mut impl Write, size: u64) -> io::Result<()> {
|
||||
if size < 0x7F {
|
||||
w.write_all(&[(size as u8) | 0x80])
|
||||
} else if size < 0x3FFF {
|
||||
w.write_all(&[((size >> 8) as u8) | 0x40, size as u8])
|
||||
} else if size < 0x1F_FFFF {
|
||||
w.write_all(&[
|
||||
((size >> 16) as u8) | 0x20,
|
||||
(size >> 8) as u8,
|
||||
size as u8,
|
||||
])
|
||||
} else if size < 0x0FFF_FFFF {
|
||||
w.write_all(&[
|
||||
((size >> 24) as u8) | 0x10,
|
||||
(size >> 16) as u8,
|
||||
(size >> 8) as u8,
|
||||
size as u8,
|
||||
])
|
||||
} else {
|
||||
// 8-byte size for large elements
|
||||
w.write_all(&[
|
||||
0x01,
|
||||
(size >> 48) as u8,
|
||||
(size >> 40) as u8,
|
||||
(size >> 32) as u8,
|
||||
(size >> 24) as u8,
|
||||
(size >> 16) as u8,
|
||||
(size >> 8) as u8,
|
||||
size as u8,
|
||||
])
|
||||
}
|
||||
}
|
||||
|
||||
/// Write an EBML "unknown size" marker (all 1s in VINT, 8 bytes).
|
||||
/// Used for the Segment element when total size isn't known upfront.
|
||||
pub fn write_unknown_size(w: &mut impl Write) -> io::Result<()> {
|
||||
w.write_all(&[0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF])
|
||||
}
|
||||
|
||||
/// Write a complete EBML unsigned integer element.
|
||||
pub fn write_uint(w: &mut impl Write, id: u32, val: u64) -> io::Result<()> {
|
||||
write_id(w, id)?;
|
||||
if val <= 0xFF {
|
||||
write_size(w, 1)?;
|
||||
w.write_all(&[val as u8])
|
||||
} else if val <= 0xFFFF {
|
||||
write_size(w, 2)?;
|
||||
w.write_all(&[(val >> 8) as u8, val as u8])
|
||||
} else if val <= 0xFF_FFFF {
|
||||
write_size(w, 3)?;
|
||||
w.write_all(&[(val >> 16) as u8, (val >> 8) as u8, val as u8])
|
||||
} else if val <= 0xFFFF_FFFF {
|
||||
write_size(w, 4)?;
|
||||
w.write_all(&[
|
||||
(val >> 24) as u8, (val >> 16) as u8,
|
||||
(val >> 8) as u8, val as u8,
|
||||
])
|
||||
} else {
|
||||
write_size(w, 8)?;
|
||||
w.write_all(&val.to_be_bytes())
|
||||
}
|
||||
}
|
||||
|
||||
/// Write a complete EBML signed integer element.
|
||||
pub fn write_int(w: &mut impl Write, id: u32, val: i64) -> io::Result<()> {
|
||||
write_uint(w, id, val as u64)
|
||||
}
|
||||
|
||||
/// Write a complete EBML float element (8-byte double).
|
||||
pub fn write_float(w: &mut impl Write, id: u32, val: f64) -> io::Result<()> {
|
||||
write_id(w, id)?;
|
||||
write_size(w, 8)?;
|
||||
w.write_all(&val.to_be_bytes())
|
||||
}
|
||||
|
||||
/// Write a complete EBML UTF-8 string element.
|
||||
pub fn write_string(w: &mut impl Write, id: u32, val: &str) -> io::Result<()> {
|
||||
write_id(w, id)?;
|
||||
write_size(w, val.len() as u64)?;
|
||||
w.write_all(val.as_bytes())
|
||||
}
|
||||
|
||||
/// Write a complete EBML binary element.
|
||||
pub fn write_binary(w: &mut impl Write, id: u32, data: &[u8]) -> io::Result<()> {
|
||||
write_id(w, id)?;
|
||||
write_size(w, data.len() as u64)?;
|
||||
w.write_all(data)
|
||||
}
|
||||
|
||||
/// Start a master element: write ID + placeholder size.
|
||||
/// Returns the file offset of the size field for later fixup.
|
||||
pub fn start_master<W: Write + Seek>(w: &mut W, id: u32) -> io::Result<u64> {
|
||||
write_id(w, id)?;
|
||||
let size_pos = w.stream_position()?;
|
||||
// 8-byte size placeholder (will be overwritten by end_master)
|
||||
w.write_all(&[0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00])?;
|
||||
Ok(size_pos)
|
||||
}
|
||||
|
||||
/// End a master element: seek back and write the actual size.
|
||||
pub fn end_master<W: Write + Seek>(w: &mut W, size_pos: u64) -> io::Result<()> {
|
||||
let end_pos = w.stream_position()?;
|
||||
let data_size = end_pos - size_pos - 8; // subtract the 8-byte size field itself
|
||||
w.seek(SeekFrom::Start(size_pos))?;
|
||||
// Write as 8-byte VINT: 0x01 followed by 7 bytes of size
|
||||
w.write_all(&[
|
||||
0x01,
|
||||
(data_size >> 48) as u8,
|
||||
(data_size >> 40) as u8,
|
||||
(data_size >> 32) as u8,
|
||||
(data_size >> 24) as u8,
|
||||
(data_size >> 16) as u8,
|
||||
(data_size >> 8) as u8,
|
||||
data_size as u8,
|
||||
])?;
|
||||
w.seek(SeekFrom::Start(end_pos))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Matroska Element IDs
|
||||
// ============================================================
|
||||
|
||||
// EBML Header
|
||||
pub const EBML: u32 = 0x1A45DFA3;
|
||||
pub const EBML_VERSION: u32 = 0x4286;
|
||||
pub const EBML_READ_VERSION: u32 = 0x42F7;
|
||||
pub const EBML_MAX_ID_LENGTH: u32 = 0x42F2;
|
||||
pub const EBML_MAX_SIZE_LENGTH: u32 = 0x42F3;
|
||||
pub const EBML_DOC_TYPE: u32 = 0x4282;
|
||||
pub const EBML_DOC_TYPE_VERSION: u32 = 0x4287;
|
||||
pub const EBML_DOC_TYPE_READ_VERSION: u32 = 0x4285;
|
||||
|
||||
// Segment
|
||||
pub const SEGMENT: u32 = 0x18538067;
|
||||
|
||||
// Seek Head
|
||||
pub const SEEK_HEAD: u32 = 0x114D9B74;
|
||||
pub const SEEK: u32 = 0x4DBB;
|
||||
pub const SEEK_ID: u32 = 0x53AB;
|
||||
pub const SEEK_POSITION: u32 = 0x53AC;
|
||||
|
||||
// Segment Info
|
||||
pub const INFO: u32 = 0x1549A966;
|
||||
pub const TIMESTAMP_SCALE: u32 = 0x2AD7B1;
|
||||
pub const DURATION: u32 = 0x4489;
|
||||
pub const MUXING_APP: u32 = 0x4D80;
|
||||
pub const WRITING_APP: u32 = 0x5741;
|
||||
pub const TITLE: u32 = 0x7BA9;
|
||||
|
||||
// Tracks
|
||||
pub const TRACKS: u32 = 0x1654AE6B;
|
||||
pub const TRACK_ENTRY: u32 = 0xAE;
|
||||
pub const TRACK_NUMBER: u32 = 0xD7;
|
||||
pub const TRACK_UID: u32 = 0x73C5;
|
||||
pub const TRACK_TYPE: u32 = 0x83;
|
||||
pub const FLAG_LACING: u32 = 0x9C;
|
||||
pub const FLAG_DEFAULT: u32 = 0x88;
|
||||
pub const FLAG_FORCED: u32 = 0x55AA;
|
||||
pub const LANGUAGE: u32 = 0x22B59C;
|
||||
pub const CODEC_ID: u32 = 0x86;
|
||||
pub const CODEC_PRIVATE: u32 = 0x63A2;
|
||||
pub const DEFAULT_DURATION: u32 = 0x23E383;
|
||||
|
||||
// Video
|
||||
pub const VIDEO: u32 = 0xE0;
|
||||
pub const PIXEL_WIDTH: u32 = 0xB0;
|
||||
pub const PIXEL_HEIGHT: u32 = 0xBA;
|
||||
pub const DISPLAY_WIDTH: u32 = 0x54B0;
|
||||
pub const DISPLAY_HEIGHT: u32 = 0x54BA;
|
||||
pub const COLOUR: u32 = 0x55B0;
|
||||
pub const TRANSFER_CHARACTERISTICS: u32 = 0x55BA;
|
||||
pub const MATRIX_COEFFICIENTS: u32 = 0x55B1;
|
||||
pub const PRIMARIES: u32 = 0x55BB;
|
||||
pub const RANGE: u32 = 0x55B9;
|
||||
|
||||
// Audio
|
||||
pub const AUDIO: u32 = 0xE1;
|
||||
pub const SAMPLING_FREQUENCY: u32 = 0xB5;
|
||||
pub const CHANNELS: u32 = 0x9F;
|
||||
pub const BIT_DEPTH: u32 = 0x6264;
|
||||
|
||||
// Cluster
|
||||
pub const CLUSTER: u32 = 0x1F43B675;
|
||||
pub const CLUSTER_TIMESTAMP: u32 = 0xE7;
|
||||
pub const SIMPLE_BLOCK: u32 = 0xA3;
|
||||
|
||||
// Cues
|
||||
pub const CUES: u32 = 0x1C53BB6B;
|
||||
pub const CUE_POINT: u32 = 0xBB;
|
||||
pub const CUE_TIME: u32 = 0xB3;
|
||||
pub const CUE_TRACK_POSITIONS: u32 = 0xB7;
|
||||
pub const CUE_TRACK: u32 = 0xF7;
|
||||
pub const CUE_CLUSTER_POSITION: u32 = 0xF1;
|
||||
|
||||
// Track types
|
||||
pub const TRACK_TYPE_VIDEO: u64 = 1;
|
||||
pub const TRACK_TYPE_AUDIO: u64 = 2;
|
||||
pub const TRACK_TYPE_SUBTITLE: u64 = 17;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::io::Cursor;
|
||||
|
||||
#[test]
|
||||
fn test_write_size() {
|
||||
let mut buf = Vec::new();
|
||||
write_size(&mut buf, 0).unwrap();
|
||||
assert_eq!(buf, [0x80]);
|
||||
|
||||
buf.clear();
|
||||
write_size(&mut buf, 127).unwrap();
|
||||
assert_eq!(buf, [0xC0, 127]); // 127 >= 0x7F, uses 2 bytes
|
||||
|
||||
buf.clear();
|
||||
write_size(&mut buf, 126).unwrap();
|
||||
assert_eq!(buf, [126 | 0x80]); // 126 < 0x7F, uses 1 byte
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_uint() {
|
||||
let mut buf = Vec::new();
|
||||
write_uint(&mut buf, 0x4286, 1).unwrap(); // EBML_VERSION = 1
|
||||
// ID: 42 86, Size: 81 (1 byte), Data: 01
|
||||
assert_eq!(buf, [0x42, 0x86, 0x81, 0x01]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_write_string() {
|
||||
let mut buf = Vec::new();
|
||||
write_string(&mut buf, 0x4282, "matroska").unwrap();
|
||||
// ID: 42 82, Size: 88 (8 bytes), Data: "matroska"
|
||||
assert_eq!(&buf[0..2], &[0x42, 0x82]);
|
||||
assert_eq!(buf[2], 0x88); // size = 8
|
||||
assert_eq!(&buf[3..], b"matroska");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_master_element() {
|
||||
let mut buf = Cursor::new(Vec::new());
|
||||
let pos = start_master(&mut buf, EBML).unwrap();
|
||||
write_uint(&mut buf, EBML_VERSION, 1).unwrap();
|
||||
end_master(&mut buf, pos).unwrap();
|
||||
let data = buf.into_inner();
|
||||
// EBML header: 1A 45 DF A3, then 8-byte size, then content
|
||||
assert_eq!(&data[0..4], &[0x1A, 0x45, 0xDF, 0xA3]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,127 @@
|
||||
//! LookaheadBuffer — generic pre-scan buffer for stream pipelines.
|
||||
//!
|
||||
//! Accumulates data up to a configurable limit. When the consumer finds
|
||||
//! what it needs, the buffer can be drained (fast path, no re-read).
|
||||
//! If the buffer fills before the consumer is satisfied, it signals
|
||||
//! overflow — the caller should discard and re-read from the source.
|
||||
//!
|
||||
//! Used by MkvStream to collect SPS/PPS before writing the MKV header.
|
||||
//! Reusable for any stream stage that needs to look ahead.
|
||||
|
||||
/// Default lookahead buffer size: 5 MB.
|
||||
pub const DEFAULT_LOOKAHEAD_SIZE: usize = 5 * 1024 * 1024;
|
||||
|
||||
/// Lookahead buffer states.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
pub enum LookaheadState {
|
||||
/// Still collecting data, haven't found what we need yet.
|
||||
Collecting,
|
||||
/// Found what we need, buffer has the data ready to drain.
|
||||
Ready,
|
||||
/// Buffer overflowed before finding what we need.
|
||||
/// Caller should discard buffer, finish scanning without buffering,
|
||||
/// then re-read from the source.
|
||||
Overflow,
|
||||
}
|
||||
|
||||
/// A bounded lookahead buffer.
|
||||
pub struct LookaheadBuffer {
|
||||
data: Vec<u8>,
|
||||
max_size: usize,
|
||||
state: LookaheadState,
|
||||
}
|
||||
|
||||
impl LookaheadBuffer {
|
||||
/// Create a new buffer with the given max size.
|
||||
/// Pass 0 for no buffering (always overflows immediately).
|
||||
pub fn new(max_size: usize) -> Self {
|
||||
Self {
|
||||
data: Vec::with_capacity(max_size.min(DEFAULT_LOOKAHEAD_SIZE)),
|
||||
max_size,
|
||||
state: LookaheadState::Collecting,
|
||||
}
|
||||
}
|
||||
|
||||
/// Push data into the buffer. Returns the new state.
|
||||
/// If the buffer would overflow, transitions to Overflow state.
|
||||
pub fn push(&mut self, chunk: &[u8]) -> LookaheadState {
|
||||
if self.state != LookaheadState::Collecting {
|
||||
return self.state;
|
||||
}
|
||||
|
||||
if self.data.len() + chunk.len() > self.max_size {
|
||||
self.state = LookaheadState::Overflow;
|
||||
return self.state;
|
||||
}
|
||||
|
||||
self.data.extend_from_slice(chunk);
|
||||
self.state
|
||||
}
|
||||
|
||||
/// Mark the buffer as ready — we found what we need.
|
||||
pub fn mark_ready(&mut self) {
|
||||
if self.state == LookaheadState::Collecting {
|
||||
self.state = LookaheadState::Ready;
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the buffered data (only valid in Ready state).
|
||||
pub fn data(&self) -> &[u8] {
|
||||
&self.data
|
||||
}
|
||||
|
||||
/// Take ownership of the buffered data, clearing the buffer.
|
||||
pub fn drain(&mut self) -> Vec<u8> {
|
||||
self.state = LookaheadState::Collecting;
|
||||
std::mem::take(&mut self.data)
|
||||
}
|
||||
|
||||
/// Current state.
|
||||
pub fn state(&self) -> LookaheadState {
|
||||
self.state
|
||||
}
|
||||
|
||||
/// How many bytes are buffered.
|
||||
pub fn len(&self) -> usize {
|
||||
self.data.len()
|
||||
}
|
||||
|
||||
/// Is the buffer empty?
|
||||
pub fn is_empty(&self) -> bool {
|
||||
self.data.is_empty()
|
||||
}
|
||||
|
||||
/// Max size this buffer can hold.
|
||||
pub fn max_size(&self) -> usize {
|
||||
self.max_size
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_basic_flow() {
|
||||
let mut buf = LookaheadBuffer::new(100);
|
||||
assert_eq!(buf.push(b"hello"), LookaheadState::Collecting);
|
||||
assert_eq!(buf.push(b"world"), LookaheadState::Collecting);
|
||||
assert_eq!(buf.len(), 10);
|
||||
buf.mark_ready();
|
||||
assert_eq!(buf.state(), LookaheadState::Ready);
|
||||
assert_eq!(buf.data(), b"helloworld");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_overflow() {
|
||||
let mut buf = LookaheadBuffer::new(5);
|
||||
assert_eq!(buf.push(b"abc"), LookaheadState::Collecting);
|
||||
assert_eq!(buf.push(b"def"), LookaheadState::Overflow);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_zero_size() {
|
||||
let mut buf = LookaheadBuffer::new(0);
|
||||
assert_eq!(buf.push(b"a"), LookaheadState::Overflow);
|
||||
}
|
||||
}
|
||||
+375
@@ -0,0 +1,375 @@
|
||||
//! Matroska (MKV) muxer.
|
||||
//!
|
||||
//! Writes EBML header, Segment with tracks, clusters, and cues.
|
||||
//! Designed for streaming writes: clusters are written as data arrives,
|
||||
//! cues and seek head are finalized at the end.
|
||||
|
||||
use std::io::{self, Write, Seek, SeekFrom};
|
||||
use super::ebml;
|
||||
use crate::disc::{VideoStream, AudioStream, SubtitleStream, Codec};
|
||||
|
||||
/// MKV track definition (built from disc stream metadata).
|
||||
pub struct MkvTrack {
|
||||
pub track_type: u64, // 1=video, 2=audio, 17=subtitle
|
||||
pub codec_id: &'static str,
|
||||
pub language: String,
|
||||
pub codec_private: Option<Vec<u8>>,
|
||||
pub is_default: bool,
|
||||
pub is_forced: bool,
|
||||
// Video-specific
|
||||
pub pixel_width: u32,
|
||||
pub pixel_height: u32,
|
||||
// Audio-specific
|
||||
pub sample_rate: f64,
|
||||
pub channels: u8,
|
||||
pub bit_depth: u8,
|
||||
}
|
||||
|
||||
impl MkvTrack {
|
||||
pub fn video(v: &VideoStream) -> Self {
|
||||
let codec_id = match v.codec {
|
||||
Codec::H264 => "V_MPEG4/ISO/AVC",
|
||||
Codec::Hevc => "V_MPEGH/ISO/HEVC",
|
||||
Codec::Vc1 => "V_MS/VFW/FOURCC",
|
||||
Codec::Mpeg2 => "V_MPEG2",
|
||||
_ => "V_MPEG2",
|
||||
};
|
||||
let (w, h) = parse_resolution(&v.resolution);
|
||||
Self {
|
||||
track_type: ebml::TRACK_TYPE_VIDEO,
|
||||
codec_id,
|
||||
language: "und".into(),
|
||||
codec_private: None, // filled later by parser
|
||||
is_default: !v.secondary,
|
||||
is_forced: false,
|
||||
pixel_width: w,
|
||||
pixel_height: h,
|
||||
sample_rate: 0.0,
|
||||
channels: 0,
|
||||
bit_depth: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn audio(a: &AudioStream) -> Self {
|
||||
let codec_id = match a.codec {
|
||||
Codec::Ac3 => "A_AC3",
|
||||
Codec::Ac3Plus => "A_EAC3",
|
||||
Codec::TrueHd => "A_TRUEHD",
|
||||
Codec::DtsHdMa | Codec::DtsHdHr | Codec::Dts => "A_DTS",
|
||||
Codec::Lpcm => "A_PCM/INT/BIG",
|
||||
_ => "A_AC3",
|
||||
};
|
||||
let sr = parse_sample_rate(&a.sample_rate);
|
||||
let ch = parse_channels(&a.channels);
|
||||
Self {
|
||||
track_type: ebml::TRACK_TYPE_AUDIO,
|
||||
codec_id,
|
||||
language: a.language.clone(),
|
||||
codec_private: None,
|
||||
is_default: !a.secondary,
|
||||
is_forced: false,
|
||||
pixel_width: 0,
|
||||
pixel_height: 0,
|
||||
sample_rate: sr,
|
||||
channels: ch,
|
||||
bit_depth: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn subtitle(s: &SubtitleStream) -> Self {
|
||||
Self {
|
||||
track_type: ebml::TRACK_TYPE_SUBTITLE,
|
||||
codec_id: "S_HDMV/PGS",
|
||||
language: s.language.clone(),
|
||||
codec_private: None,
|
||||
is_default: false,
|
||||
is_forced: s.forced,
|
||||
pixel_width: 0,
|
||||
pixel_height: 0,
|
||||
sample_rate: 0.0,
|
||||
channels: 0,
|
||||
bit_depth: 0,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Cue point for seeking.
|
||||
struct CuePoint {
|
||||
timestamp_ms: i64,
|
||||
track: usize,
|
||||
cluster_pos: u64, // relative to Segment start
|
||||
}
|
||||
|
||||
/// MKV muxer. Call write_frame() for each frame, then finish() at the end.
|
||||
pub struct MkvMuxer<W: Write + Seek> {
|
||||
writer: W,
|
||||
segment_start: u64,
|
||||
cluster_open: bool,
|
||||
cluster_pos: u64,
|
||||
cluster_size_pos: u64,
|
||||
cluster_ts_ms: i64,
|
||||
cues: Vec<CuePoint>,
|
||||
frame_count: u64,
|
||||
/// File positions of codecPrivate placeholders (track_idx → offset, max_size).
|
||||
/// Used to seek back and fill in SPS/PPS after first keyframe.
|
||||
codec_private_slots: Vec<Option<(u64, usize)>>,
|
||||
codec_private_filled: Vec<bool>,
|
||||
}
|
||||
|
||||
/// New cluster every 5 seconds.
|
||||
const CLUSTER_DURATION_MS: i64 = 5000;
|
||||
|
||||
impl<W: Write + Seek> MkvMuxer<W> {
|
||||
/// Create a new MKV muxer: writes EBML header, Segment start, Info, Tracks.
|
||||
pub fn new(mut writer: W, tracks: &[MkvTrack], title: Option<&str>, duration_secs: f64) -> io::Result<Self> {
|
||||
// EBML Header
|
||||
let ebml_pos = ebml::start_master(&mut writer, ebml::EBML)?;
|
||||
ebml::write_uint(&mut writer, ebml::EBML_VERSION, 1)?;
|
||||
ebml::write_uint(&mut writer, ebml::EBML_READ_VERSION, 1)?;
|
||||
ebml::write_uint(&mut writer, ebml::EBML_MAX_ID_LENGTH, 4)?;
|
||||
ebml::write_uint(&mut writer, ebml::EBML_MAX_SIZE_LENGTH, 8)?;
|
||||
ebml::write_string(&mut writer, ebml::EBML_DOC_TYPE, "matroska")?;
|
||||
ebml::write_uint(&mut writer, ebml::EBML_DOC_TYPE_VERSION, 4)?;
|
||||
ebml::write_uint(&mut writer, ebml::EBML_DOC_TYPE_READ_VERSION, 2)?;
|
||||
ebml::end_master(&mut writer, ebml_pos)?;
|
||||
|
||||
// Segment (unknown size — we'll write cues at the end)
|
||||
ebml::write_id(&mut writer, ebml::SEGMENT)?;
|
||||
ebml::write_unknown_size(&mut writer)?;
|
||||
let segment_start = writer.stream_position()?;
|
||||
|
||||
// Info
|
||||
let info_pos = ebml::start_master(&mut writer, ebml::INFO)?;
|
||||
ebml::write_uint(&mut writer, ebml::TIMESTAMP_SCALE, 1_000_000)?; // 1ms precision
|
||||
if duration_secs > 0.0 {
|
||||
ebml::write_float(&mut writer, ebml::DURATION, duration_secs * 1000.0)?; // in ms
|
||||
}
|
||||
ebml::write_string(&mut writer, ebml::MUXING_APP, "freemkv")?;
|
||||
ebml::write_string(&mut writer, ebml::WRITING_APP, "freemkv")?;
|
||||
if let Some(t) = title {
|
||||
ebml::write_string(&mut writer, ebml::TITLE, t)?;
|
||||
}
|
||||
ebml::end_master(&mut writer, info_pos)?;
|
||||
|
||||
// Tracks
|
||||
let mut codec_private_slots: Vec<Option<(u64, usize)>> = Vec::new();
|
||||
let mut codec_private_filled: Vec<bool> = Vec::new();
|
||||
let tracks_pos = ebml::start_master(&mut writer, ebml::TRACKS)?;
|
||||
for (i, track) in tracks.iter().enumerate() {
|
||||
let entry_pos = ebml::start_master(&mut writer, ebml::TRACK_ENTRY)?;
|
||||
ebml::write_uint(&mut writer, ebml::TRACK_NUMBER, (i + 1) as u64)?;
|
||||
ebml::write_uint(&mut writer, ebml::TRACK_UID, (i + 1) as u64)?;
|
||||
ebml::write_uint(&mut writer, ebml::TRACK_TYPE, track.track_type)?;
|
||||
ebml::write_uint(&mut writer, ebml::FLAG_LACING, 0)?;
|
||||
ebml::write_string(&mut writer, ebml::CODEC_ID, track.codec_id)?;
|
||||
ebml::write_string(&mut writer, ebml::LANGUAGE, &track.language)?;
|
||||
|
||||
if !track.is_default {
|
||||
ebml::write_uint(&mut writer, ebml::FLAG_DEFAULT, 0)?;
|
||||
}
|
||||
if track.is_forced {
|
||||
ebml::write_uint(&mut writer, ebml::FLAG_FORCED, 1)?;
|
||||
}
|
||||
|
||||
if let Some(ref cp) = track.codec_private {
|
||||
ebml::write_binary(&mut writer, ebml::CODEC_PRIVATE, cp)?;
|
||||
codec_private_slots.push(None); // already filled
|
||||
codec_private_filled.push(true);
|
||||
} else if track.track_type == ebml::TRACK_TYPE_VIDEO {
|
||||
// Reserve space for codecPrivate — will be filled after first keyframe
|
||||
// Reserve 256 bytes (enough for SPS+PPS or VPS+SPS+PPS)
|
||||
let cp_pos = writer.stream_position()?;
|
||||
let placeholder = vec![0u8; 256];
|
||||
ebml::write_binary(&mut writer, ebml::CODEC_PRIVATE, &placeholder)?;
|
||||
codec_private_slots.push(Some((cp_pos, 256)));
|
||||
codec_private_filled.push(false);
|
||||
} else {
|
||||
codec_private_slots.push(None);
|
||||
codec_private_filled.push(true);
|
||||
}
|
||||
|
||||
// Video-specific
|
||||
if track.track_type == ebml::TRACK_TYPE_VIDEO && track.pixel_width > 0 {
|
||||
let vid_pos = ebml::start_master(&mut writer, ebml::VIDEO)?;
|
||||
ebml::write_uint(&mut writer, ebml::PIXEL_WIDTH, track.pixel_width as u64)?;
|
||||
ebml::write_uint(&mut writer, ebml::PIXEL_HEIGHT, track.pixel_height as u64)?;
|
||||
ebml::end_master(&mut writer, vid_pos)?;
|
||||
}
|
||||
|
||||
// Audio-specific
|
||||
if track.track_type == ebml::TRACK_TYPE_AUDIO && track.sample_rate > 0.0 {
|
||||
let aud_pos = ebml::start_master(&mut writer, ebml::AUDIO)?;
|
||||
ebml::write_float(&mut writer, ebml::SAMPLING_FREQUENCY, track.sample_rate)?;
|
||||
ebml::write_uint(&mut writer, ebml::CHANNELS, track.channels as u64)?;
|
||||
if track.bit_depth > 0 {
|
||||
ebml::write_uint(&mut writer, ebml::BIT_DEPTH, track.bit_depth as u64)?;
|
||||
}
|
||||
ebml::end_master(&mut writer, aud_pos)?;
|
||||
}
|
||||
|
||||
ebml::end_master(&mut writer, entry_pos)?;
|
||||
}
|
||||
ebml::end_master(&mut writer, tracks_pos)?;
|
||||
|
||||
Ok(Self {
|
||||
writer,
|
||||
segment_start,
|
||||
cluster_open: false,
|
||||
cluster_pos: 0,
|
||||
cluster_size_pos: 0,
|
||||
cluster_ts_ms: 0,
|
||||
cues: Vec::new(),
|
||||
frame_count: 0,
|
||||
codec_private_slots,
|
||||
codec_private_filled,
|
||||
})
|
||||
}
|
||||
|
||||
/// Write a single frame.
|
||||
pub fn write_frame(&mut self, track_idx: usize, pts_ns: i64, keyframe: bool, data: &[u8]) -> io::Result<()> {
|
||||
let pts_ms = pts_ns / 1_000_000;
|
||||
|
||||
// Start new cluster if needed
|
||||
if !self.cluster_open || (pts_ms - self.cluster_ts_ms) >= CLUSTER_DURATION_MS {
|
||||
if self.cluster_open {
|
||||
// Close current cluster (it's a master with unknown size — we use known size)
|
||||
// Actually, for streaming we keep clusters open-ended. Just start a new one.
|
||||
}
|
||||
self.start_cluster(pts_ms)?;
|
||||
|
||||
// Add cue point at cluster start for keyframes (video track 0)
|
||||
if keyframe && track_idx == 0 {
|
||||
self.cues.push(CuePoint {
|
||||
timestamp_ms: pts_ms,
|
||||
track: track_idx + 1,
|
||||
cluster_pos: self.cluster_pos - self.segment_start,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Write SimpleBlock
|
||||
let relative_ts = (pts_ms - self.cluster_ts_ms) as i16;
|
||||
self.write_simple_block(track_idx + 1, relative_ts, keyframe, data)?;
|
||||
self.frame_count += 1;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Finish the MKV file: write Cues element.
|
||||
pub fn finish(mut self) -> io::Result<()> {
|
||||
// Close final cluster
|
||||
self.end_cluster()?;
|
||||
|
||||
// Write Cues
|
||||
if !self.cues.is_empty() {
|
||||
let cues_pos = ebml::start_master(&mut self.writer, ebml::CUES)?;
|
||||
for cue in &self.cues {
|
||||
let cp_pos = ebml::start_master(&mut self.writer, ebml::CUE_POINT)?;
|
||||
ebml::write_uint(&mut self.writer, ebml::CUE_TIME, cue.timestamp_ms as u64)?;
|
||||
let ctp_pos = ebml::start_master(&mut self.writer, ebml::CUE_TRACK_POSITIONS)?;
|
||||
ebml::write_uint(&mut self.writer, ebml::CUE_TRACK, cue.track as u64)?;
|
||||
ebml::write_uint(&mut self.writer, ebml::CUE_CLUSTER_POSITION, cue.cluster_pos)?;
|
||||
ebml::end_master(&mut self.writer, ctp_pos)?;
|
||||
ebml::end_master(&mut self.writer, cp_pos)?;
|
||||
}
|
||||
ebml::end_master(&mut self.writer, cues_pos)?;
|
||||
}
|
||||
|
||||
self.writer.flush()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Fill in a deferred codecPrivate for a track.
|
||||
/// Seeks back to the placeholder, writes the actual data, restores position.
|
||||
pub fn fill_codec_private(&mut self, track_idx: usize, data: &[u8]) -> io::Result<()> {
|
||||
if track_idx >= self.codec_private_filled.len() || self.codec_private_filled[track_idx] {
|
||||
return Ok(());
|
||||
}
|
||||
if let Some((pos, max_size)) = self.codec_private_slots[track_idx] {
|
||||
if data.len() > max_size {
|
||||
// Data too large for reserved space — can't fill in place
|
||||
// This shouldn't happen with 256 bytes reserved
|
||||
return Ok(());
|
||||
}
|
||||
let current = self.writer.stream_position()?;
|
||||
self.writer.seek(SeekFrom::Start(pos))?;
|
||||
// Rewrite: element ID + size + data + zero-pad remainder
|
||||
let mut padded = data.to_vec();
|
||||
padded.resize(max_size, 0);
|
||||
ebml::write_binary(&mut self.writer, ebml::CODEC_PRIVATE, &padded)?;
|
||||
self.writer.seek(SeekFrom::Start(current))?;
|
||||
self.codec_private_filled[track_idx] = true;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn start_cluster(&mut self, ts_ms: i64) -> io::Result<()> {
|
||||
// Close previous cluster if open
|
||||
if self.cluster_open {
|
||||
self.end_cluster()?;
|
||||
}
|
||||
self.cluster_pos = self.writer.stream_position()?;
|
||||
self.cluster_size_pos = ebml::start_master(&mut self.writer, ebml::CLUSTER)?;
|
||||
ebml::write_uint(&mut self.writer, ebml::CLUSTER_TIMESTAMP, ts_ms as u64)?;
|
||||
self.cluster_ts_ms = ts_ms;
|
||||
self.cluster_open = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn end_cluster(&mut self) -> io::Result<()> {
|
||||
if self.cluster_open {
|
||||
ebml::end_master(&mut self.writer, self.cluster_size_pos)?;
|
||||
self.cluster_open = false;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn write_simple_block(&mut self, track_num: usize, relative_ts: i16, keyframe: bool, data: &[u8]) -> io::Result<()> {
|
||||
// SimpleBlock: [track_number VINT] [relative_ts i16] [flags u8] [data]
|
||||
// Track number as EBML VINT
|
||||
let track_vint = if track_num < 0x80 {
|
||||
vec![(track_num as u8) | 0x80]
|
||||
} else {
|
||||
vec![0x40 | ((track_num >> 8) as u8), track_num as u8]
|
||||
};
|
||||
|
||||
let flags: u8 = if keyframe { 0x80 } else { 0x00 };
|
||||
|
||||
let block_size = track_vint.len() + 2 + 1 + data.len(); // vint + ts(2) + flags(1) + data
|
||||
ebml::write_id(&mut self.writer, ebml::SIMPLE_BLOCK)?;
|
||||
ebml::write_size(&mut self.writer, block_size as u64)?;
|
||||
self.writer.write_all(&track_vint)?;
|
||||
self.writer.write_all(&relative_ts.to_be_bytes())?;
|
||||
self.writer.write_all(&[flags])?;
|
||||
self.writer.write_all(data)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ============================================================
|
||||
// Helpers
|
||||
// ============================================================
|
||||
|
||||
fn parse_resolution(s: &str) -> (u32, u32) {
|
||||
if s.contains("2160") { (3840, 2160) }
|
||||
else if s.contains("1080") { (1920, 1080) }
|
||||
else if s.contains("720") { (1280, 720) }
|
||||
else if s.contains("576") { (720, 576) }
|
||||
else if s.contains("480") { (720, 480) }
|
||||
else { (1920, 1080) }
|
||||
}
|
||||
|
||||
fn parse_sample_rate(s: &str) -> f64 {
|
||||
if s.contains("96") { 96000.0 }
|
||||
else if s.contains("192") { 192000.0 }
|
||||
else { 48000.0 }
|
||||
}
|
||||
|
||||
fn parse_channels(s: &str) -> u8 {
|
||||
if s.contains("7.1") { 8 }
|
||||
else if s.contains("5.1") { 6 }
|
||||
else if s.contains("stereo") || s.contains("2.0") { 2 }
|
||||
else if s.contains("mono") { 1 }
|
||||
else { 6 }
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
//! MKV muxing pipeline.
|
||||
//!
|
||||
//! Provides BD transport stream → MKV remuxing via composable streams.
|
||||
//!
|
||||
//! The main type is `MkvStream` — wraps any `Write + Seek` output,
|
||||
//! receives raw BD-TS bytes via `write()`, outputs MKV.
|
||||
//!
|
||||
//! ```text
|
||||
//! disc.rip(title, MkvStream::new(file, &title))
|
||||
//! ```
|
||||
//!
|
||||
//! Components (for advanced use):
|
||||
//! - `ts`: BD transport stream demuxer (192-byte packets → PES frames)
|
||||
//! - `ebml`: EBML write primitives for Matroska container
|
||||
//! - `mkv`: MKV muxer (tracks, clusters, blocks, cues)
|
||||
//! - `codec`: Elementary stream parsers (frame boundaries, codec headers)
|
||||
|
||||
pub mod ebml;
|
||||
pub mod ts;
|
||||
pub mod mkv;
|
||||
pub mod codec;
|
||||
pub mod lookahead;
|
||||
pub mod stream;
|
||||
|
||||
pub use stream::MkvStream;
|
||||
@@ -0,0 +1,218 @@
|
||||
//! MkvStream — a Write adapter that demuxes BD-TS and writes MKV.
|
||||
//!
|
||||
//! ```rust
|
||||
//! let output = MkvStream::new(file)
|
||||
//! .title(&disc.titles[0])
|
||||
//! .max_buffer(20 * 1024 * 1024);
|
||||
//!
|
||||
//! disc.rip(0, output)?;
|
||||
//! ```
|
||||
|
||||
use std::io::{self, Write, Seek};
|
||||
use super::ts::{TsDemuxer, PesPacket};
|
||||
use super::mkv::{MkvMuxer, MkvTrack};
|
||||
use super::codec::{self, CodecParser};
|
||||
use super::lookahead::{LookaheadBuffer, LookaheadState, DEFAULT_LOOKAHEAD_SIZE};
|
||||
use crate::disc::{Stream, Title};
|
||||
|
||||
/// Phase of the MkvStream.
|
||||
#[derive(Debug, Clone, Copy, PartialEq)]
|
||||
enum Phase {
|
||||
/// Collecting data in lookahead buffer, scanning for codec setup.
|
||||
Scanning,
|
||||
/// Header written, streaming directly to muxer.
|
||||
Streaming,
|
||||
}
|
||||
|
||||
/// MKV output stream. Implements `Write`.
|
||||
pub struct MkvStream<W: Write + Seek> {
|
||||
demuxer: TsDemuxer,
|
||||
muxer: Option<MkvMuxer<W>>,
|
||||
writer: Option<W>,
|
||||
parsers: Vec<(u16, Box<dyn CodecParser>)>,
|
||||
pid_to_track: Vec<(u16, usize)>,
|
||||
tracks: Vec<MkvTrack>,
|
||||
title_name: String,
|
||||
duration_secs: f64,
|
||||
lookahead: LookaheadBuffer,
|
||||
phase: Phase,
|
||||
video_tracks_pending: usize,
|
||||
}
|
||||
|
||||
impl<W: Write + Seek> MkvStream<W> {
|
||||
/// Create a new MkvStream wrapping an output writer.
|
||||
pub fn new(writer: W) -> Self {
|
||||
Self {
|
||||
demuxer: TsDemuxer::new(&[]),
|
||||
muxer: None,
|
||||
writer: Some(writer),
|
||||
parsers: Vec::new(),
|
||||
pid_to_track: Vec::new(),
|
||||
tracks: Vec::new(),
|
||||
title_name: String::new(),
|
||||
duration_secs: 0.0,
|
||||
lookahead: LookaheadBuffer::new(DEFAULT_LOOKAHEAD_SIZE),
|
||||
phase: Phase::Scanning,
|
||||
video_tracks_pending: 0,
|
||||
}
|
||||
}
|
||||
|
||||
/// Set the title metadata (streams, duration, name). Returns self.
|
||||
pub fn title(mut self, title: &Title) -> Self {
|
||||
let mut pids = Vec::new();
|
||||
|
||||
for stream in &title.streams {
|
||||
let (pid, track, parser) = match stream {
|
||||
Stream::Video(v) => {
|
||||
self.video_tracks_pending += 1;
|
||||
(v.pid, MkvTrack::video(v), codec::parser_for_codec(v.codec))
|
||||
}
|
||||
Stream::Audio(a) => (a.pid, MkvTrack::audio(a), codec::parser_for_codec(a.codec)),
|
||||
Stream::Subtitle(s) => (s.pid, MkvTrack::subtitle(s), codec::parser_for_codec(s.codec)),
|
||||
};
|
||||
let track_idx = self.tracks.len();
|
||||
pids.push(pid);
|
||||
self.pid_to_track.push((pid, track_idx));
|
||||
self.parsers.push((pid, parser));
|
||||
self.tracks.push(track);
|
||||
}
|
||||
|
||||
self.demuxer = TsDemuxer::new(&pids);
|
||||
self.title_name = title.playlist.clone();
|
||||
self.duration_secs = title.duration_secs;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the lookahead buffer size in bytes. Default 5 MB. Returns self.
|
||||
pub fn max_buffer(mut self, size: usize) -> Self {
|
||||
self.lookahead = LookaheadBuffer::new(size);
|
||||
self
|
||||
}
|
||||
|
||||
/// Finalize the MKV file — close cluster, write cues.
|
||||
pub fn finish(mut self) -> io::Result<()> {
|
||||
if let Some(ref mut muxer) = self.muxer {
|
||||
let remaining = self.demuxer.flush();
|
||||
for pes in &remaining {
|
||||
Self::process_one_pes(&self.pid_to_track, &mut self.parsers, muxer, pes)?;
|
||||
}
|
||||
}
|
||||
if let Some(muxer) = self.muxer {
|
||||
muxer.finish()?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn check_codec_private(&mut self) -> bool {
|
||||
if self.video_tracks_pending == 0 {
|
||||
return true;
|
||||
}
|
||||
for (pid, parser) in &self.parsers {
|
||||
if let Some(cp) = parser.codec_private() {
|
||||
if let Some((_, track_idx)) = self.pid_to_track.iter().find(|(p, _)| p == pid) {
|
||||
if self.tracks[*track_idx].codec_private.is_none() {
|
||||
self.tracks[*track_idx].codec_private = Some(cp);
|
||||
self.video_tracks_pending -= 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
self.video_tracks_pending == 0
|
||||
}
|
||||
|
||||
fn start_streaming(&mut self) -> io::Result<()> {
|
||||
let writer = self.writer.take().ok_or_else(|| {
|
||||
io::Error::new(io::ErrorKind::Other, "writer already consumed")
|
||||
})?;
|
||||
|
||||
let muxer = MkvMuxer::new(writer, &self.tracks, Some(&self.title_name), self.duration_secs)?;
|
||||
self.muxer = Some(muxer);
|
||||
self.phase = Phase::Streaming;
|
||||
|
||||
// Re-parse and write buffered data
|
||||
let buffered = self.lookahead.drain();
|
||||
if !buffered.is_empty() {
|
||||
let pids: Vec<u16> = self.pid_to_track.iter().map(|(pid, _)| *pid).collect();
|
||||
let mut temp_demuxer = TsDemuxer::new(&pids);
|
||||
let mut packets = temp_demuxer.feed(&buffered);
|
||||
packets.extend(temp_demuxer.flush());
|
||||
|
||||
if let Some(ref mut muxer) = self.muxer {
|
||||
for pes in &packets {
|
||||
Self::process_one_pes(&self.pid_to_track, &mut self.parsers, muxer, pes)?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn process_one_pes(
|
||||
pid_to_track: &[(u16, usize)],
|
||||
parsers: &mut [(u16, Box<dyn CodecParser>)],
|
||||
muxer: &mut MkvMuxer<W>,
|
||||
pes: &PesPacket,
|
||||
) -> io::Result<()> {
|
||||
let track_idx = match pid_to_track.iter().find(|(pid, _)| *pid == pes.pid) {
|
||||
Some((_, idx)) => *idx,
|
||||
None => return Ok(()),
|
||||
};
|
||||
let parser = match parsers.iter_mut().find(|(pid, _)| *pid == pes.pid) {
|
||||
Some((_, p)) => p,
|
||||
None => return Ok(()),
|
||||
};
|
||||
let frames = parser.parse(pes);
|
||||
for frame in frames {
|
||||
muxer.write_frame(track_idx, frame.pts_ns, frame.keyframe, &frame.data)?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl<W: Write + Seek> Write for MkvStream<W> {
|
||||
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
|
||||
match self.phase {
|
||||
Phase::Scanning => {
|
||||
// Parse for codec info
|
||||
let packets = self.demuxer.feed(buf);
|
||||
for pes in &packets {
|
||||
if let Some((_, parser)) = self.parsers.iter_mut().find(|(p, _)| *p == pes.pid) {
|
||||
let _ = parser.parse(pes);
|
||||
}
|
||||
}
|
||||
|
||||
// Try to buffer
|
||||
let state = self.lookahead.push(buf);
|
||||
|
||||
// Check if we have everything
|
||||
if self.check_codec_private() {
|
||||
self.lookahead.mark_ready();
|
||||
self.start_streaming()?;
|
||||
return Ok(buf.len());
|
||||
}
|
||||
|
||||
match state {
|
||||
LookaheadState::Collecting => Ok(buf.len()),
|
||||
LookaheadState::Overflow => Err(io::Error::new(
|
||||
io::ErrorKind::OutOfMemory,
|
||||
"MKV lookahead buffer overflow — no codec data found within buffer limit",
|
||||
)),
|
||||
LookaheadState::Ready => Ok(buf.len()),
|
||||
}
|
||||
}
|
||||
Phase::Streaming => {
|
||||
let packets = self.demuxer.feed(buf);
|
||||
if let Some(ref mut muxer) = self.muxer {
|
||||
for pes in &packets {
|
||||
Self::process_one_pes(&self.pid_to_track, &mut self.parsers, muxer, pes)?;
|
||||
}
|
||||
}
|
||||
Ok(buf.len())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn flush(&mut self) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
+269
@@ -0,0 +1,269 @@
|
||||
//! BD Transport Stream demuxer.
|
||||
//!
|
||||
//! Blu-ray uses 192-byte TS packets (not standard 188):
|
||||
//! - 4-byte TP_extra_header (arrival timestamp + copy permission)
|
||||
//! - 188-byte standard MPEG-TS packet
|
||||
//!
|
||||
//! This demuxer extracts PES packets from selected PIDs, with PTS/DTS timestamps.
|
||||
|
||||
/// BD transport stream packet size (4-byte extra header + 188-byte TS).
|
||||
const BD_TS_PACKET_SIZE: usize = 192;
|
||||
|
||||
/// Standard TS packet size.
|
||||
const TS_PACKET_SIZE: usize = 188;
|
||||
|
||||
/// TS sync byte.
|
||||
const SYNC_BYTE: u8 = 0x47;
|
||||
|
||||
/// A reassembled PES packet with timestamp info.
|
||||
#[derive(Debug)]
|
||||
pub struct PesPacket {
|
||||
/// MPEG-TS PID this packet belongs to.
|
||||
pub pid: u16,
|
||||
/// Presentation timestamp in 90kHz ticks (if present).
|
||||
pub pts: Option<i64>,
|
||||
/// Decode timestamp in 90kHz ticks (if present).
|
||||
pub dts: Option<i64>,
|
||||
/// Elementary stream data (video frame, audio frame, subtitle segment, etc.).
|
||||
pub data: Vec<u8>,
|
||||
}
|
||||
|
||||
/// Per-PID PES reassembly state.
|
||||
struct PesAssembler {
|
||||
pid: u16,
|
||||
buffer: Vec<u8>,
|
||||
pts: Option<i64>,
|
||||
dts: Option<i64>,
|
||||
active: bool,
|
||||
}
|
||||
|
||||
impl PesAssembler {
|
||||
fn new(pid: u16) -> Self {
|
||||
Self {
|
||||
pid,
|
||||
buffer: Vec::with_capacity(256 * 1024),
|
||||
pts: None,
|
||||
dts: None,
|
||||
active: false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Start a new PES packet. Returns the completed previous packet (if any).
|
||||
fn start(&mut self, pts: Option<i64>, dts: Option<i64>) -> Option<PesPacket> {
|
||||
let completed = if self.active && !self.buffer.is_empty() {
|
||||
Some(PesPacket {
|
||||
pid: self.pid,
|
||||
pts: self.pts,
|
||||
dts: self.dts,
|
||||
data: std::mem::replace(&mut self.buffer, Vec::with_capacity(256 * 1024)),
|
||||
})
|
||||
} else {
|
||||
self.buffer.clear();
|
||||
None
|
||||
};
|
||||
self.pts = pts;
|
||||
self.dts = dts;
|
||||
self.active = true;
|
||||
completed
|
||||
}
|
||||
|
||||
/// Append payload data to the current PES packet.
|
||||
fn push(&mut self, data: &[u8]) {
|
||||
if self.active {
|
||||
self.buffer.extend_from_slice(data);
|
||||
}
|
||||
}
|
||||
|
||||
/// Flush remaining data as a PES packet.
|
||||
fn flush(&mut self) -> Option<PesPacket> {
|
||||
if self.active && !self.buffer.is_empty() {
|
||||
self.active = false;
|
||||
Some(PesPacket {
|
||||
pid: self.pid,
|
||||
pts: self.pts,
|
||||
dts: self.dts,
|
||||
data: std::mem::take(&mut self.buffer),
|
||||
})
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// BD Transport Stream demuxer.
|
||||
pub struct TsDemuxer {
|
||||
assemblers: Vec<PesAssembler>,
|
||||
pid_index: [i16; 8192], // PID → index into assemblers, -1 = not tracked
|
||||
}
|
||||
|
||||
impl TsDemuxer {
|
||||
/// Create a new demuxer tracking the given PIDs.
|
||||
pub fn new(pids: &[u16]) -> Self {
|
||||
let mut pid_index = [-1i16; 8192];
|
||||
let mut assemblers = Vec::with_capacity(pids.len());
|
||||
for (i, &pid) in pids.iter().enumerate() {
|
||||
pid_index[pid as usize] = i as i16;
|
||||
assemblers.push(PesAssembler::new(pid));
|
||||
}
|
||||
Self { assemblers, pid_index }
|
||||
}
|
||||
|
||||
/// Feed a chunk of BD transport stream data (must be aligned to 192-byte packets).
|
||||
/// Returns completed PES packets.
|
||||
pub fn feed(&mut self, data: &[u8]) -> Vec<PesPacket> {
|
||||
let mut completed = Vec::new();
|
||||
let mut offset = 0;
|
||||
|
||||
while offset + BD_TS_PACKET_SIZE <= data.len() {
|
||||
let packet = &data[offset..offset + BD_TS_PACKET_SIZE];
|
||||
offset += BD_TS_PACKET_SIZE;
|
||||
|
||||
// Skip 4-byte TP_extra_header, check sync byte
|
||||
if packet[4] != SYNC_BYTE {
|
||||
continue;
|
||||
}
|
||||
|
||||
let ts = &packet[4..]; // 188-byte standard TS packet
|
||||
|
||||
// Parse TS header
|
||||
let pid = (((ts[1] & 0x1F) as u16) << 8) | ts[2] as u16;
|
||||
let pusi = ts[1] & 0x40 != 0; // Payload Unit Start Indicator
|
||||
let adaptation = (ts[3] >> 4) & 0x03;
|
||||
|
||||
// Check if we're tracking this PID
|
||||
let idx = self.pid_index[pid as usize];
|
||||
if idx < 0 {
|
||||
continue;
|
||||
}
|
||||
let asm = &mut self.assemblers[idx as usize];
|
||||
|
||||
// Find payload start (skip adaptation field if present)
|
||||
let payload_start = if adaptation == 0x03 || adaptation == 0x02 {
|
||||
// Adaptation field present
|
||||
let af_len = ts[4] as usize;
|
||||
5 + af_len
|
||||
} else {
|
||||
4
|
||||
};
|
||||
|
||||
if payload_start >= TS_PACKET_SIZE {
|
||||
continue;
|
||||
}
|
||||
|
||||
// No payload
|
||||
if adaptation == 0x02 {
|
||||
continue;
|
||||
}
|
||||
|
||||
let payload = &ts[payload_start..];
|
||||
|
||||
if pusi {
|
||||
// New PES packet starts here — parse PES header
|
||||
let (pts, dts, pes_data_start) = parse_pes_header(payload);
|
||||
if let Some(prev) = asm.start(pts, dts) {
|
||||
completed.push(prev);
|
||||
}
|
||||
if pes_data_start < payload.len() {
|
||||
asm.push(&payload[pes_data_start..]);
|
||||
}
|
||||
} else {
|
||||
// Continuation of current PES packet
|
||||
asm.push(payload);
|
||||
}
|
||||
}
|
||||
|
||||
completed
|
||||
}
|
||||
|
||||
/// Flush all assemblers, returning any remaining PES packets.
|
||||
pub fn flush(&mut self) -> Vec<PesPacket> {
|
||||
let mut completed = Vec::new();
|
||||
for asm in &mut self.assemblers {
|
||||
if let Some(pkt) = asm.flush() {
|
||||
completed.push(pkt);
|
||||
}
|
||||
}
|
||||
completed
|
||||
}
|
||||
}
|
||||
|
||||
/// Parse a PES packet header, extracting PTS and DTS.
|
||||
/// Returns (pts, dts, offset_to_elementary_stream_data).
|
||||
fn parse_pes_header(data: &[u8]) -> (Option<i64>, Option<i64>, usize) {
|
||||
// PES packet: 00 00 01 [stream_id] [length:2] [flags...]
|
||||
if data.len() < 9 || data[0] != 0x00 || data[1] != 0x00 || data[2] != 0x01 {
|
||||
return (None, None, 0);
|
||||
}
|
||||
|
||||
let stream_id = data[3];
|
||||
|
||||
// Some stream IDs don't have the standard PES header extension
|
||||
// (program_stream_map, padding, private_stream_2, ECM, EMM, etc.)
|
||||
if stream_id == 0xBC || stream_id == 0xBE || stream_id == 0xBF
|
||||
|| stream_id == 0xF0 || stream_id == 0xF1 || stream_id == 0xFF
|
||||
{
|
||||
return (None, None, 6);
|
||||
}
|
||||
|
||||
// Standard PES header: [6] = flags1, [7] = flags2, [8] = header_data_length
|
||||
if data.len() < 9 {
|
||||
return (None, None, 6);
|
||||
}
|
||||
|
||||
let pts_dts_flags = (data[7] >> 6) & 0x03;
|
||||
let header_data_len = data[8] as usize;
|
||||
let data_start = 9 + header_data_len;
|
||||
|
||||
let mut pts = None;
|
||||
let mut dts = None;
|
||||
|
||||
if pts_dts_flags >= 2 && data.len() >= 14 {
|
||||
pts = Some(parse_timestamp(&data[9..14]));
|
||||
}
|
||||
if pts_dts_flags == 3 && data.len() >= 19 {
|
||||
dts = Some(parse_timestamp(&data[14..19]));
|
||||
}
|
||||
|
||||
(pts, dts, data_start)
|
||||
}
|
||||
|
||||
/// Parse a 5-byte PTS/DTS timestamp (33 bits in 90kHz).
|
||||
fn parse_timestamp(data: &[u8]) -> i64 {
|
||||
let b0 = data[0] as i64;
|
||||
let b1 = data[1] as i64;
|
||||
let b2 = data[2] as i64;
|
||||
let b3 = data[3] as i64;
|
||||
let b4 = data[4] as i64;
|
||||
|
||||
((b0 >> 1) & 0x07) << 30
|
||||
| b1 << 22
|
||||
| (b2 >> 1) << 15
|
||||
| b3 << 7
|
||||
| b4 >> 1
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_parse_timestamp() {
|
||||
// Example: PTS = 0 → encoded as 21 00 01 00 01
|
||||
let data = [0x21, 0x00, 0x01, 0x00, 0x01];
|
||||
assert_eq!(parse_timestamp(&data), 0);
|
||||
|
||||
// Example: PTS = 90000 (1 second at 90kHz)
|
||||
// Manual encoding: 33 bits = 0x00015F90
|
||||
// This is just a sanity check that the parser doesn't crash
|
||||
let data2 = [0x21, 0x00, 0x07, 0xE9, 0x01]; // approximate
|
||||
let pts = parse_timestamp(&data2);
|
||||
assert!(pts >= 0);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_demuxer_empty() {
|
||||
let mut demux = TsDemuxer::new(&[0x1011]);
|
||||
let result = demux.feed(&[]);
|
||||
assert!(result.is_empty());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user