Lumotia/crates/transcription/src/streaming/rms_vad.rs

//! RMS-energy-backed VAD chunker.
//!
//! This is the fallback backend the plan (`docs/whisper-ecosystem/
//! workstream-A.md`, Phase A.3 "Known unknowns") permits while the ort
//! 2.0.0-rc.10 vs rc.12 ecosystem conflict prevents a drop-in Silero
//! dep. The `VadChunker` trait surface is identical to what a Silero
//! backend will present, so the live-session path does not change when
//! Silero lands.
//!
//! The chunker emits a `VadChunk` when a sustained-speech region ends
//! (RMS drops below `exit_threshold` for `silence_close_samples`) or
//! when an in-progress region exceeds `max_chunk_samples` (so Whisper
//! is not fed a 30-second monolith). It applies hysteresis — an
//! `enter_threshold` higher than `exit_threshold` — so a VAD score
//! bouncing around the threshold does not toggle state every frame.

use super::{VadChunk, VadChunker};

/// Sample window used to compute a single RMS reading. 50 ms at 16
/// kHz. Shorter windows twitch on transients; longer windows blur the
/// speech-onset boundary.
const FRAME_SAMPLES: usize = 800;

/// Default thresholds tuned to match the existing `evaluate_speech_gate`
/// behaviour in `src-tauri/src/commands/live.rs`. The underlying
/// constants live in that file; this chunker exposes them as fields so
/// they can be tuned per-session without a recompile.
const DEFAULT_ENTER_RMS_THRESHOLD: f32 = 0.003;
const DEFAULT_EXIT_RMS_THRESHOLD: f32 = 0.0014;
/// Frames of sustained speech required before the chunker enters the
/// "in-speech" state. Filters out single-frame transients (keyboard
/// clicks, door closes).
const DEFAULT_SPEECH_ONSET_FRAMES: usize = 3;
/// Silence duration that closes an in-progress chunk, in samples.
/// 500 ms = 10 frames at 16 kHz / 50 ms-frames.
const DEFAULT_SILENCE_CLOSE_SAMPLES: usize = 8_000;
/// Hard cap on a single chunk. Matches the existing `CHUNK_SAMPLES`
/// (2 s) so the live-streaming experience is not delayed arbitrarily
/// by a user speaking continuously.
const DEFAULT_MAX_CHUNK_SAMPLES: usize = 32_000;
/// Sample rate the thresholds above assume. Exposed so future backends
/// (Parakeet, Moonshine) at different rates can construct a chunker
/// matching their native sample rate.
const DEFAULT_SAMPLE_RATE_HZ: u32 = 16_000;

#[derive(Debug, Clone, Copy, PartialEq)]
enum State {
    /// Nothing buffered. Waiting for the first RMS excursion over
    /// `enter_threshold`.
    Idle,
    /// In-progress speech. Samples accumulate; closes on
    /// `silence_close_samples` of sub-threshold audio or on
    /// `max_chunk_samples`.
    InSpeech,
}

pub struct RmsVadChunker {
    // Tunables
    enter_threshold: f32,
    exit_threshold: f32,
    speech_onset_frames: usize,
    silence_close_samples: usize,
    max_chunk_samples: usize,

    // Running state
    state: State,
    /// Frame-boundary reassembly: samples that did not complete a
    /// frame on the previous `push`. Always shorter than `FRAME_SAMPLES`.
    pending: Vec<f32>,
    /// Samples belonging to the current in-progress chunk (State::InSpeech).
    active_chunk: Vec<f32>,
    /// Trailing silence sample count inside the current chunk. Resets
    /// to zero whenever a speech frame is seen.
    silent_tail_samples: usize,
    /// Consecutive speech frames observed while `State::Idle`. When
    /// this hits `speech_onset_frames`, state transitions to InSpeech.
    pending_onset_frames: usize,
    /// Samples buffered from the onset window that should be attached
    /// to the front of the emitted chunk so Whisper sees the speech
    /// onset itself, not just the post-onset audio.
    onset_buffer: Vec<f32>,
    /// Absolute sample index of the next sample `push` will consume.
    next_sample_index: u64,
    /// Absolute sample index where the current in-progress chunk
    /// started. Valid only while `state == InSpeech`.
    active_chunk_start: u64,
}

impl RmsVadChunker {
    pub fn new() -> Self {
        Self::with_thresholds(
            DEFAULT_ENTER_RMS_THRESHOLD,
            DEFAULT_EXIT_RMS_THRESHOLD,
            DEFAULT_SPEECH_ONSET_FRAMES,
            DEFAULT_SILENCE_CLOSE_SAMPLES,
            DEFAULT_MAX_CHUNK_SAMPLES,
        )
    }

