SoulStream/crates/player/src/loudness.rs

//! Integrated-loudness measurement (ITU-R BS.1770-4 / EBU R128), the basis
//! of the player's volume normalization: every track is measured in LUFS
//! and played back with a gain that brings it to one target loudness, so a
//! quiet 90s master and a brickwalled modern one come out level.
//!
//! Dependency-free implementation of the standard: per-channel K-weighting
//! (a high-shelf that models head diffraction, then a high-pass), energy
//! collected in 400 ms blocks at 75 % overlap, and two-stage gating
//! (absolute −70 LUFS, then relative −10 LU below the ungated mean). The
//! meter is INCREMENTAL — `push` accepts interleaved frames as they are
//! decoded, and `integrated_lufs()` can be read at any point for the
//! loudness of everything seen so far — which is what lets a Soulseek
//! stream be normalized during its very first play instead of after it.

/// A biquad in direct form 1, processing one channel.
#[derive(Clone, Copy, Default)]
struct Biquad {
    b0: f64,
    b1: f64,
    b2: f64,
    a1: f64,
    a2: f64,
    x1: f64,
    x2: f64,
    y1: f64,
    y2: f64,
}

impl Biquad {
    #[inline]
    fn process(&mut self, x: f64) -> f64 {
        let y = self.b0 * x + self.b1 * self.x1 + self.b2 * self.x2
            - self.a1 * self.y1
            - self.a2 * self.y2;
        self.x2 = self.x1;
        self.x1 = x;
        self.y2 = self.y1;
        self.y1 = y;
        y
    }
}

/// Stage-1 K-weighting shelf (BS.1770 spherical-head model), designed for
/// an arbitrary sample rate with the reference parameters the standard's
/// 48 kHz table was derived from (the libebur128 derivation).
fn shelf(fs: f64) -> Biquad {
    let f0 = 1681.974450955533;
    let gain_db = 3.999843853973347;
    let q = 0.7071752369554196;

    let k = (std::f64::consts::PI * f0 / fs).tan();
    let vh = 10f64.powf(gain_db / 20.0);
    let vb = vh.powf(0.4996667741545416);
    let a0 = 1.0 + k / q + k * k;
    Biquad {
        b0: (vh + vb * k / q + k * k) / a0,
        b1: 2.0 * (k * k - vh) / a0,
        b2: (vh - vb * k / q + k * k) / a0,
        a1: 2.0 * (k * k - 1.0) / a0,
        a2: (1.0 - k / q + k * k) / a0,
        ..Default::default()
    }
}

/// Stage-2 K-weighting high-pass (rumble removal below ~38 Hz).
fn highpass(fs: f64) -> Biquad {
    let f0 = 38.13547087602444;
    let q = 0.5003270373238773;

    let k = (std::f64::consts::PI * f0 / fs).tan();
    let a0 = 1.0 + k / q + k * k;
    Biquad {
        b0: 1.0,
        b1: -2.0,
        b2: 1.0,
        a1: 2.0 * (k * k - 1.0) / a0,
        a2: (1.0 - k / q + k * k) / a0,
        ..Default::default()
    }
}

const MAX_CHANNELS: usize = 8;
/// Gating block = 400 ms as 4 overlapping 100 ms sub-blocks (75 % overlap).
const SUBBLOCKS_PER_BLOCK: usize = 4;
const ABSOLUTE_GATE_LUFS: f64 = -70.0;

pub struct LoudnessMeter {
    sample_rate: u32,
    channels: usize,
    shelves: [Biquad; MAX_CHANNELS],
    highpasses: [Biquad; MAX_CHANNELS],
    /// Samples per 100 ms sub-block (per channel).
    subblock_len: usize,
    /// Running energy sum of the current sub-block (all channels, weighted).
    cur_energy: f64,
    cur_count: usize,
    /// Completed 100 ms sub-block mean energies.
    subblocks: Vec<f64>,
    /// Pre-filter sample peak, for clipping headroom when boosting.
    peak: f32,
}

impl LoudnessMeter {
    pub fn new(sample_rate: u32, channels: usize) -> Self {
        let fs = sample_rate.max(1) as f64;
        let ch = channels.clamp(1, MAX_CHANNELS);
        LoudnessMeter {
            sample_rate: sample_rate.max(1),
            channels: ch,
            shelves: [shelf(fs); MAX_CHANNELS],
            highpasses: [highpass(fs); MAX_CHANNELS],
            subblock_len: (fs / 10.0).round().max(1.0) as usize,
            cur_energy: 0.0,
            cur_count: 0,
            subblocks: Vec::new(),
            peak: 0.0,
        }
    }

