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, Terminal};

/// 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;

/// How an analysis ended.
///
/// The distinction is the whole point: integrated loudness is a property of
/// a COMPLETE programme, and the numbers a half-decoded track produces look
/// exactly like real ones. Only [`Analysis::Complete`] may be recorded as a
/// track's measurement.
pub enum Analysis {
    /// Decoded from the first sample to the natural end of the media.
    Complete { lufs: f64, peak: f32 },
    /// Ended before the media did — whatever was measured describes a
    /// prefix, not the track.
    Incomplete(Incomplete),
}

/// Why an analysis stopped short.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Incomplete {
    /// A newer track superseded this one.
    Cancelled,
    /// The stream was stopped or switched away from.
    Aborted,
    /// The transfer died mid-file.
    TransferFailed,
    /// The bytes stopped arriving without the transfer reporting either
    /// completion or failure.
    Truncated,
    /// Unreadable container, no supported track, or a decoder that gave up.
    DecodeFailed,
    /// Never enough audio above the gate to produce a measurement.
    TooShort,
}

/// Why the decode loop stopped, before the transfer gets a say.
enum Exit {
    /// Ran out of packets. For a growing file this may be a real end of
    /// media OR an abort wearing its clothes — the transfer decides.
    Ended,
    Failed(Incomplete),
}