    pub fn with_thresholds(
        enter_threshold: f32,
        exit_threshold: f32,
        speech_onset_frames: usize,
        silence_close_samples: usize,
        max_chunk_samples: usize,
    ) -> Self {
        debug_assert!(
            exit_threshold <= enter_threshold,
            "exit_threshold must not exceed enter_threshold (hysteresis requires enter >= exit)"
        );
        Self {
            enter_threshold,
            exit_threshold,
            speech_onset_frames,
            silence_close_samples,
            max_chunk_samples,
            state: State::Idle,
            pending: Vec::new(),
            active_chunk: Vec::new(),
            silent_tail_samples: 0,
            pending_onset_frames: 0,
            onset_buffer: Vec::new(),
            next_sample_index: 0,
        active_chunk_start: 0,
        }
    }

    pub fn sample_rate_hz(&self) -> u32 {
        DEFAULT_SAMPLE_RATE_HZ
    }

    fn frame_rms(frame: &[f32]) -> f32 {
        if frame.is_empty() {
            return 0.0;
        }
        let sum_sq: f32 = frame.iter().map(|x| x * x).sum();
        (sum_sq / frame.len() as f32).sqrt()
    }

    /// Consume one complete frame's worth of samples and update state.
    /// `frame_start` is the absolute sample index of `frame[0]` in the
    /// stream fed since `reset`. Returns a `VadChunk` if this frame
    /// closed the in-progress chunk.
    fn consume_frame(&mut self, frame: Vec<f32>, frame_start: u64) -> Option<VadChunk> {
        let rms = Self::frame_rms(&frame);
        match self.state {
            State::Idle => self.consume_frame_idle(frame, frame_start, rms),
            State::InSpeech => self.consume_frame_in_speech(frame, rms),
        }
    }

    fn consume_frame_idle(
        &mut self,
        frame: Vec<f32>,
        frame_start: u64,
        rms: f32,
    ) -> Option<VadChunk> {
        if rms >= self.enter_threshold {
            self.pending_onset_frames += 1;
            // Keep a rolling buffer of onset audio so once we confirm
            // speech, the emitted chunk contains the speech attack
            // rather than starting mid-syllable.
            self.onset_buffer.extend_from_slice(&frame);
            let onset_cap = self.speech_onset_frames * FRAME_SAMPLES;
            if self.onset_buffer.len() > onset_cap {
                let overflow = self.onset_buffer.len() - onset_cap;
                self.onset_buffer.drain(..overflow);
            }

            if self.pending_onset_frames >= self.speech_onset_frames {
                // Transition: flush the onset buffer into active_chunk
                // and begin accumulating. The onset buffer includes
                // the current frame, so its start index is
                // `frame_start + FRAME_SAMPLES - onset_buffer.len()`.
                self.state = State::InSpeech;
                self.active_chunk_start = frame_start
                    .saturating_add(FRAME_SAMPLES as u64)
                    .saturating_sub(self.onset_buffer.len() as u64);
                self.active_chunk.clear();
                self.active_chunk.append(&mut self.onset_buffer);
                self.silent_tail_samples = 0;
                self.pending_onset_frames = 0;
            }
        } else {
            // Sub-threshold frame while idle — reset the onset counter
            // and drop any onset buffer. The gate demands *sustained*
            // speech, not a single frame over threshold.
            self.pending_onset_frames = 0;
            self.onset_buffer.clear();
        }
        None
    }

    fn consume_frame_in_speech(&mut self, frame: Vec<f32>, rms: f32) -> Option<VadChunk> {
        self.active_chunk.extend_from_slice(&frame);
        if rms >= self.exit_threshold {
            self.silent_tail_samples = 0;
        } else {
            self.silent_tail_samples += frame.len();
        }

        let end_of_utterance = self.silent_tail_samples >= self.silence_close_samples;
        let hit_max = self.active_chunk.len() >= self.max_chunk_samples;
        if end_of_utterance || hit_max {
            return Some(self.emit_active_chunk());
        }
        None
    }

    fn emit_active_chunk(&mut self) -> VadChunk {
        // Trim the trailing silence from the chunk so Whisper isn't
        // fed dead air at the end. Matches the "tighten chunk on
        // commit" property #25 will later tie to commit points.
        let mut samples = std::mem::take(&mut self.active_chunk);
        if self.silent_tail_samples > 0 && samples.len() > self.silent_tail_samples {
            let keep = samples.len() - self.silent_tail_samples;
            samples.truncate(keep);
        }
        let start_sample = self.active_chunk_start;

        self.state = State::Idle;
        self.silent_tail_samples = 0;
        self.pending_onset_frames = 0;
        self.onset_buffer.clear();

        VadChunk {
            start_sample,
            samples,
        }
    }
}

