SoulStream/crates/player/src/waveform.rs

//! Waveform extraction: decode an audio file and downsample it to a small
//! array of normalized 0..1 loudness values, for a waveform scrubber.
//!
//! Uses RMS energy per window, not peak amplitude: modern masters are
//! limited so hard that nearly every window's *peak* touches the digital
//! ceiling, which renders as a clipped brick. RMS tracks perceived
//! loudness and keeps the verse/chorus dynamics visible.

use std::path::Path;

use symphonia::core::audio::SampleBuffer;
use symphonia::core::codecs::DecoderOptions;
use symphonia::core::formats::FormatOptions;
use symphonia::core::io::MediaSourceStream;
use symphonia::core::meta::MetadataOptions;
use symphonia::core::probe::Hint;

/// Peaks folded over this many frames before resampling to the final bucket
/// count. Keeps memory bounded regardless of track length.
const WINDOW: usize = 1024;

/// Decode `path` and return `buckets` normalized peak values (0.0..=1.0).
/// Returns an empty vec on any decode failure.
pub fn compute_waveform(path: &Path, buckets: usize) -> Vec<f32> {
    match try_compute(path, buckets) {
        Ok(v) => v,
        Err(e) => {
            log::debug!("waveform for {} failed: {e}", path.display());
            Vec::new()
        }
    }
}

fn try_compute(path: &Path, buckets: usize) -> Result<Vec<f32>, String> {
    let file = std::fs::File::open(path).map_err(|e| e.to_string())?;
    let mss = MediaSourceStream::new(Box::new(file), Default::default());

    let mut hint = Hint::new();
    if let Some(ext) = path.extension().and_then(|e| e.to_str()) {
        // Stream partials are "<name>.<ext>.partial" — hint the codec ext.
        let ext = if ext.eq_ignore_ascii_case("partial") {
            path.file_stem()
                .map(Path::new)
                .and_then(|p| p.extension())
                .and_then(|e| e.to_str())
                .unwrap_or(ext)
        } else {
            ext
        };
        hint.with_extension(ext);
    }

    let probed = symphonia::default::get_probe()
        .format(
            &hint,
            mss,
            &FormatOptions::default(),
            &MetadataOptions::default(),
        )
        .map_err(|e| e.to_string())?;
    let mut format = probed.format;

    let track = format
        .default_track()
        .ok_or_else(|| "no default track".to_string())?;
    let track_id = track.id;
    let mut decoder = symphonia::default::get_codecs()
        .make(&track.codec_params, &DecoderOptions::default())
        .map_err(|e| e.to_string())?;

    let mut sample_buf: Option<SampleBuffer<f32>> = None;
    let mut coarse: Vec<f32> = Vec::new();
    let mut win_energy = 0f64;
    let mut win_count = 0usize;

    loop {
        let packet = match format.next_packet() {
            Ok(p) => p,
            Err(_) => break, // EOF or error → stop with what we have
        };
        if packet.track_id() != track_id {
            continue;
        }
        let decoded = match decoder.decode(&packet) {
            Ok(d) => d,
            Err(symphonia::core::errors::Error::DecodeError(_)) => continue,
            Err(_) => break,
        };

        let spec = *decoded.spec();
        let channels = spec.channels.count().max(1);
        if sample_buf.is_none() {
            sample_buf = Some(SampleBuffer::new(decoded.capacity() as u64, spec));
        }
        let sb = sample_buf.as_mut().unwrap();
        sb.copy_interleaved_ref(decoded);

        for frame in sb.samples().chunks(channels) {
            let energy =
                frame.iter().map(|&s| (s as f64) * (s as f64)).sum::<f64>() / channels as f64;
            win_energy += energy;
            win_count += 1;
            if win_count >= WINDOW {
                coarse.push((win_energy / WINDOW as f64).sqrt() as f32);
                win_energy = 0.0;
                win_count = 0;
            }
        }
    }
    if win_count > 0 {
        coarse.push((win_energy / win_count as f64).sqrt() as f32);
    }

    Ok(resample(&coarse, buckets))
}

/// Downsample `coarse` RMS values to exactly `buckets`, averaging each
/// span, then normalize so the loudest bucket is 1.0.
fn resample(coarse: &[f32], buckets: usize) -> Vec<f32> {
    let mut out = vec![0f32; buckets];
    if coarse.is_empty() || buckets == 0 {
        return out;
    }
    for (b, slot) in out.iter_mut().enumerate() {
        let start = b * coarse.len() / buckets;
        let end = (((b + 1) * coarse.len() / buckets).max(start + 1)).min(coarse.len());
        let span = &coarse[start..end];
        *slot = span.iter().sum::<f32>() / span.len() as f32;
    }
    let max = out.iter().fold(0f32, |m, &v| m.max(v));
    if max > 0.0 {
        for v in out.iter_mut() {
            *v /= max;
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn resample_normalizes() {
        let coarse = vec![0.2, 0.4, 0.8, 0.1, 0.6, 0.3];
        let out = resample(&coarse, 3);
        assert_eq!(out.len(), 3);
        assert!((out.iter().cloned().fold(0f32, f32::max) - 1.0).abs() < 1e-6);
    }

    #[test]
    fn empty_input() {
        assert_eq!(resample(&[], 5), vec![0.0; 5]);
    }
}