use std::collections::BTreeMap;
use std::io::{Read, Seek, SeekFrom};
use crate::error::{Error, Result};
use crate::mkv::cluster::{ClusterEnd, RawFrame, scan_cluster};
use crate::mkv::ebml::{Element, child_path, indexed_path};
use crate::mkv::ids;
use crate::mkv::segment::{element, open_segment, read_whole};
use crate::mkv::source::Source;
use crate::track::TrackKind;
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub struct TrackTiming {
pub track_id: u32,
pub codec_delay: u64,
pub discard_padding: u64,
}
pub fn first_track_timing<R: Read + Seek>(
mut reader: R,
kind: TrackKind,
) -> Result<Option<TrackTiming>> {
let input_len = reader
.seek(SeekFrom::End(0))
.map_err(|err| Error::io("the input length", 0, err))?;
let mut src = Source::new(&mut reader, input_len);
let start = open_segment(&mut src)?;
let end = start.declared_end.min(input_len);
let mut tracks = None;
let mut padding = BTreeMap::new();
let mut frames = Vec::new();
let mut clusters = 0u32;
let mut at = start.data;
while at < end {
let Some(header) = src.header_at(at, "Segment")? else {
break;
};
if matches!(header.id, ids::EBML | ids::SEGMENT) {
break;
}
if header.id == ids::CLUSTER {
clusters = clusters.saturating_add(1);
let scanned = scan_cluster(&mut src, at, header, clusters, &mut frames)?;
add_padding(&mut frames, &mut padding);
match scanned {
ClusterEnd::At(next) => {
at = next;
continue;
}
ClusterEnd::Truncated => break,
}
}
let path = format!("Segment/{}", ids::name(header.id));
let next = (at + header.len).saturating_add(header.known_size(&path, at)?);
if next > input_len {
break;
}
if header.id == ids::TRACKS && tracks.is_none() {
let (body, bytes) = read_whole(&mut src, at, header, &path)?;
tracks = Some(first_of_kind(&element(header.id, at, body, &bytes), kind)?);
}
at = next;
}
let tracks = tracks.ok_or_else(|| Error::mkv_missing("Segment/Tracks"))?;
Ok(tracks.map(|(track_id, codec_delay)| TrackTiming {
track_id,
codec_delay,
discard_padding: padding.get(&u64::from(track_id)).copied().unwrap_or(0),
}))
}
fn add_padding(frames: &mut Vec<RawFrame>, padding: &mut BTreeMap<u64, u64>) {
for frame in frames.drain(..) {
if let Some(value) = frame.discard_padding.filter(|&value| value > 0) {
let sum = padding.entry(frame.track).or_default();
*sum = sum.saturating_add(value.unsigned_abs());
}
}
}
fn first_of_kind(tracks: &Element<'_>, kind: TrackKind) -> Result<Option<(u32, u64)>> {
let path = "Segment/Tracks";
let mut entries = 0usize;
for child in tracks.children(path) {
let child = child?;
if child.id != ids::TRACK_ENTRY {
continue;
}
entries += 1;
let entry_path = indexed_path(path, ids::TRACK_ENTRY, entries);
let (mut number, mut track_type, mut codec_delay) = (None, None, 0);
for field in child.children(&entry_path) {
let field = field?;
let at = child_path(&entry_path, field.id);
match field.id {
ids::TRACK_NUMBER => number = Some(field.uint(&at)?),
ids::TRACK_TYPE => track_type = Some(field.uint(&at)?),
ids::CODEC_DELAY => codec_delay = field.uint(&at)?,
_ => {}
}
}
if track_type.map(TrackKind::from_matroska_type) != Some(kind) {
continue;
}
let at = format!("{entry_path}/TrackNumber");
let number = number.ok_or_else(|| Error::mkv_missing(at.clone()))?;
let id = u32::try_from(number)
.ok()
.filter(|&id| id != 0)
.ok_or_else(|| {
Error::mkv_malformed(
&at,
child.offset,
format!("track number {number} is outside 1 to {}", u32::MAX),
)
})?;
return Ok(Some((id, codec_delay)));
}
Ok(None)
}