impl Default for RmsVadChunker {
    fn default() -> Self {
        Self::new()
    }
}

impl VadChunker for RmsVadChunker {
    fn push(&mut self, samples: &[f32]) -> Vec<VadChunk> {
        if samples.is_empty() {
            return Vec::new();
        }
        self.pending.extend_from_slice(samples);
        self.next_sample_index = self.next_sample_index.saturating_add(samples.len() as u64);

        let mut emitted = Vec::new();
        while self.pending.len() >= FRAME_SAMPLES {
            // Absolute index of the first sample in the frame we are
            // about to consume: total fed minus what is still pending.
            let frame_start = self
                .next_sample_index
                .saturating_sub(self.pending.len() as u64);
            let frame: Vec<f32> = self.pending.drain(..FRAME_SAMPLES).collect();
            if let Some(chunk) = self.consume_frame(frame, frame_start) {
                emitted.push(chunk);
            }
        }
        emitted
    }

    fn flush(&mut self) -> Option<VadChunk> {
        // Consume any tail of fewer-than-frame samples so the last
        // utterance is not lost when a user stops recording mid-word.
        if !self.pending.is_empty() {
            // The tail's absolute start index is captured BEFORE we
            // consume — same formula as in `push`.
            let frame_start = self
                .next_sample_index
                .saturating_sub(self.pending.len() as u64);
            let pad_len = FRAME_SAMPLES - self.pending.len();
            let mut padded = std::mem::take(&mut self.pending);
            padded.extend(std::iter::repeat(0.0_f32).take(pad_len));
            let _ = self.consume_frame(padded, frame_start);
        }
        if self.state == State::InSpeech {
            Some(self.emit_active_chunk())
        } else {
            None
        }
    }

    fn reset(&mut self) {
        self.state = State::Idle;
        self.pending.clear();
        self.active_chunk.clear();
        self.silent_tail_samples = 0;
        self.pending_onset_frames = 0;
        self.onset_buffer.clear();
        self.next_sample_index = 0;
        self.active_chunk_start = 0;
    }

