peerspeak/src/audio/eq.rs

//! Per-peer listener-side voice EQ.
//!
//! The EQ is deliberately small and local: three RBJ cookbook biquads at fixed
//! voice-oriented frequencies, with only gain exposed to the UI. State lives per
//! peer in the playout mixer so filter delay registers are continuous across 20ms
//! Opus frames; flat settings are treated as bypass so the default path is cheap
//! and sample-exact.

use serde::{Deserialize, Serialize};

const DEFAULT_SAMPLE_RATE: f32 = 48_000.0;
const LOW_SHELF_HZ: f32 = 160.0;
const MID_PEAK_HZ: f32 = 2_400.0;
const HIGH_SHELF_HZ: f32 = 6_500.0;
const MID_Q: f32 = 1.0;
const SHELF_Q: f32 = std::f32::consts::FRAC_1_SQRT_2;
const FLAT_EPSILON_DB: f32 = 0.001;

/// UI and config clamp for each band. Wide enough to be useful for voice, narrow
/// enough that a peer cannot accidentally make the listener-side limiter do all
/// the work.
pub const EQ_GAIN_DB_MIN: f32 = -12.0;
pub const EQ_GAIN_DB_MAX: f32 = 12.0;

/// Persisted per-peer EQ gains, in decibels. `Default` is flat/bypassed.
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq)]
pub struct EqSettings {
    #[serde(default)]
    pub low_gain_db: f32,
    #[serde(default)]
    pub mid_gain_db: f32,
    #[serde(default)]
    pub high_gain_db: f32,
}

impl Default for EqSettings {
    fn default() -> Self {
        Self {
            low_gain_db: 0.0,
            mid_gain_db: 0.0,
            high_gain_db: 0.0,
        }
    }
}

impl EqSettings {
    pub fn flat() -> Self {
        Self::default()
    }

    /// Clamp all public gains to the supported UI/DSP range.
    pub fn clamped(self) -> Self {
        Self {
            low_gain_db: self.low_gain_db.clamp(EQ_GAIN_DB_MIN, EQ_GAIN_DB_MAX),
            mid_gain_db: self.mid_gain_db.clamp(EQ_GAIN_DB_MIN, EQ_GAIN_DB_MAX),
            high_gain_db: self.high_gain_db.clamp(EQ_GAIN_DB_MIN, EQ_GAIN_DB_MAX),
        }
    }

    /// True when the EQ should be bypassed entirely.
    pub fn is_flat(self) -> bool {
        self.low_gain_db.abs() <= FLAT_EPSILON_DB
            && self.mid_gain_db.abs() <= FLAT_EPSILON_DB
            && self.high_gain_db.abs() <= FLAT_EPSILON_DB
    }
}

/// A stateful three-band EQ. One instance belongs to one decoded peer stream.
pub struct Eq {
    settings: EqSettings,
    low: Biquad,
    mid: Biquad,
    high: Biquad,
}

impl Eq {
    /// Build an EQ at the application's audio rate (48 kHz).
    pub fn new(settings: EqSettings) -> Self {
        Self::with_sample_rate(settings, DEFAULT_SAMPLE_RATE)
    }

    fn with_sample_rate(settings: EqSettings, sample_rate: f32) -> Self {
        let settings = settings.clamped();
        Self {
            settings,
            low: Biquad::low_shelf(sample_rate, LOW_SHELF_HZ, settings.low_gain_db, SHELF_Q),
            mid: Biquad::peaking(sample_rate, MID_PEAK_HZ, settings.mid_gain_db, MID_Q),
            high: Biquad::high_shelf(sample_rate, HIGH_SHELF_HZ, settings.high_gain_db, SHELF_Q),
        }
    }

    pub fn settings(&self) -> EqSettings {
        self.settings
    }

    /// Process one mono PCM frame in place. Flat settings are sample-exact bypass.
    pub fn process_frame(&mut self, frame: &mut [i16]) {
        if self.settings.is_flat() {
            return;
        }
        for sample in frame {
            let x = *sample as f32;
            let y = self.high.process(self.mid.process(self.low.process(x)));
            *sample = y.round().clamp(i16::MIN as f32, i16::MAX as f32) as i16;
        }
    }
}

#[derive(Debug, Clone, Copy)]
struct Coeffs {
    b0: f32,
    b1: f32,
    b2: f32,
    a1: f32,
    a2: f32,
}

