peerspeak/src/audio/gate.rs

//! A smooth, stateful noise gate for the capture path.
//!
//! The original gate was a per-frame hard cut: compute the frame's RMS and, if it
//! fell below the threshold, drop the whole 20ms frame. That chops word onsets and
//! tails, chatters when the level sits right at the threshold, and offers no partial
//! attenuation. This replaces it with a proper gate envelope:
//!
//! - **Hysteresis** — the gate opens at `open_threshold` but only closes once the
//!   level falls below a lower `close_threshold` (a fixed ratio of the open one), so
//!   speech hovering near the threshold doesn't flap the gate on and off.
//! - **Attack / release** — when opening, the gain ramps 0→1 over a few ms; when
//!   closing, it ramps 1→0 over a longer window. The ramp is applied per sample, so
//!   the gate fades rather than clicking.
//! - **Hangover (hold)** — after the level drops, the gate stays fully open for a
//!   hold window before it begins to release, so brief inter-word pauses and quiet
//!   word tails survive instead of being clipped.
//!
//! A fully-closed frame (gain pinned at 0 with nothing to release) is reported as
//! "don't transmit" so the gate keeps the original bandwidth win of not sending
//! pure silence — the receiver's jitter buffer conceals the gap.

/// Gate timing/shape constants, in milliseconds. Tuned for voice.
const ATTACK_MS: f32 = 5.0;
const RELEASE_MS: f32 = 80.0;
const HOLD_MS: f32 = 200.0;
/// The close threshold as a fraction of the open threshold (hysteresis).
const CLOSE_RATIO: f32 = 0.6;
/// Below this open threshold the gate is considered disabled (pass-through).
const DISABLED_EPSILON: f32 = 0.0001;

/// A smooth noise gate. One instance lives in the capture thread and processes
/// each PCM frame in place, carrying its envelope state across frames.
pub struct NoiseGate {
    /// Per-sample gain increment while opening (1.0 / attack_samples).
    attack_step: f32,
    /// Per-sample gain decrement while releasing (1.0 / release_samples).
    release_step: f32,
    /// How long (in samples) to hold the gate open after the level drops.
    hold_samples: u32,
    /// Current envelope gain in `0.0..=1.0`, carried across frames.
    gain: f32,
    /// Whether the gate currently considers the signal "present".
    open: bool,
    /// Remaining hold (in samples) before an open gate begins to release.
    hold_counter: u32,
}

impl NoiseGate {
    /// Builds a gate for the given sample rate (Hz). Thresholds are passed per
    /// frame to [`process`](Self::process) so the live slider value applies
    /// immediately without rebuilding the gate.
    pub fn new(sample_rate: u32) -> Self {
        let sr = sample_rate as f32;
        let attack_samples = (ATTACK_MS / 1000.0 * sr).max(1.0);
        let release_samples = (RELEASE_MS / 1000.0 * sr).max(1.0);
        Self {
            attack_step: 1.0 / attack_samples,
            release_step: 1.0 / release_samples,
            hold_samples: (HOLD_MS / 1000.0 * sr) as u32,
            gain: 0.0,
            open: false,
            hold_counter: 0,
        }
    }

    /// Applies the gate to one PCM frame in place. `open_threshold` is the live
    /// slider value (normalized RMS, `0.0..`); pass `<= DISABLED_EPSILON` to
    /// disable gating (pass-through). Returns `true` if the frame should be
    /// transmitted, `false` only when the gate is fully closed (so the caller can
    /// skip sending pure silence).
    pub fn process(&mut self, pcm: &mut [i16], open_threshold: f32) -> bool {
        // Disabled: pass through untouched, and make sure the envelope is parked
        // open so re-enabling mid-stream doesn't start with a spurious fade-in.
        if open_threshold <= DISABLED_EPSILON {
            self.gain = 1.0;
            self.open = true;
            self.hold_counter = self.hold_samples;
            return true;
        }
        if pcm.is_empty() {
            return self.open || self.gain > 0.0;
        }

        let close_threshold = open_threshold * CLOSE_RATIO;
        let rms = frame_rms(pcm);

        // Update open/closed state with hysteresis + hold. Detection is per frame;
        // the gain ramp below is per sample.
        if rms >= open_threshold {
            self.open = true;
            self.hold_counter = self.hold_samples;
        } else if self.open {
            if rms >= close_threshold {
                // Still above the close threshold — refresh the hold window.
                self.hold_counter = self.hold_samples;
            } else {
                // Below close: spend the hold window, then begin releasing.
                self.hold_counter = self.hold_counter.saturating_sub(pcm.len() as u32);
                if self.hold_counter == 0 {
                    self.open = false;
                }
            }
        }

        let target = if self.open { 1.0 } else { 0.0 };

        // Per-sample gain ramp toward the target, applied to the frame.
        for sample in pcm.iter_mut() {
            if self.gain < target {
                self.gain = (self.gain + self.attack_step).min(target);
            } else if self.gain > target {
                self.gain = (self.gain - self.release_step).max(target);
            }
            *sample = (*sample as f32 * self.gain).round() as i16;
        }

        // Transmit unless the gate is fully closed with nothing left to release.
        self.open || self.gain > 0.0
    }
}