    fn next_sample_index(&self) -> u64 {
        self.next_sample_index
    }
}

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

    /// Generate a vector of `len` samples at amplitude `amp`. The
    /// signal is a constant DC offset, which gives a deterministic
    /// RMS of exactly `amp.abs()` — simpler than a sinusoid for
    /// threshold-crossing tests.
    fn constant_signal(len: usize, amp: f32) -> Vec<f32> {
        vec![amp; len]
    }

    #[test]
    fn pure_silence_emits_nothing() {
        let mut c = RmsVadChunker::new();
        let silence = constant_signal(16_000, 0.0); // 1 s of zero
        let chunks = c.push(&silence);
        assert!(chunks.is_empty());
        assert_eq!(c.flush().is_none(), true);
    }

    #[test]
    fn below_enter_threshold_does_not_trigger() {
        let mut c = RmsVadChunker::new();
        // 0.002 is between the default exit (0.0014) and enter (0.003)
        // thresholds — must NOT transition Idle → InSpeech.
        let hum = constant_signal(16_000, 0.002);
        let chunks = c.push(&hum);
        assert!(
            chunks.is_empty(),
            "samples below enter_threshold must not trigger onset"
        );
    }

    #[test]
    fn single_loud_frame_does_not_trigger_onset() {
        let mut c = RmsVadChunker::new();
        // One frame above enter, surrounded by silence. With
        // speech_onset_frames=3 this should NOT transition.
        let mut signal = Vec::new();
        signal.extend(constant_signal(FRAME_SAMPLES, 0.0));
        signal.extend(constant_signal(FRAME_SAMPLES, 0.01)); // loud, one frame
        signal.extend(constant_signal(FRAME_SAMPLES * 4, 0.0));
        let chunks = c.push(&signal);
        assert!(
            chunks.is_empty(),
            "single-frame transient must not cross sustained-speech onset"
        );
    }

    #[test]
    fn sustained_speech_followed_by_silence_emits_one_chunk() {
        let mut c = RmsVadChunker::new();
        // 8 frames of speech (well over onset) followed by 12 frames of
        // silence (well over silence_close). Must emit exactly one
        // chunk.
        let mut signal = Vec::new();
        signal.extend(constant_signal(FRAME_SAMPLES * 8, 0.01));
        signal.extend(constant_signal(FRAME_SAMPLES * 12, 0.0));
        let chunks = c.push(&signal);
        assert_eq!(chunks.len(), 1, "one speech region → one chunk");
        let chunk = &chunks[0];
        assert!(
            !chunk.samples.is_empty(),
            "emitted chunk must contain samples"
        );
    }

    #[test]
    fn hysteresis_prevents_mid_utterance_close_on_brief_dip() {
        let mut c = RmsVadChunker::new();
        // Onset → loud → brief dip between enter and exit → loud again
        // → silence. The dip is above exit_threshold so the chunk must
        // NOT close across it.
        let loud = constant_signal(FRAME_SAMPLES * 4, 0.01);
        let dip = constant_signal(FRAME_SAMPLES, 0.002);
        let more_loud = constant_signal(FRAME_SAMPLES * 4, 0.01);
        let silence = constant_signal(FRAME_SAMPLES * 12, 0.0);
        let mut signal = Vec::new();
        signal.extend(loud);
        signal.extend(dip);
        signal.extend(more_loud);
        signal.extend(silence);
        let chunks = c.push(&signal);
        assert_eq!(
            chunks.len(),
            1,
            "hysteresis dip between enter and exit thresholds must not split a chunk"
        );
    }

    #[test]
    fn max_chunk_samples_caps_continuous_speech() {
        let mut c = RmsVadChunker::with_thresholds(
            DEFAULT_ENTER_RMS_THRESHOLD,
            DEFAULT_EXIT_RMS_THRESHOLD,
            DEFAULT_SPEECH_ONSET_FRAMES,
            DEFAULT_SILENCE_CLOSE_SAMPLES,
            FRAME_SAMPLES * 4, // tighter cap for the test
        );
        // Feed 12 frames of sustained speech with no silence break.
        // The 4-frame cap must cause at least one emission mid-stream.
        let signal = constant_signal(FRAME_SAMPLES * 12, 0.01);
        let chunks = c.push(&signal);
        assert!(
            !chunks.is_empty(),
            "continuous speech over the cap must emit at least one chunk"
        );
        for chunk in &chunks {
            assert!(
                chunk.samples.len() <= FRAME_SAMPLES * 4,
                "emitted chunk exceeded max_chunk_samples"
            );
        }
    }

    #[test]
    fn flush_emits_in_flight_speech() {
        let mut c = RmsVadChunker::new();
        // Sustained speech with NO closing silence. Without flush this
        // stays buffered; flush must surface it as a final chunk.
        let signal = constant_signal(FRAME_SAMPLES * 5, 0.01);
        let chunks = c.push(&signal);
        assert!(
            chunks.is_empty(),
            "in-progress speech with no silence close stays buffered until flush"
        );
        let final_chunk = c.flush();
        assert!(final_chunk.is_some(), "flush must emit the in-flight chunk");
    }

    #[test]
    fn flush_returns_none_when_idle() {
        let mut c = RmsVadChunker::new();
        assert!(c.flush().is_none());
        let _ = c.push(&constant_signal(16_000, 0.0));
        assert!(c.flush().is_none(), "flushing pure silence emits nothing");
    }

    #[test]
    fn reset_clears_state() {
        let mut c = RmsVadChunker::new();
        let signal = constant_signal(FRAME_SAMPLES * 5, 0.01);
        let _ = c.push(&signal);
        c.reset();
        assert_eq!(c.next_sample_index(), 0);
        // After reset, silence must not emit a chunk derived from pre-reset state.
        let silence = constant_signal(FRAME_SAMPLES * 12, 0.0);
        let chunks = c.push(&silence);
        assert!(chunks.is_empty());
        assert!(c.flush().is_none());
    }

    #[test]
    fn start_sample_includes_onset_audio() {
        let mut c = RmsVadChunker::new();
        // First 2 frames silent (so next_sample_index is advanced but
        // no onset). Then speech.
        let silence = constant_signal(FRAME_SAMPLES * 2, 0.0);
        let _ = c.push(&silence);
        assert_eq!(c.next_sample_index(), (FRAME_SAMPLES * 2) as u64);

        let speech = constant_signal(FRAME_SAMPLES * 5, 0.01);
        let closing_silence = constant_signal(FRAME_SAMPLES * 12, 0.0);
        let mut signal = Vec::new();
        signal.extend(speech);
        signal.extend(closing_silence);
        let chunks = c.push(&signal);
        assert_eq!(chunks.len(), 1);
        let chunk = &chunks[0];
        // The chunk's start_sample should reflect the absolute index
        // of the first onset-buffered sample, NOT the post-onset index.
        assert!(
            chunk.start_sample >= (FRAME_SAMPLES * 2) as u64,
            "start_sample must be at or after the pre-speech silence"
        );
        assert!(
            chunk.start_sample
                <= (FRAME_SAMPLES * 2 + FRAME_SAMPLES * DEFAULT_SPEECH_ONSET_FRAMES) as u64,
            "start_sample must not skip past the onset frames"
        );
    }
}