From e45a036dba5eb99ed814bde98689b9a331a190a9 Mon Sep 17 00:00:00 2001 From: Serial <69764315+Serial-ATA@users.noreply.github.com> Date: Tue, 15 Oct 2024 09:03:23 -0400 Subject: [PATCH] EBML: Support duration and bitrate calculation --- lofty/src/ebml/element_reader.rs | 29 ++++ lofty/src/ebml/properties.rs | 72 ++++++++-- lofty/src/ebml/read.rs | 1 + lofty/src/ebml/read/segment.rs | 9 +- lofty/src/ebml/read/segment_cluster.rs | 175 +++++++++++++++++++++++++ lofty/src/ebml/read/segment_info.rs | 37 ++++-- lofty/src/properties/tests.rs | 2 + 7 files changed, 305 insertions(+), 20 deletions(-) create mode 100644 lofty/src/ebml/read/segment_cluster.rs diff --git a/lofty/src/ebml/element_reader.rs b/lofty/src/ebml/element_reader.rs index 4b882e37..1bf0a870 100644 --- a/lofty/src/ebml/element_reader.rs +++ b/lofty/src/ebml/element_reader.rs @@ -109,6 +109,7 @@ ebml_master_elements! { children: [ // SeekHead: { 0x114D_9B74, Master }, Info: { 0x1549_A966, Master }, + Cluster: { 0x1F43_B675, Master }, Tracks: { 0x1654_AE6B, Master }, Tags: { 0x1254_C367, Master }, Attachments: { 0x1941_A469, Master }, @@ -131,9 +132,28 @@ ebml_master_elements! { TimecodeScale: { 0x2AD7_B1, UnsignedInt }, MuxingApp: { 0x4D80, Utf8 }, WritingApp: { 0x5741, Utf8 }, + Duration: { 0x4489, Float }, ], }, + // segment.cluster + Cluster: { + id: 0x1F43_B675, + children: [ + Timestamp: { 0xE7, UnsignedInt }, + SimpleBlock: { 0xA3, Binary }, + BlockGroup: { 0xA0, Master }, + ], + }, + + // segment.cluster.blockGroup + BlockGroup: { + id: 0xA0, + children: [ + Block: { 0xA1, Binary }, + ] + }, + // segment.tracks Tracks: { id: 0x1654_AE6B, @@ -726,6 +746,15 @@ where } } +impl<'a, R> Read for ElementChildIterator<'a, R> +where + R: Read, +{ + fn read(&mut self, buf: &mut [u8]) -> io::Result { + self.reader.read(buf) + } +} + impl<'a, R> Deref for ElementChildIterator<'a, R> where R: Read, diff --git a/lofty/src/ebml/properties.rs b/lofty/src/ebml/properties.rs index 74c697b1..48c4b976 100644 --- a/lofty/src/ebml/properties.rs +++ b/lofty/src/ebml/properties.rs @@ -1,6 +1,8 @@ use super::Language; use crate::properties::FileProperties; +use std::time::Duration; + /// Properties from the EBML header /// /// These are present for all EBML formats. @@ -77,6 +79,7 @@ pub struct SegmentInfo { pub(crate) timestamp_scale: u64, pub(crate) muxing_app: String, pub(crate) writing_app: String, + pub(crate) duration: Option, } impl SegmentInfo { @@ -100,6 +103,14 @@ impl SegmentInfo { pub fn writing_app(&self) -> &str { &self.writing_app } + + /// The duration of the segment + /// + /// NOTE: This information is not always present in the segment, in which case + /// [`EbmlProperties::duration`] should be used. + pub fn duration(&self) -> Option { + self.duration + } } impl Default for SegmentInfo { @@ -109,6 +120,7 @@ impl Default for SegmentInfo { timestamp_scale: 1_000_000, muxing_app: String::new(), writing_app: String::new(), + duration: None, } } } @@ -210,11 +222,15 @@ impl AudioTrackDescriptor { /// Settings for an audio track #[derive(Debug, Clone, PartialEq, Default)] pub struct AudioTrackSettings { + // Provided to us for free pub(crate) sampling_frequency: f64, pub(crate) output_sampling_frequency: f64, pub(crate) channels: u8, pub(crate) bit_depth: Option, pub(crate) emphasis: Option, + + // Need to be calculated + pub(crate) bitrate: Option, } impl AudioTrackSettings { @@ -322,27 +338,67 @@ impl EbmlProperties { /// Information about the default audio track /// - /// The information is extracted from the first audio track with its default flag set - /// in the `\Segment\Tracks` element. + /// The "default" track is selected as: + /// 1. The first audio track with its `default` flag set + /// 2. If 1 fails, just grab the first audio track with its `enabled` flag set pub fn default_audio_track(&self) -> Option<&AudioTrackDescriptor> { - self.audio_tracks.iter().find(|track| track.default) + if let Some(position) = self.default_audio_track_position() { + return self.audio_tracks.get(position); + } + + None + } + + // TODO: Actually calculate from cluster + /// The duration of the default audio track + /// + /// NOTE: see [`EbmlProperties::default_audio_track`] + /// + /// This will always use the duration written in `\Segment\Info` if present. Otherwise, it will + /// be manually calculated using `\Segment\Cluster` data. + pub fn duration(&self) -> Duration { + self.segment_info.duration().unwrap() + } + + /// Audio bitrate (kbps) + /// + /// NOTE: This is the bitrate of the default audio track see [`EbmlProperties::default_audio_track`] + /// for what this means. + pub fn bitrate(&self) -> Option { + self.default_audio_track() + .and_then(|track| track.settings.bitrate) + } + + pub(crate) fn default_audio_track_position(&self) -> Option { + self.audio_tracks + .iter() + .position(|track| track.default) + .or_else(|| { + // Otherwise, it's normal to just pick the first enabled track + self.audio_tracks.iter().position(|track| track.enabled) + }) } } impl From for FileProperties { fn from(input: EbmlProperties) -> Self { let Some(default_audio_track) = input.default_audio_track() else { - return FileProperties::default(); + let mut properties = FileProperties::default(); + if let Some(duration) = input.segment_info.duration { + properties.duration = duration; + } + + return properties; }; Self { - duration: todo!("Support duration"), - overall_bitrate: todo!("Support bitrate"), - audio_bitrate: todo!("Support bitrate"), + duration: input.duration(), + overall_bitrate: input.bitrate(), + audio_bitrate: input.bitrate(), sample_rate: Some(default_audio_track.settings.sampling_frequency as u32), bit_depth: default_audio_track.settings.bit_depth, channels: Some(default_audio_track.settings.channels), - channel_mask: todo!("Channel mask"), + channel_mask: None, // TODO: Will require reading into track data } } } diff --git a/lofty/src/ebml/read.rs b/lofty/src/ebml/read.rs index 2a6856d8..0b032505 100644 --- a/lofty/src/ebml/read.rs +++ b/lofty/src/ebml/read.rs @@ -1,6 +1,7 @@ mod segment; mod segment_attachments; mod segment_chapters; +mod segment_cluster; mod segment_info; mod segment_tags; mod segment_tracks; diff --git a/lofty/src/ebml/read/segment.rs b/lofty/src/ebml/read/segment.rs index 1709ccb1..c2718f2d 100644 --- a/lofty/src/ebml/read/segment.rs +++ b/lofty/src/ebml/read/segment.rs @@ -1,4 +1,4 @@ -use super::{segment_attachments, segment_info, segment_tags, segment_tracks}; +use super::{segment_attachments, segment_cluster, segment_info, segment_tags, segment_tracks}; use crate::config::ParseOptions; use crate::ebml::element_reader::{ElementHeader, ElementIdent, ElementReader, ElementReaderYield}; use crate::ebml::properties::EbmlProperties; @@ -30,6 +30,13 @@ where properties, )?; }, + ElementIdent::Cluster if parse_options.read_properties => { + segment_cluster::read_from( + &mut children_reader.children(), + parse_options, + properties, + )? + }, ElementIdent::Tracks if parse_options.read_properties => { segment_tracks::read_from( &mut children_reader.children(), diff --git a/lofty/src/ebml/read/segment_cluster.rs b/lofty/src/ebml/read/segment_cluster.rs new file mode 100644 index 00000000..d82853a9 --- /dev/null +++ b/lofty/src/ebml/read/segment_cluster.rs @@ -0,0 +1,175 @@ +use crate::config::ParseOptions; +use crate::ebml::element_reader::{ + ChildElementDescriptor, ElementChildIterator, ElementIdent, ElementReaderYield, +}; +use crate::ebml::properties::EbmlProperties; +use crate::ebml::{AudioTrackDescriptor, VInt}; +use crate::error::Result; + +use std::io::{Read, Seek}; + +pub(super) fn read_from( + children_reader: &mut ElementChildIterator<'_, R>, + parse_options: ParseOptions, + properties: &mut EbmlProperties, +) -> Result<()> +where + R: Read + Seek, +{ + // TODO: Support Tracks appearing after Cluster (should implement SeekHead first) + let Some(default_audio_track_position) = properties.default_audio_track_position() else { + log::warn!( + "No default audio track found (does \\Segment\\Cluster appear before \ + \\Segment\\Tracks?)" + ); + children_reader.exhaust_current_master()?; + return Ok(()); + }; + + let default_audio_track = &properties.audio_tracks[default_audio_track_position]; + + let target_track_number = default_audio_track.number(); + let mut total_audio_data_size = 0u64; + + while let Some(child) = children_reader.next()? { + let ident; + let size; + match child { + ElementReaderYield::Master((master_ident, master_size)) => { + ident = master_ident; + size = master_size; + }, + ElementReaderYield::Child((descriptor, child_size)) => { + ident = descriptor.ident; + size = child_size; + }, + ElementReaderYield::Unknown(unknown) => { + children_reader.skip_element(unknown)?; + continue; + }, + ElementReaderYield::Eof => break, + } + + match ident { + ElementIdent::Timestamp => { + // TODO: Fancy timestamp durations + children_reader.skip(size.value())?; + continue; + }, + ElementIdent::SimpleBlock => { + let (block_is_applicable, header_size) = check_block( + children_reader, + parse_options, + size.value(), + target_track_number, + properties.header.max_size_length, + )?; + + if !block_is_applicable { + continue; + } + + total_audio_data_size += (size.value() - header_size as u64); + }, + ElementIdent::BlockGroup => read_block_group( + &mut children_reader.children(), + parse_options, + properties, + target_track_number, + &mut total_audio_data_size, + )?, + _ => unreachable!("Unhandled child element in \\Segment\\Cluster: {child:?}"), + } + } + + if total_audio_data_size == 0 { + log::warn!("No audio data found, audio bitrate will be 0, duration may be 0"); + return Ok(()); + } + + let duration_millis = properties.duration().as_secs() as u128; + if duration_millis == 0 { + log::warn!("Duration is zero, cannot calculate bitrate"); + return Ok(()); + } + + let default_audio_track = &mut properties.audio_tracks[default_audio_track_position]; // TODO + + let bitrate_bps = (((total_audio_data_size as u128) * 8) / duration_millis) as u32; + default_audio_track.settings.bitrate = Some(bitrate_bps / 1000); + + Ok(()) +} + +fn read_block_group( + children_reader: &mut ElementChildIterator<'_, R>, + parse_options: ParseOptions, + properties: &mut EbmlProperties, + target_track_number: u64, + total_audio_data_size: &mut u64, +) -> Result<()> +where + R: Read + Seek, +{ + while let Some(child) = children_reader.next()? { + let size; + match child { + ElementReaderYield::Child(( + ChildElementDescriptor { + ident: ElementIdent::Block, + .. + }, + child_size, + )) => { + size = child_size; + }, + ElementReaderYield::Unknown(unknown) => { + children_reader.skip_element(unknown)?; + continue; + }, + _ => unimplemented!( + "Unhandled child element in \\Segment\\Cluster\\BlockGroup: {child:?}" + ), + } + + let (block_is_applicable, header_size) = check_block( + children_reader, + parse_options, + size.value(), + target_track_number, + properties.header.max_size_length, + )?; + + if !block_is_applicable { + continue; + } + + *total_audio_data_size += (size.value() - header_size as u64); + } + + Ok(()) +} + +fn check_block( + children_reader: &mut ElementChildIterator<'_, R>, + _parse_options: ParseOptions, + block_size: u64, + target_track_number: u64, + max_size_length: u8, +) -> Result<(bool, u8)> +where + R: Read + Seek, +{ + // The block header is Track number (variable), timestamp (i16), and flags (u8) + const NON_VARIABLE_BLOCK_HEADER_SIZE: u8 = 2 /* Timestamp */ + 1 /* Flags */; + + let track_number = VInt::::parse(children_reader, max_size_length)?; + let track_number_octets = track_number.octet_length(); + + children_reader.skip(block_size - track_number_octets as u64)?; + if track_number != target_track_number { + return Ok((false, track_number_octets + NON_VARIABLE_BLOCK_HEADER_SIZE)); + } + + Ok((true, track_number_octets + NON_VARIABLE_BLOCK_HEADER_SIZE)) +} diff --git a/lofty/src/ebml/read/segment_info.rs b/lofty/src/ebml/read/segment_info.rs index 840dc932..a049817c 100644 --- a/lofty/src/ebml/read/segment_info.rs +++ b/lofty/src/ebml/read/segment_info.rs @@ -5,6 +5,7 @@ use crate::error::Result; use crate::macros::decode_err; use std::io::{Read, Seek}; +use std::time::Duration; pub(super) fn read_from( children_reader: &mut ElementChildIterator<'_, R>, @@ -14,6 +15,10 @@ pub(super) fn read_from( where R: Read + Seek, { + // Deal with duration after parsing, in case the timestamp scale appears after it + // for some reason. + let mut duration = None; + while let Some(child) = children_reader.next()? { match child { ElementReaderYield::Master((id, size)) => { @@ -29,21 +34,17 @@ where ElementIdent::TimecodeScale => { properties.segment_info.timestamp_scale = children_reader.read_unsigned_int(size.value())?; - - if properties.segment_info.timestamp_scale == 0 { - log::warn!("Segment.Info.TimecodeScale is 0, which is invalid"); - if parse_options.parsing_mode == ParsingMode::Strict { - decode_err!(@BAIL Ebml, "Segment.Info.TimecodeScale must be nonzero"); - } - } }, ElementIdent::MuxingApp => { - properties.segment_info.muxing_app = - children_reader.read_utf8(size.value())? + let muxing_app = children_reader.read_utf8(size.value())?; + properties.segment_info.muxing_app = muxing_app; }, ElementIdent::WritingApp => { - properties.segment_info.writing_app = - children_reader.read_utf8(size.value())? + let writing_app = children_reader.read_utf8(size.value())?; + properties.segment_info.writing_app = writing_app; + }, + ElementIdent::Duration => { + duration = Some(children_reader.read_float(size.value())?); }, _ => { // We do not end up using information from all of the segment @@ -63,5 +64,19 @@ where } } + if properties.segment_info.timestamp_scale == 0 { + log::warn!("Segment.Info.TimecodeScale is 0, which is invalid"); + if parse_options.parsing_mode == ParsingMode::Strict { + decode_err!(@BAIL Ebml, "Segment.Info.TimecodeScale must be non-zero"); + } + + return Ok(()); + } + + if let Some(duration) = duration { + let scaled_duration = duration * properties.segment_info.timestamp_scale as f64; + properties.segment_info.duration = Some(Duration::from_nanos(scaled_duration as u64)); + } + Ok(()) } diff --git a/lofty/src/properties/tests.rs b/lofty/src/properties/tests.rs index e8abdba7..a458665e 100644 --- a/lofty/src/properties/tests.rs +++ b/lofty/src/properties/tests.rs @@ -88,6 +88,7 @@ fn MKA_PROPERTIES() -> EbmlProperties { timestamp_scale: 1000000, muxing_app: String::from("Lavf60.3.100"), writing_app: String::from("Lavf60.3.100"), + duration: Some(Duration::from_millis(1431)), }, audio_tracks: vec![AudioTrackDescriptor { number: 1, @@ -105,6 +106,7 @@ fn MKA_PROPERTIES() -> EbmlProperties { channels: 2, bit_depth: Some(32), emphasis: None, + bitrate: Some(99), // TODO: FFmpeg reports 97, not bad }, }], }