    /// Re-arm the filters for a new spec mid-stream (rare); accumulated
    /// blocks are kept — they are already energies, spec-independent.
    pub fn rearm(&mut self, sample_rate: u32, channels: usize) {
        if sample_rate == self.sample_rate && channels.clamp(1, MAX_CHANNELS) == self.channels {
            return;
        }
        let keep = std::mem::take(&mut self.subblocks);
        *self = LoudnessMeter::new(sample_rate, channels);
        self.subblocks = keep;
    }

    /// Feed interleaved f32 frames.
    pub fn push(&mut self, interleaved: &[f32]) {
        let ch = self.channels;
        for frame in interleaved.chunks_exact(ch) {
            let mut frame_energy = 0.0f64;
            for (c, &s) in frame.iter().enumerate() {
                self.peak = self.peak.max(s.abs());
                let w = self.shelves[c].process(s as f64);
                let k = self.highpasses[c].process(w);
                // Channel weights: 1.0 (stereo/mono; surround rears would be
                // 1.41 — practically all content here is stereo).
                frame_energy += k * k;
            }
            self.cur_energy += frame_energy;
            self.cur_count += 1;
            if self.cur_count >= self.subblock_len {
                self.subblocks.push(self.cur_energy / self.cur_count as f64);
                self.cur_energy = 0.0;
                self.cur_count = 0;
            }
        }
    }

    /// Seconds of audio consumed so far.
    pub fn seconds_seen(&self) -> f64 {
        (self.subblocks.len() as f64) * 0.1 + self.cur_count as f64 / self.sample_rate as f64
    }

    pub fn sample_peak(&self) -> f32 {
        self.peak
    }

    /// Gated integrated loudness of everything pushed so far, or None while
    /// there is not yet one full 400 ms block above the absolute gate.
    pub fn integrated_lufs(&self) -> Option<f64> {
        if self.subblocks.len() < SUBBLOCKS_PER_BLOCK {
            return None;
        }
        let block_energy = |i: usize| -> f64 {
            self.subblocks[i..i + SUBBLOCKS_PER_BLOCK]
                .iter()
                .sum::<f64>()
                / SUBBLOCKS_PER_BLOCK as f64
        };
        let lufs = |e: f64| -0.691 + 10.0 * e.max(f64::MIN_POSITIVE).log10();
        let n_blocks = self.subblocks.len() - SUBBLOCKS_PER_BLOCK + 1;

        // Absolute gate.
        let mut sum = 0.0;
        let mut count = 0usize;
        for i in 0..n_blocks {
            let e = block_energy(i);
            if lufs(e) > ABSOLUTE_GATE_LUFS {
                sum += e;
                count += 1;
            }
        }
        if count == 0 {
            return None;
        }
        // Relative gate: −10 LU below the absolute-gated mean.
        let relative_gate = lufs(sum / count as f64) - 10.0;
        let mut sum2 = 0.0;
        let mut count2 = 0usize;
        for i in 0..n_blocks {
            let e = block_energy(i);
            let l = lufs(e);
            if l > ABSOLUTE_GATE_LUFS && l > relative_gate {
                sum2 += e;
                count2 += 1;
            }
        }
        if count2 == 0 {
            return None;
        }
        Some(lufs(sum2 / count2 as f64))
    }
}

// ------------------------------------------------------------ analyzer ----

use std::path::PathBuf;
use std::sync::Arc;

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

use crate::source::{GrowingFile, StreamProgress};

/// What the loudness analyzer reads. `Growing` is the streaming case: the
/// same blocking growing-file reader playback uses, so analysis proceeds at
/// download speed — normally far ahead of the play cursor — and the FIRST
/// play of a stream converges on its true loudness within seconds.
pub enum AnalysisSource {
    Local(PathBuf),
    Growing {
        path: PathBuf,
        total: u64,
        progress: Arc<StreamProgress>,
    },
}

/// How often (in analyzed audio time) partial measurements are published.
/// 1 s keeps the play-gate responsive — the caller decides when it has
/// heard enough to commit a pre-playback gain.
const PUBLISH_EVERY_SECS: f64 = 1.0;