/// Decode `source` from the top and measure it. `on_update` receives
/// (partial LUFS, peak, seconds analyzed) roughly every second of audio, for
/// a caller that wants to act before the end; the RETURN value is the only
/// thing that says whether the measurement covers the whole track.
///
/// Blocking — run on a worker thread.
pub fn analyze_loudness(
    source: AnalysisSource,
    cancelled: &dyn Fn() -> bool,
    on_update: &mut dyn FnMut(f64, f32, f64),
) -> Analysis {
    use Analysis::Incomplete as Inc;

    // Kept so a growing stream can be asked, once the packets stop, whether
    // that was the end of the file or the end of the transfer.
    let mut transfer: Option<Arc<StreamProgress>> = None;
    let (mss, path) = match source {
        AnalysisSource::Local(path) => {
            let Ok(file) = std::fs::File::open(&path) else {
                return Inc(Incomplete::DecodeFailed);
            };
            (
                MediaSourceStream::new(Box::new(file), Default::default()),
                path,
            )
        }
        AnalysisSource::Growing {
            path,
            total,
            progress,
        } => {
            transfer = Some(progress.clone());
            let Ok(growing) = GrowingFile::open(&path, total, progress) else {
                return Inc(Incomplete::DecodeFailed);
            };
            (
                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 Ok(probed) = symphonia::default::get_probe().format(
        &hint,
        mss,
        &FormatOptions::default(),
        &MetadataOptions::default(),
    ) else {
        return Inc(Incomplete::DecodeFailed);
    };
    let mut format = probed.format;
    let Some(track) = format.default_track() else {
        return Inc(Incomplete::DecodeFailed);
    };
    let track_id = track.id;
    let Ok(mut decoder) = symphonia::default::get_codecs()
        .make(&track.codec_params, &DecoderOptions::default())
    else {
        return Inc(Incomplete::DecodeFailed);
    };

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

    let exit = loop {
        if cancelled() {
            return Inc(Incomplete::Cancelled);
        }
        let packet = match format.next_packet() {
            Ok(p) => p,
            // Out of packets. Symphonia signals a clean end of media as an
            // unexpected-EOF io error; anything else is the file or the
            // demuxer failing, which is not an end of programme.
            Err(symphonia::core::errors::Error::IoError(e))
                if e.kind() == std::io::ErrorKind::UnexpectedEof =>
            {
                break Exit::Ended
            }
            Err(_) => break Exit::Failed(Incomplete::DecodeFailed),
        };
        if packet.track_id() != track_id {
            continue;
        }
        let decoded = match decoder.decode(&packet) {
            Ok(d) => d,
            // A corrupt packet is recoverable; the standard says skip it.
            Err(symphonia::core::errors::Error::DecodeError(_)) => continue,
            Err(_) => break Exit::Failed(Incomplete::DecodeFailed),
        };
        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;
        }
    };

    // Running out of packets only means the whole track was seen if the
    // BYTES also ran out for the right reason. An aborted growing file
    // returns a graceful EOF, so without this the head of a stream the user
    // skipped away from would be indistinguishable from a finished track.
    match (exit, transfer.as_deref().map(StreamProgress::terminal)) {
        // A transfer that stopped explains the exit whichever way the
        // decoder happened to experience it — an abort arrives as a clean
        // end of file, a failure as an unreadable one.
        (_, Some(Terminal::Aborted)) => Inc(Incomplete::Aborted),
        (_, Some(Terminal::Failed)) => Inc(Incomplete::TransferFailed),
        (Exit::Failed(why), _) => Inc(why),
        // Out of packets while bytes were still expected: the reader gave up
        // early, so this is a prefix even though nothing reported an error.
        (Exit::Ended, Some(Terminal::Running)) => Inc(Incomplete::Truncated),
        (Exit::Ended, None | Some(Terminal::Done)) => {
            match meter.as_ref().and_then(LoudnessMeter::integrated_lufs) {
                Some(lufs) => Analysis::Complete {
                    lufs,
                    peak: meter.expect("measured").sample_peak(),
                },
                None => Inc(Incomplete::TooShort),
            }
        }
    }
}

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

    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());
    }

    // ------------------------------------------------- completion ----

    /// A 16-bit PCM WAV of `secs` seconds, so the analyzer has a real
    /// container to decode rather than a synthetic packet stream.
    fn wav(secs: f64, amplitude: f64) -> Vec<u8> {
        let (fs, ch) = (48_000u32, 1usize);
        let pcm = sine(997.0, fs, secs, amplitude, ch);
        let data: Vec<u8> = pcm
            .iter()
            .flat_map(|s| ((s * 32767.0) as i16).to_le_bytes())
            .collect();
        let mut out = Vec::with_capacity(44 + data.len());
        out.extend(b"RIFF");
        out.extend(((36 + data.len()) as u32).to_le_bytes());
        out.extend(b"WAVEfmt ");
        out.extend(16u32.to_le_bytes());
        out.extend(1u16.to_le_bytes()); // PCM
        out.extend((ch as u16).to_le_bytes());
        out.extend(fs.to_le_bytes());
        out.extend((fs * ch as u32 * 2).to_le_bytes()); // byte rate
        out.extend(((ch * 2) as u16).to_le_bytes()); // block align
        out.extend(16u16.to_le_bytes()); // bits
        out.extend(b"data");
        out.extend((data.len() as u32).to_le_bytes());
        out.extend(data);
        out
    }

    fn fixture(name: &str, bytes: &[u8]) -> PathBuf {
        let dir = std::env::temp_dir().join("ssplus-player-test");
        std::fs::create_dir_all(&dir).unwrap();
        let path = dir.join(name);
        std::fs::write(&path, bytes).unwrap();
        path
    }

    /// A transfer that finished is a whole-track measurement.
    #[test]
    fn a_finished_stream_measures_complete() {
        let bytes = wav(5.0, 0.5);
        let path = fixture("complete.wav", &bytes);
        let progress = StreamProgress::new();
        progress.update(bytes.len() as u64, true, None);
        let got = analyze_loudness(
            AnalysisSource::Growing {
                path,
                total: bytes.len() as u64,
                progress,
            },
            &|| false,
            &mut |_, _, _| {},
        );
        assert!(matches!(got, Analysis::Complete { .. }), "expected Complete");
    }

    /// The cache-poisoning path: a stream stopped during its opening reads
    /// as a graceful EOF, and must NOT come back as a track measurement.
    /// Before this, skipping a song mid-download recorded whatever its intro
    /// measured under the full declared file size, and every later play
    /// boosted the whole track off it.
    #[test]
    fn an_aborted_stream_is_not_a_measurement() {
        let bytes = wav(30.0, 0.5);
        // Only the first fifth has arrived — plenty for a gated reading.
        let head = bytes.len() / 5;
        let path = fixture("aborted.wav", &bytes);
        let progress = StreamProgress::new();
        progress.update(head as u64, false, None);

        let killer = {
            let progress = progress.clone();
            std::thread::spawn(move || {
                std::thread::sleep(Duration::from_millis(300));
                progress.abort();
            })
        };
        let mut measured_a_prefix = false;
        let got = analyze_loudness(
            AnalysisSource::Growing {
                path,
                total: bytes.len() as u64,
                progress,
            },
            &|| false,
            &mut |_, _, _| measured_a_prefix = true,
        );
        killer.join().unwrap();

        // The point is precisely that there WAS a usable-looking number.
        assert!(measured_a_prefix, "fixture never produced a partial reading");
        assert!(
            matches!(got, Analysis::Incomplete(Incomplete::Aborted)),
            "an aborted stream must not report a whole-track measurement"
        );
    }

    /// A transfer that died mid-file is likewise not a measurement.
    #[test]
    fn a_failed_transfer_is_not_a_measurement() {
        let bytes = wav(30.0, 0.5);
        let head = bytes.len() / 5;
        let path = fixture("failed.wav", &bytes);
        let progress = StreamProgress::new();
        progress.update(head as u64, false, None);
        let killer = {
            let progress = progress.clone();
            std::thread::spawn(move || {
                std::thread::sleep(Duration::from_millis(300));
                progress.update(head as u64, false, Some("peer went away".into()));
            })
        };
        let got = analyze_loudness(
            AnalysisSource::Growing {
                path,
                total: bytes.len() as u64,
                progress,
            },
            &|| false,
            &mut |_, _, _| {},
        );
        killer.join().unwrap();
        assert!(
            matches!(got, Analysis::Incomplete(Incomplete::TransferFailed)),
            "a dead transfer must not report a whole-track measurement"
        );
    }

    /// A complete local file is the ordinary case and must still measure.
    #[test]
    fn a_local_file_measures_complete() {
        let path = fixture("local.wav", &wav(5.0, 0.5));
        let got = analyze_loudness(AnalysisSource::Local(path), &|| false, &mut |_, _, _| {});
        let Analysis::Complete { lufs, .. } = got else {
            panic!("expected Complete");
        };
        // -6 dBFS sine, mono: ~-9 LUFS.
        assert!((lufs - (-9.0)).abs() < 0.5, "got {lufs}");
    }

    /// Cancellation stops the work and reports why.
    #[test]
    fn cancellation_reports_itself() {
        let path = fixture("cancel.wav", &wav(5.0, 0.5));
        let got = analyze_loudness(AnalysisSource::Local(path), &|| true, &mut |_, _, _| {});
        assert!(matches!(got, Analysis::Incomplete(Incomplete::Cancelled)));
    }
}