impl Coeffs {
    fn normalized(b0: f32, b1: f32, b2: f32, a0: f32, a1: f32, a2: f32) -> Self {
        let inv_a0 = 1.0 / a0;
        Self {
            b0: b0 * inv_a0,
            b1: b1 * inv_a0,
            b2: b2 * inv_a0,
            a1: a1 * inv_a0,
            a2: a2 * inv_a0,
        }
    }

    fn all_finite(self) -> bool {
        self.b0.is_finite()
            && self.b1.is_finite()
            && self.b2.is_finite()
            && self.a1.is_finite()
            && self.a2.is_finite()
    }
}

/// Direct Form II transposed biquad. The two delay registers are the state that
/// must survive across frames.
struct Biquad {
    coeffs: Coeffs,
    z1: f32,
    z2: f32,
}

impl Biquad {
    fn new(coeffs: Coeffs) -> Self {
        debug_assert!(coeffs.all_finite());
        Self {
            coeffs,
            z1: 0.0,
            z2: 0.0,
        }
    }

    fn low_shelf(sample_rate: f32, freq: f32, gain_db: f32, q: f32) -> Self {
        let (a, cos_w0, alpha) = rbj_terms(sample_rate, freq, gain_db, q);
        let sqrt_a = a.sqrt();
        let b0 = a * ((a + 1.0) - (a - 1.0) * cos_w0 + 2.0 * sqrt_a * alpha);
        let b1 = 2.0 * a * ((a - 1.0) - (a + 1.0) * cos_w0);
        let b2 = a * ((a + 1.0) - (a - 1.0) * cos_w0 - 2.0 * sqrt_a * alpha);
        let a0 = (a + 1.0) + (a - 1.0) * cos_w0 + 2.0 * sqrt_a * alpha;
        let a1 = -2.0 * ((a - 1.0) + (a + 1.0) * cos_w0);
        let a2 = (a + 1.0) + (a - 1.0) * cos_w0 - 2.0 * sqrt_a * alpha;
        Self::new(Coeffs::normalized(b0, b1, b2, a0, a1, a2))
    }

    fn peaking(sample_rate: f32, freq: f32, gain_db: f32, q: f32) -> Self {
        let (a, cos_w0, alpha) = rbj_terms(sample_rate, freq, gain_db, q);
        let b0 = 1.0 + alpha * a;
        let b1 = -2.0 * cos_w0;
        let b2 = 1.0 - alpha * a;
        let a0 = 1.0 + alpha / a;
        let a1 = -2.0 * cos_w0;
        let a2 = 1.0 - alpha / a;
        Self::new(Coeffs::normalized(b0, b1, b2, a0, a1, a2))
    }

    fn high_shelf(sample_rate: f32, freq: f32, gain_db: f32, q: f32) -> Self {
        let (a, cos_w0, alpha) = rbj_terms(sample_rate, freq, gain_db, q);
        let sqrt_a = a.sqrt();
        let b0 = a * ((a + 1.0) + (a - 1.0) * cos_w0 + 2.0 * sqrt_a * alpha);
        let b1 = -2.0 * a * ((a - 1.0) + (a + 1.0) * cos_w0);
        let b2 = a * ((a + 1.0) + (a - 1.0) * cos_w0 - 2.0 * sqrt_a * alpha);
        let a0 = (a + 1.0) - (a - 1.0) * cos_w0 + 2.0 * sqrt_a * alpha;
        let a1 = 2.0 * ((a - 1.0) - (a + 1.0) * cos_w0);
        let a2 = (a + 1.0) - (a - 1.0) * cos_w0 - 2.0 * sqrt_a * alpha;
        Self::new(Coeffs::normalized(b0, b1, b2, a0, a1, a2))
    }

    fn process(&mut self, x: f32) -> f32 {
        let y = self.coeffs.b0 * x + self.z1;
        self.z1 = self.coeffs.b1 * x - self.coeffs.a1 * y + self.z2;
        self.z2 = self.coeffs.b2 * x - self.coeffs.a2 * y;

        // Avoid carrying denormal-sized state forever on long quiet tails.
        if self.z1.abs() < 1.0e-20 {
            self.z1 = 0.0;
        }
        if self.z2.abs() < 1.0e-20 {
            self.z2 = 0.0;
        }
        y
    }
}