/// Decode `source` from the top and measure it. `on_update` receives
/// (partial LUFS, peak, seconds analyzed) roughly every 5 s of audio;
/// returns the final (integrated LUFS, peak) at EOF, or None on decode
/// failure / cancellation / too little audio. Blocking — run on a worker
/// thread. An aborted stream unblocks the reader and lands here as EOF.
pub fn analyze_loudness(
    source: AnalysisSource,
    cancelled: &dyn Fn() -> bool,
    on_update: &mut dyn FnMut(f64, f32, f64),
) -> Option<(f64, f32)> {
    let (mss, path) = match source {
        AnalysisSource::Local(path) => {
            let file = std::fs::File::open(&path).ok()?;
            (
                MediaSourceStream::new(Box::new(file), Default::default()),
                path,
            )
        }
        AnalysisSource::Growing {
            path,
            total,
            progress,
        } => {
            let growing = GrowingFile::open(&path, total, progress).ok()?;
            (
                MediaSourceStream::new(Box::new(ReadOnlySource::new(growing)), Default::default()),
                path,
            )
        }
    };

    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(std::path::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(),
        )
        .ok()?;
    let mut format = probed.format;
    let track = format.default_track()?;
    let track_id = track.id;
    let mut decoder = symphonia::default::get_codecs()
        .make(&track.codec_params, &DecoderOptions::default())
        .ok()?;

    let mut sample_buf: Option<SampleBuffer<f32>> = None;
    let mut meter: Option<LoudnessMeter> = None;
    let mut last_published = 0.0f64;

    loop {
        if cancelled() {
            return None;
        }
        let packet = match format.next_packet() {
            Ok(p) => p,
            Err(_) => break, // EOF (or aborted stream) → finish 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);
        match &mut meter {
            None => meter = Some(LoudnessMeter::new(spec.rate, channels)),
            Some(m) => m.rearm(spec.rate, channels),
        }
        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);
        let m = meter.as_mut().unwrap();
        m.push(sb.samples());

        let secs = m.seconds_seen();
        if secs - last_published >= PUBLISH_EVERY_SECS {
            if let Some(lufs) = m.integrated_lufs() {
                on_update(lufs, m.sample_peak(), secs);
            }
            last_published = secs;
        }
    }

    let m = meter?;
    m.integrated_lufs().map(|l| (l, m.sample_peak()))
}

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

    fn sine(freq: f64, fs: u32, secs: f64, amplitude: f64, channels: usize) -> Vec<f32> {
        let frames = (fs as f64 * secs) as usize;
        let mut out = Vec::with_capacity(frames * channels);
        for n in 0..frames {
            let s = (amplitude
                * (2.0 * std::f64::consts::PI * freq * n as f64 / fs as f64).sin())
                as f32;
            for _ in 0..channels {
                out.push(s);
            }
        }
        out
    }

    /// BS.1770's own reference point: a full-scale 997 Hz sine reads
    /// −3.01 LUFS (per channel of a stereo pair it sums +3 dB, so use mono).
    #[test]
    fn full_scale_sine_reads_reference_level() {
        let mut m = LoudnessMeter::new(48_000, 1);
        m.push(&sine(997.0, 48_000, 5.0, 1.0, 1));
        let l = m.integrated_lufs().expect("enough audio");
        assert!((l - (-3.01)).abs() < 0.1, "got {l}");
    }

    /// −20 dB quieter input must read 20 LU lower — the meter is linear in
    /// level, which is what makes gain = target − measured valid.
    #[test]
    fn level_linearity() {
        let mut loud = LoudnessMeter::new(44_100, 2);
        loud.push(&sine(997.0, 44_100, 4.0, 0.5, 2));
        let mut quiet = LoudnessMeter::new(44_100, 2);
        quiet.push(&sine(997.0, 44_100, 4.0, 0.05, 2));
        let dl = loud.integrated_lufs().unwrap();
        let dq = quiet.integrated_lufs().unwrap();
        assert!(((dl - dq) - 20.0).abs() < 0.2, "loud {dl} quiet {dq}");
    }

    /// Silence padding must not drag the integrated value down — that is
    /// the absolute gate doing its job.
    #[test]
    fn silence_is_gated_out() {
        let mut padded = LoudnessMeter::new(48_000, 1);
        padded.push(&sine(997.0, 48_000, 3.0, 0.25, 1));
        padded.push(&vec![0.0f32; 48_000 * 3]);
        let mut bare = LoudnessMeter::new(48_000, 1);
        bare.push(&sine(997.0, 48_000, 3.0, 0.25, 1));
        let lp = padded.integrated_lufs().unwrap();
        let lb = bare.integrated_lufs().unwrap();
        assert!((lp - lb).abs() < 0.3, "padded {lp} vs bare {lb}");
    }

    #[test]
    fn peak_tracks_prefilter_samples() {
        let mut m = LoudnessMeter::new(48_000, 2);
        m.push(&sine(440.0, 48_000, 1.0, 0.8, 2));
        assert!((m.sample_peak() - 0.8).abs() < 0.01);
    }

    #[test]
    fn too_little_audio_reads_none() {
        let mut m = LoudnessMeter::new(48_000, 2);
        m.push(&sine(997.0, 48_000, 0.2, 0.5, 2));
        assert!(m.integrated_lufs().is_none());
    }
}