/// RMS of a PCM frame, normalized to `0.0..=1.0` (full-scale i16 == 1.0).
fn frame_rms(pcm: &[i16]) -> f32 {
    if pcm.is_empty() {
        return 0.0;
    }
    let mut sum_sq = 0.0f32;
    for &s in pcm {
        let n = s as f32 / 32768.0;
        sum_sq += n * n;
    }
    (sum_sq / pcm.len() as f32).sqrt()
}

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

    const SR: u32 = 48000;
    const FRAME: usize = 960; // 20ms @ 48kHz mono

    /// A frame of constant amplitude (a crude tone) at the given i16 level.
    fn frame(amp: i16) -> Vec<i16> {
        vec![amp; FRAME]
    }

    /// Peak absolute sample in a frame — a proxy for "how open" the gate was.
    fn peak(pcm: &[i16]) -> i16 {
        pcm.iter().copied().map(|s| s.abs()).max().unwrap_or(0)
    }

    #[test]
    fn disabled_threshold_passes_through_untouched() {
        let mut g = NoiseGate::new(SR);
        let mut f = frame(5000);
        let original = f.clone();
        assert!(g.process(&mut f, 0.0));
        assert_eq!(f, original, "a disabled gate must not alter samples");
    }

    #[test]
    fn loud_signal_opens_and_reaches_full_gain() {
        let mut g = NoiseGate::new(SR);
        // amp 10000 -> rms ~0.305, well above a 0.05 threshold.
        // After a couple of frames the attack ramp should be complete.
        let mut last = 0;
        for _ in 0..3 {
            let mut f = frame(10000);
            assert!(g.process(&mut f, 0.05), "loud frame must transmit");
            last = peak(&f);
        }
        assert!(last >= 9900, "gain should reach ~1.0 on sustained loud input, got peak {last}");
    }

    #[test]
    fn attack_is_gradual_not_a_hard_jump() {
        let mut g = NoiseGate::new(SR);
        let mut f = frame(10000);
        g.process(&mut f, 0.05);
        // 5ms attack @48k = 240 samples; across a 960-sample frame the gain ramps
        // 0->1, so the early samples are well below full scale (no instant click).
        assert!(f[0].abs() < 5000, "attack should start near zero, got {}", f[0]);
        assert!(f[FRAME - 1].abs() > 9000, "attack should complete within the frame");
    }

    #[test]
    fn quiet_after_loud_is_held_open_then_released() {
        let mut g = NoiseGate::new(SR);
        // Open it.
        for _ in 0..3 {
            let mut f = frame(10000);
            g.process(&mut f, 0.05);
        }
        // First quiet frame right after speech: hold keeps it open (not chopped).
        let mut q = frame(50); // rms ~0.0015, below close (0.03)
        assert!(g.process(&mut q, 0.05), "first quiet frame must stay open (hangover)");
        assert!(peak(&q) > 0, "held-open frame must not be silenced immediately");

        // Hold is 200ms = 10 frames; keep feeding quiet until it fully closes.
        let mut closed = false;
        for _ in 0..40 {
            let mut q = frame(0);
            if !g.process(&mut q, 0.05) {
                closed = true;
                break;
            }
        }
        assert!(closed, "gate must eventually close and stop transmitting after sustained silence");
    }

    #[test]
    fn hysteresis_keeps_gate_open_between_thresholds() {
        let mut g = NoiseGate::new(SR);
        // Open with a loud frame.
        let mut f = frame(10000);
        g.process(&mut f, 0.05); // open=0.05, close=0.03
        // A frame between close and open thresholds: rms ~0.04 (amp ~1310).
        let mut mid = frame(1310);
        assert!(g.process(&mut mid, 0.05), "between-threshold frame must keep an open gate open");
        assert!(g.open, "hysteresis: gate stays open above the close threshold");
    }

    #[test]
    fn closed_gate_does_not_transmit_silence() {
        let mut g = NoiseGate::new(SR);
        // Never opened; feed silence — should report don't-transmit promptly.
        let mut f = frame(0);
        assert!(!g.process(&mut f, 0.05), "an unopened gate on silence must not transmit");
    }

    #[test]
    fn frame_rms_known_values() {
        assert_eq!(frame_rms(&[]), 0.0);
        assert_eq!(frame_rms(&frame(0)), 0.0);

        let f_high = frame(16384);
        let rms_high = frame_rms(&f_high);
        assert!((rms_high - 0.5).abs() < 1e-4, "rms_high was {}", rms_high);

        let f_low = frame(3277);
        let rms_low = frame_rms(&f_low);
        assert!((rms_low - 0.1).abs() < 1e-3, "rms_low was {}", rms_low);
    }

    #[test]
    fn empty_frame_transmit_follows_gate_state() {
        let mut g = NoiseGate::new(SR);
        // fresh gate (never opened)
        assert!(!g.process(&mut [], 0.05));

        // open it with loud signals
        for _ in 0..3 {
            let mut f = frame(10000);
            assert!(g.process(&mut f, 0.05));
        }

        // now empty frame should transmit
        assert!(g.process(&mut [], 0.05));
    }

    #[test]
    fn disabled_parks_envelope_so_reenable_has_no_fade_in() {
        let mut g = NoiseGate::new(SR);
        let mut f1 = frame(5000);
        assert!(g.process(&mut f1, 0.0)); // disabled

        let mut f2 = frame(10000);
        assert!(g.process(&mut f2, 0.05)); // enabled
        assert!(f2[0].abs() > 9000, "expected first sample of enabled frame to have no fade-in, got {}", f2[0]);
    }

    #[test]
    fn hold_keeps_open_through_window_then_releases_to_closed() {
        let mut g = NoiseGate::new(SR);
        // open it
        for _ in 0..3 {
            let mut f = frame(10000);
            g.process(&mut f, 0.05);
        }

        let mut results = Vec::new();
        for _ in 0..25 {
            let mut f = frame(0);
            results.push(g.process(&mut f, 0.05));
        }

        assert!(results[4], "should still transmit at the 5th silent frame");
        assert!(!results[19], "should not transmit at the 20th silent frame");
    }

    #[test]
    fn sustained_mid_level_refreshes_hold_and_stays_open() {
        let mut g = NoiseGate::new(SR);
        // open it loud
        for _ in 0..3 {
            let mut f = frame(10000);
            g.process(&mut f, 0.05);
        }

        // feed 30x frame(1310) (rms ~0.04, between close 0.03 and open 0.05)
        for _ in 0..30 {
            let mut f = frame(1310);
            assert!(g.process(&mut f, 0.05));
        }
        assert!(g.open, "gate must stay open (hold refreshed by mid-level input)");
    }

    #[test]
    fn loud_signal_reopens_a_releasing_gate() {
        let mut g = NoiseGate::new(SR);
        // open it
        for _ in 0..3 {
            let mut f = frame(10000);
            g.process(&mut f, 0.05);
        }

        // feed silent frames to fully close
        let mut closed = false;
        for _ in 0..40 {
            let mut f = frame(0);
            if !g.process(&mut f, 0.05) {
                closed = true;
                break;
            }
        }
        assert!(closed);

        // feed frame(10000) @ 0.05 a few times
        let mut last_peak = 0;
        for _ in 0..3 {
            let mut f = frame(10000);
            assert!(g.process(&mut f, 0.05));
            last_peak = peak(&f);
        }
        assert!(g.open);
        assert!(last_peak >= 9900, "peak of the 3rd reopened frame must be >= 9900, got {}", last_peak);
    }
}