Lumotia/crates/audio/src/streaming_resample.rs

// Streaming resampler used by the live transcription session.
//
// Microphones expose whatever native rate the device supports (commonly
// 44 100 or 48 000 Hz). whisper.cpp wants 16 kHz mono `f32`. The live
// session calls `push_samples()` with each capture chunk as it arrives
// and gets back zero-or-more 16 kHz samples to enqueue into the model
// input buffer. At end-of-session it calls `flush()` once to drain any
// residual input and the resampler's internal tail.
//
// Implementation notes:
//
//   - We use rubato's `SincFixedIn` (same engine the file-level
//     `resample::resample_to_16khz` uses) so behaviour stays consistent
//     across live + file paths.
//   - rubato's fixed-in API requires a constant-size input chunk. We
//     buffer captured samples in a residual `Vec<f32>` and only feed
//     the resampler when we have a full chunk.
//   - When the input rate already matches 16 kHz we skip rubato
//     entirely and pass samples straight through (zero allocations
//     beyond the returned `Vec`).
//   - `flush()` zero-pads the residual to one final chunk, processes
//     it, then truncates the output to the proportion that came from
//     real (non-padded) samples — otherwise the trailing silence
//     produced by the padding leaks into the saved audio file.

use rubato::{
    Resampler, SincFixedIn, SincInterpolationParameters, SincInterpolationType, WindowFunction,
};

use kon_core::constants::WHISPER_SAMPLE_RATE;
use kon_core::error::{KonError, Result};

/// Number of input samples the rubato resampler consumes per `process()`
/// call. Matches the chunk size used in `resample::resample_to_16khz`.
const INPUT_CHUNK: usize = 1024;

pub enum StreamingResampler {
    /// Source is already at 16 kHz — emit input verbatim.
    Passthrough,
    /// Source is at some other rate — feed via rubato.
    Sinc {
        resampler: SincFixedIn<f32>,
        residual: Vec<f32>,
        ratio: f64,
    },
}

impl StreamingResampler {
    /// Construct a resampler that converts `from_rate` Hz mono input to
    /// 16 kHz mono output. Returns an error if `from_rate` is zero or if
    /// rubato rejects the requested ratio.
    pub fn new(from_rate: u32) -> Result<Self> {
        if from_rate == 0 {
            return Err(KonError::AudioDecodeFailed(
                "StreamingResampler: input sample rate is 0".into(),
            ));
        }

        if from_rate == WHISPER_SAMPLE_RATE {
            return Ok(Self::Passthrough);
        }

        let ratio = WHISPER_SAMPLE_RATE as f64 / from_rate as f64;

        let params = SincInterpolationParameters {
            sinc_len: 256,
            f_cutoff: 0.95,
            oversampling_factor: 128,
            interpolation: SincInterpolationType::Cubic,
            window: WindowFunction::Blackman,
        };

        let resampler = SincFixedIn::<f32>::new(
            ratio,
            1.1, // max relative jitter; mirrors the file-level resampler
            params,
            INPUT_CHUNK,
            1, // mono
        )
        .map_err(|e| KonError::AudioDecodeFailed(format!("StreamingResampler init failed: {e}")))?;

        Ok(Self::Sinc {
            resampler,
            residual: Vec::new(),
            ratio,
        })
    }

    /// Feed a fresh capture chunk and return any 16 kHz samples that are
    /// ready to dispatch. The caller may pass any length; samples that
    /// don't yet form a complete `INPUT_CHUNK` are buffered internally
    /// and emitted on a later call (or on `flush()`).
    pub fn push_samples(&mut self, mono: &[f32]) -> Result<Vec<f32>> {
        match self {
            Self::Passthrough => Ok(mono.to_vec()),
            Self::Sinc {
                resampler,
                residual,
                ..
            } => {
                if mono.is_empty() {
                    return Ok(Vec::new());
                }
                residual.extend_from_slice(mono);

                let mut out: Vec<f32> = Vec::new();
                while residual.len() >= INPUT_CHUNK {
                    let chunk: Vec<f32> = residual.drain(..INPUT_CHUNK).collect();
                    let input = vec![chunk];
                    let result = resampler.process(&input, None).map_err(|e| {
                        KonError::AudioDecodeFailed(format!(
                            "StreamingResampler process failed: {e}"
                        ))
                    })?;
                    if let Some(channel) = result.into_iter().next() {
                        out.extend_from_slice(&channel);
                    }
                }
                Ok(out)
            }
        }
    }

    /// Drain any residual samples and return the final 16 kHz output.
    /// Called once when the live session is stopping. Subsequent calls
    /// return an empty `Vec`.
    pub fn flush(&mut self) -> Result<Vec<f32>> {
        match self {
            Self::Passthrough => Ok(Vec::new()),
            Self::Sinc {
                resampler,
                residual,
                ratio,
            } => {
                if residual.is_empty() {
                    return Ok(Vec::new());
                }

                let leftover = residual.len();
                let mut chunk = std::mem::take(residual);
                chunk.resize(INPUT_CHUNK, 0.0);

                let input = vec![chunk];
                let result = resampler.process(&input, None).map_err(|e| {
                    KonError::AudioDecodeFailed(format!("StreamingResampler flush failed: {e}"))
                })?;

                let Some(mut out) = result.into_iter().next() else {
                    return Ok(Vec::new());
                };

                // Trim padding-induced output: keep only the proportion
                // of samples that came from real input, not from the
                // zeros we used to fill the chunk.
                let real_out = ((leftover as f64) * *ratio).round() as usize;
                if real_out < out.len() {
                    out.truncate(real_out);
                }
                Ok(out)
            }
        }
    }
}

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

    #[test]
    fn passthrough_at_16khz() {
        let mut r = StreamingResampler::new(16_000).unwrap();
        let out = r.push_samples(&[0.1, 0.2, 0.3]).unwrap();
        assert_eq!(out, vec![0.1, 0.2, 0.3]);
        assert!(r.flush().unwrap().is_empty());
    }

    #[test]
    fn rejects_zero_rate() {
        assert!(StreamingResampler::new(0).is_err());
    }

    #[test]
    fn streaming_48k_to_16k_preserves_duration() {
        let from_rate = 48_000u32;
        let secs = 1.0;
        let n = (from_rate as f64 * secs) as usize;
        let samples: Vec<f32> = (0..n).map(|i| (i as f32 * 0.001).sin()).collect();

        let mut r = StreamingResampler::new(from_rate).unwrap();

        // Push in irregular chunks to exercise the residual buffer.
        let mut produced: Vec<f32> = Vec::new();
        for window in samples.chunks(700) {
            produced.extend(r.push_samples(window).unwrap());
        }
        produced.extend(r.flush().unwrap());

        let out_secs = produced.len() as f64 / WHISPER_SAMPLE_RATE as f64;
        assert!(
            (out_secs - secs).abs() < 0.05,
            "expected ~{secs}s of 16 kHz output, got {out_secs}s ({} samples)",
            produced.len(),
        );
    }

    #[test]
    fn flush_after_no_input_is_empty() {
        let mut r = StreamingResampler::new(48_000).unwrap();
        assert!(r.flush().unwrap().is_empty());
    }
}