diff --git a/crates/audio/src/audio_data.rs b/crates/audio/src/audio_data.rs index 8a02642a14..061e02993b 100644 --- a/crates/audio/src/audio_data.rs +++ b/crates/audio/src/audio_data.rs @@ -89,6 +89,7 @@ impl AudioData { stream_start_time }; let stream_time_base = f64::from(stream_time_base); + let stream_duration = input_stream.duration(); let source_start_sample = range.as_ref().map_or(0, |window| window.start); let covered_source_end_sample = range.as_ref().map_or(usize::MAX, |window| window.end); let mut sought = false; @@ -105,6 +106,14 @@ impl AudioData { let mut decoded_frame = ffmpeg::frame::Audio::empty(); let mut resampler_output = ResamplerOutput::new(); let mut samples = Vec::new(); + let compressed_bytes = std::fs::metadata(path).map(|meta| meta.len()).unwrap_or(0); + let _ = samples.try_reserve(decoded_sample_capacity( + stream_duration, + stream_time_base, + target_channels, + range.as_ref(), + compressed_bytes, + )); let mut resampled_samples = Vec::new(); let mut next_source_sample = if sought { None } else { Some(0usize) }; let mut complete = source_start_sample == covered_source_end_sample; @@ -309,6 +318,47 @@ fn target_channels_for_source(channels: u16) -> u16 { if channels <= 1 { 1 } else { 2 } } +// Packed f32 bytes per compressed byte. ~6 kbps speech expands to 48 kHz stereo +// f32 by ~512×; 1024× leaves a 2× margin for container overhead and resampler delay. +const MAX_F32_BYTES_PER_COMPRESSED_BYTE: u64 = 1024; +const MAX_DECODED_SAMPLES: usize = 8 * 3_600 * AudioData::SAMPLE_RATE as usize * 2; + +fn decoded_samples_from_compressed_bytes(compressed_bytes: u64) -> usize { + let decoded_bytes = compressed_bytes.saturating_mul(MAX_F32_BYTES_PER_COMPRESSED_BYTE); + usize::try_from(decoded_bytes / std::mem::size_of::() as u64).unwrap_or(usize::MAX) +} + +fn decoded_sample_capacity( + duration: i64, + stream_time_base: f64, + target_channels: u16, + range: Option<&Range>, + compressed_bytes: u64, +) -> usize { + let channels = usize::from(target_channels); + let from_file = + decoded_samples_from_compressed_bytes(compressed_bytes).min(MAX_DECODED_SAMPLES); + if duration <= 0 || !stream_time_base.is_finite() || stream_time_base <= 0.0 { + return 0; + } + let frames = duration as f64 * stream_time_base * f64::from(AudioData::SAMPLE_RATE); + if !frames.is_finite() || frames <= 0.0 { + return 0; + } + let from_duration = (frames.ceil() as usize) + .saturating_add(64) + .saturating_mul(channels) + .min(from_file); + let Some(window) = range else { + return from_duration; + }; + window + .end + .saturating_sub(window.start) + .saturating_mul(channels) + .min(from_duration) +} + pub(crate) struct ResamplerOutput { frame: FFAudio, capacity: usize, @@ -781,4 +831,106 @@ mod tests { assert_eq!(data.sample_count(), 0); assert!(data.covers_source_range(1_000, 1_000)); } + + #[test] + fn decoded_sample_capacity_uses_range_or_duration() { + assert_eq!( + decoded_sample_capacity(48_000, 1.0 / 48_000.0, 2, Some(&(1_000..1_250)), u64::MAX), + 500 + ); + assert_eq!( + decoded_sample_capacity(0, 1.0 / 48_000.0, 2, Some(&(1_000..1_250)), u64::MAX), + 0 + ); + assert_eq!( + decoded_sample_capacity(-1, 1.0 / 48_000.0, 2, None, u64::MAX), + 0 + ); + assert_eq!( + decoded_sample_capacity(48_000, 1.0 / 48_000.0, 2, None, u64::MAX), + (48_000 + 64) * 2 + ); + } + + #[test] + fn decoded_sample_capacity_caps_load_track_open_ended_end() { + let end = usize::MAX / 4; + assert_eq!( + decoded_sample_capacity(48_000, 1.0 / 48_000.0, 1, Some(&(0..end)), u64::MAX), + 48_000 + 64 + ); + assert_eq!( + decoded_sample_capacity(48_000, 1.0 / 48_000.0, 2, Some(&(1_000..end)), u64::MAX), + (48_000 + 64) * 2 + ); + assert_eq!( + decoded_sample_capacity(-1, 1.0 / 48_000.0, 2, Some(&(0..end)), u64::MAX), + 0 + ); + assert_eq!( + decoded_sample_capacity(48_000, f64::NAN, 2, Some(&(0..end)), u64::MAX), + 0 + ); + } + + #[test] + fn decoded_sample_capacity_clamps_forged_duration_to_file_size() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tiny.wav"); + write_pcm_wav(&path, AudioData::SAMPLE_RATE, 8, &[1]); + let file_len = std::fs::metadata(&path).unwrap().len(); + assert!(file_len < 1_024); + + let ten_hours = 10 * 3_600 * i64::from(AudioData::SAMPLE_RATE); + let open_ended = 0..(usize::MAX / 4); + let reserved = + decoded_sample_capacity(ten_hours, 1.0 / 48_000.0, 2, Some(&open_ended), file_len); + let max_from_file = + decoded_samples_from_compressed_bytes(file_len).min(MAX_DECODED_SAMPLES); + + assert_eq!(reserved, max_from_file); + assert!(reserved > 0); + assert!(reserved < 1_000_000); + + let unclamped_hours = (10 * 3_600 * AudioData::SAMPLE_RATE as usize + 64) * 2; + assert!(unclamped_hours > 3_000_000_000usize); + assert!(reserved < unclamped_hours / 1_000); + } + + #[test] + fn from_file_range_accepts_load_track_open_ended_end() { + let _ = ffmpeg::init(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("open_ended.wav"); + let frames = AudioData::SAMPLE_RATE as usize; + write_pcm_wav(&path, AudioData::SAMPLE_RATE, frames, &[2_000]); + + let data = AudioData::from_file_range(&path, 0, usize::MAX / 4).unwrap(); + assert_eq!(data.channels(), 1); + assert!(data.sample_count().abs_diff(frames) <= 64); + assert!(data.covers_source_range(0, usize::MAX / 4)); + } + + #[test] + #[ignore] + fn benchmark_from_file_decode() { + let _ = ffmpeg::init(); + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("stereo_60s.wav"); + let frames = AudioData::SAMPLE_RATE as usize * 60; + write_pcm_wav(&path, AudioData::SAMPLE_RATE, frames, &[2_000, -2_000]); + let mut runs = Vec::new(); + for _ in 0..5 { + let start = std::time::Instant::now(); + let data = AudioData::from_file(&path).unwrap(); + runs.push(start.elapsed().as_secs_f64() * 1000.0); + assert_eq!(data.sample_count(), frames); + } + runs.sort_by(f64::total_cmp); + println!( + "from_file_60s_stereo_wav median_ms={:.1} min_ms={:.1}", + runs[runs.len() / 2], + runs[0] + ); + } }