fn rbj_terms(sample_rate: f32, freq: f32, gain_db: f32, q: f32) -> (f32, f32, f32) {
    let sr = sample_rate.max(1.0);
    let f = freq.clamp(1.0, sr * 0.49);
    let w0 = 2.0 * std::f32::consts::PI * f / sr;
    let a = 10.0f32.powf(gain_db / 40.0);
    let alpha = w0.sin() / (2.0 * q.max(0.001));
    (a, w0.cos(), alpha)
}

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

    fn sine(freq: f32, len: usize, amp: f32) -> Vec<i16> {
        (0..len)
            .map(|n| {
                let t = n as f32 / DEFAULT_SAMPLE_RATE;
                (amp * (2.0 * std::f32::consts::PI * freq * t).sin()).round() as i16
            })
            .collect()
    }

    fn rms(frame: &[i16]) -> f32 {
        let sum: f32 = frame.iter().map(|&s| (s as f32).powi(2)).sum();
        (sum / frame.len().max(1) as f32).sqrt()
    }

    #[test]
    fn flat_eq_is_sample_exact_identity() {
        let mut eq = Eq::new(EqSettings::flat());
        let mut frame: Vec<i16> = (-480..480).map(|n| (n * 31) as i16).collect();
        let original = frame.clone();
        eq.process_frame(&mut frame);
        assert_eq!(frame, original);
    }

    #[test]
    fn low_shelf_boost_raises_low_frequency_energy() {
        let mut eq = Eq::new(EqSettings {
            low_gain_db: 9.0,
            ..EqSettings::flat()
        });
        let mut low = sine(100.0, 48_000, 3_000.0);
        let before = rms(&low);
        eq.process_frame(&mut low);
        let after = rms(&low);
        assert!(
            after > before * 1.6,
            "low shelf should boost low RMS: {before} -> {after}"
        );
    }

    #[test]
    fn high_shelf_boost_raises_high_frequency_energy() {
        let mut eq = Eq::new(EqSettings {
            high_gain_db: 9.0,
            ..EqSettings::flat()
        });
        let mut high = sine(8_000.0, 48_000, 3_000.0);
        let before = rms(&high);
        eq.process_frame(&mut high);
        let after = rms(&high);
        assert!(
            after > before * 1.6,
            "high shelf should boost high RMS: {before} -> {after}"
        );
    }

    #[test]
    fn coefficients_are_finite_across_supported_gain_range() {
        for gain in [EQ_GAIN_DB_MIN, -6.0, 0.0, 6.0, EQ_GAIN_DB_MAX] {
            for b in [
                Biquad::low_shelf(DEFAULT_SAMPLE_RATE, LOW_SHELF_HZ, gain, SHELF_Q),
                Biquad::peaking(DEFAULT_SAMPLE_RATE, MID_PEAK_HZ, gain, MID_Q),
                Biquad::high_shelf(DEFAULT_SAMPLE_RATE, HIGH_SHELF_HZ, gain, SHELF_Q),
            ] {
                assert!(
                    b.coeffs.all_finite(),
                    "coefficients must be finite at {gain} dB"
                );
            }
        }
    }

    #[test]
    fn hot_signal_does_not_nan_or_wrap() {
        let mut eq = Eq::new(EqSettings {
            low_gain_db: 12.0,
            mid_gain_db: 12.0,
            high_gain_db: 12.0,
        });
        let mut frame = sine(1_000.0, 48_000, 30_000.0);
        eq.process_frame(&mut frame);
        let peak = frame.iter().map(|&s| i32::from(s).abs()).max().unwrap_or(0);
        assert!(
            peak > 1_000,
            "processed signal should retain audible energy"
        );
        assert!(
            frame.iter().any(|&s| s > 0) && frame.iter().any(|&s| s < 0),
            "a boosted sine should retain both polarities"
        );
    }

    #[test]
    fn settings_are_clamped() {
        let s = EqSettings {
            low_gain_db: -99.0,
            mid_gain_db: 2.0,
            high_gain_db: 99.0,
        }
        .clamped();
        assert_eq!(s.low_gain_db, EQ_GAIN_DB_MIN);
        assert_eq!(s.mid_gain_db, 2.0);
        assert_eq!(s.high_gain_db, EQ_GAIN_DB_MAX);
    }
}