EBML: Support duration and bitrate calculation

Author: Serial-ATA

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<usize> {
+		self.reader.read(buf)
+	}
+}
+
 impl<'a, R> Deref for ElementChildIterator<'a, R>
 where
 	R: Read,

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<Duration>,
 }
 
 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<Duration> {
+		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<u8>,
 	pub(crate) emphasis: Option<EbmlAudioTrackEmphasis>,
+
+	// Need to be calculated
+	pub(crate) bitrate: Option<u32>,
 }
 
 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<u32> {
+		self.default_audio_track()
+			.and_then(|track| track.settings.bitrate)
+	}
+
+	pub(crate) fn default_audio_track_position(&self) -> Option<usize> {
+		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<EbmlProperties> 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
 		}
 	}
 }

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;

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(),

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<R>(
+	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<R>(
+	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<R>(
+	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::<u64>::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))
+}