gemini_adk_fluent_rs/voice/dsp/
mod.rs

1//! DSP foundation for the mic chain: a float audio bus, a stage contract,
2//! and a metered chain runner.
3//!
4//! The original mic chain passed `Vec<i16>` from stage to stage — every hop
5//! re-quantized (adding noise floor) and every boundary was a hidden clip
6//! point. This module is the engineer's version: samples are converted to
7//! `f32` **once** on entry, every stage processes in float with headroom,
8//! and one saturating conversion happens at the exit — where clipping is
9//! *counted*, not silent.
10//!
11//! ```ignore
12//! // `ignore`: the resampler/STFT stages need the `dsp` feature and the
13//! // denoiser the `denoise` feature; `live` is a `Live` builder.
14//! let chain = DspChain::new(16_000)
15//!     .stage(HighPass::speech_default(16_000))   // DC / rumble removal
16//!     .stage(IntStage::new(Denoiser::new(16_000))) // legacy i16 stage, one boundary
17//!     .stage(Agc::default_speech())
18//!     .stage(Limiter::default_ceiling());
19//! let metrics = chain.metrics();                  // live per-stage meters
20//! live.mic_processor(chain);                      // drop into the existing seam
21//! ```
22//!
23//! # Design rules
24//!
25//! - **Allocation-free steady state**: scratch buffers are owned by the
26//!   chain and stages; the hot path only does arithmetic. (A stage may
27//!   resize its output — resamplers legitimately change length — but must
28//!   not allocate per call once warmed.)
29//! - **Uniform measurement**: the chain, not the stages, meters peak/RMS
30//!   in and out of every stage plus exit clipping, so every stage is
31//!   observed identically and stages stay pure.
32//! - **Latency is declared**: every stage reports its group delay via
33//!   [`DspStage::latency_samples`]; [`DspChain::total_latency_samples`]
34//!   sums the chain's causal budget so turn-commit timestamps can cite it.
35//!
36//! # Canonical stage order
37//!
38//! `HPF → AEC → denoise → AGC → gate → limiter` — each stage assumes what
39//! the previous one guarantees: echo cancellation needs the *linear* signal
40//! (before the nonlinear denoiser breaks the echo-path model), gain control
41//! wants denoised speech so it does not amplify noise, and the limiter is
42//! last so nothing after it can clip.
43
44#[cfg(feature = "dsp")]
45pub mod aec;
46#[cfg(feature = "dsp")]
47pub mod resample;
48pub mod stages;
49#[cfg(feature = "dsp")]
50pub mod stft;
51
52#[cfg(feature = "dsp")]
53pub use aec::{Aec, AecConfig, AecFarEnd};
54#[cfg(feature = "dsp")]
55pub use resample::SincResampler;
56pub use stages::{Agc, HighPass, Limiter};
57#[cfg(feature = "dsp")]
58pub use stft::{Identity, SpectralFloor, SpectralStage, Stft};
59
60use std::sync::Arc;
61use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
62
63use serde::Serialize;
64
65use gemini_adk_rs::live::InputAudioProcessor;
66
67/// One block of audio moving through the chain: `f32` samples in
68/// `[-1.0, 1.0]` nominal range (headroom above is legal between stages)
69/// plus the sample rate a stage may change (resamplers).
70pub struct AudioBus<'a> {
71    /// The samples. Stages mutate in place and may change the length.
72    pub samples: &'a mut Vec<f32>,
73    /// Sample rate of `samples` in Hz.
74    pub sample_rate: u32,
75}
76
77/// A single processing stage on the float bus.
78pub trait DspStage: Send {
79    /// Short stable name, shown in metrics snapshots.
80    fn name(&self) -> &'static str;
81    /// Process one block in place.
82    fn process(&mut self, bus: &mut AudioBus);
83    /// Group delay this stage introduces, in samples at the bus rate
84    /// (lookahead, filter delay, block buffering). Default 0.
85    fn latency_samples(&self) -> usize {
86        0
87    }
88}
89
90/// Wrap a legacy integer-domain [`InputAudioProcessor`] (e.g. the RNNoise
91/// [`Denoiser`](crate::voice::Denoiser) or [`NoiseGate`](crate::voice::NoiseGate))
92/// as a [`DspStage`]. This is the *one* deliberate int boundary in a float
93/// chain — the cost of reusing a proven stage unchanged.
94pub struct IntStage<P: InputAudioProcessor> {
95    inner: P,
96    name: &'static str,
97    scratch: Vec<i16>,
98    latency: usize,
99}
100
101impl<P: InputAudioProcessor> IntStage<P> {
102    /// Wrap `inner`, reported under `name` in metrics.
103    pub fn named(inner: P, name: &'static str) -> Self {
104        Self {
105            inner,
106            name,
107            scratch: Vec::new(),
108            latency: 0,
109        }
110    }
111
112    /// Declare the wrapped processor's internal buffering (it cannot
113    /// declare it itself — the integer trait has no latency contract).
114    /// E.g. the RNNoise denoiser buffers one 10 ms block: 160 samples.
115    pub fn with_latency(mut self, samples: usize) -> Self {
116        self.latency = samples;
117        self
118    }
119}
120
121impl<P: InputAudioProcessor> DspStage for IntStage<P> {
122    fn name(&self) -> &'static str {
123        self.name
124    }
125
126    fn latency_samples(&self) -> usize {
127        self.latency
128    }
129
130    fn process(&mut self, bus: &mut AudioBus) {
131        self.scratch.clear();
132        self.scratch.extend(
133            bus.samples
134                .iter()
135                .map(|&s| (s * 32768.0).round().clamp(-32768.0, 32767.0) as i16),
136        );
137        self.inner.process_frame(&mut self.scratch);
138        bus.samples.clear();
139        bus.samples
140            .extend(self.scratch.iter().map(|&s| f32::from(s) / 32768.0));
141    }
142}
143
144/// Live meters for one stage, updated per block, readable concurrently.
145#[derive(Default)]
146struct StageMeter {
147    /// Max |sample| seen at stage output since start (f32 bits).
148    peak_out: AtomicU32,
149    /// EWMA of block RMS at stage output (f32 bits; single writer).
150    rms_out: AtomicU32,
151    /// Blocks processed.
152    blocks: AtomicU64,
153}
154
155/// Point-in-time view of one stage's meters.
156#[derive(Debug, Clone, Serialize)]
157pub struct StageSnapshot {
158    /// Stage name.
159    pub name: &'static str,
160    /// Max |sample| at the stage output since start.
161    pub peak_out: f32,
162    /// Smoothed RMS at the stage output.
163    pub rms_out: f32,
164    /// Declared group delay in samples.
165    pub latency_samples: usize,
166}
167
168/// Point-in-time view of the whole chain.
169#[derive(Debug, Clone, Serialize)]
170pub struct ChainSnapshot {
171    /// Per-stage meters, in chain order.
172    pub stages: Vec<StageSnapshot>,
173    /// Samples clipped at the exit conversion since start.
174    pub exit_clipped: u64,
175    /// Total declared group delay in samples at the bus rate.
176    pub total_latency_samples: usize,
177    /// Blocks processed.
178    pub blocks: u64,
179}
180
181/// Shared metrics handle — clone freely; reading never blocks the chain.
182#[derive(Clone)]
183pub struct ChainMetrics {
184    inner: Arc<MetricsInner>,
185}
186
187struct MetricsInner {
188    meters: Vec<StageMeter>,
189    names: Vec<&'static str>,
190    latencies: Vec<usize>,
191    exit_clipped: AtomicU64,
192    blocks: AtomicU64,
193}
194
195impl ChainMetrics {
196    /// A point-in-time snapshot of every stage's meters.
197    pub fn snapshot(&self) -> ChainSnapshot {
198        let stages = self
199            .inner
200            .meters
201            .iter()
202            .zip(&self.inner.names)
203            .zip(&self.inner.latencies)
204            .map(|((meter, name), &latency)| StageSnapshot {
205                name,
206                peak_out: f32::from_bits(meter.peak_out.load(Ordering::Relaxed)),
207                rms_out: f32::from_bits(meter.rms_out.load(Ordering::Relaxed)),
208                latency_samples: latency,
209            })
210            .collect();
211        ChainSnapshot {
212            stages,
213            exit_clipped: self.inner.exit_clipped.load(Ordering::Relaxed),
214            total_latency_samples: self.inner.latencies.iter().sum(),
215            blocks: self.inner.blocks.load(Ordering::Relaxed),
216        }
217    }
218}
219
220/// The metered float chain. Build with [`stage`](Self::stage), hand to the
221/// existing `mic_processor(..)` seam — it implements [`InputAudioProcessor`].
222pub struct DspChain {
223    stages: Vec<Box<dyn DspStage>>,
224    sample_rate: u32,
225    bus: Vec<f32>,
226    metrics: Option<ChainMetrics>,
227}
228
229impl DspChain {
230    /// An empty chain at `sample_rate` (an empty chain is bit-transparent).
231    pub fn new(sample_rate: u32) -> Self {
232        Self {
233            stages: Vec::new(),
234            sample_rate,
235            bus: Vec::new(),
236            metrics: None,
237        }
238    }
239
240    /// Append a stage (chain order is processing order).
241    pub fn stage(mut self, stage: impl DspStage + 'static) -> Self {
242        self.stages.push(Box::new(stage));
243        self.metrics = None; // rebuilt lazily to match the stage list
244        self
245    }
246
247    /// Shared live-metrics handle for this chain.
248    pub fn metrics(&mut self) -> ChainMetrics {
249        self.ensure_metrics();
250        self.metrics.as_ref().expect("just built").clone()
251    }
252
253    /// Total declared group delay across all stages, in samples.
254    pub fn total_latency_samples(&self) -> usize {
255        self.stages.iter().map(|s| s.latency_samples()).sum()
256    }
257
258    fn ensure_metrics(&mut self) {
259        if self.metrics.is_none() {
260            self.metrics = Some(ChainMetrics {
261                inner: Arc::new(MetricsInner {
262                    meters: self.stages.iter().map(|_| StageMeter::default()).collect(),
263                    names: self.stages.iter().map(|s| s.name()).collect(),
264                    latencies: self.stages.iter().map(|s| s.latency_samples()).collect(),
265                    exit_clipped: AtomicU64::new(0),
266                    blocks: AtomicU64::new(0),
267                }),
268            });
269        }
270    }
271}
272
273impl InputAudioProcessor for DspChain {
274    fn process_frame(&mut self, frame: &mut Vec<i16>) {
275        self.ensure_metrics();
276        // Entry: one int -> float conversion.
277        self.bus.clear();
278        self.bus
279            .extend(frame.iter().map(|&s| f32::from(s) / 32768.0));
280
281        let metrics = self.metrics.as_ref().expect("ensured").inner.clone();
282        // Rate changes (a resampler mid-chain) propagate stage-to-stage
283        // WITHIN this frame only; the chain's own input rate is fixed, so
284        // the next frame's fresh PCM is labeled correctly again.
285        let mut rate = self.sample_rate;
286        for (stage, meter) in self.stages.iter_mut().zip(&metrics.meters) {
287            let mut bus = AudioBus {
288                samples: &mut self.bus,
289                sample_rate: rate,
290            };
291            stage.process(&mut bus);
292            rate = bus.sample_rate;
293
294            // Uniform metering at the stage output.
295            let mut peak = 0.0f32;
296            let mut energy = 0.0f64;
297            for &s in self.bus.iter() {
298                peak = peak.max(s.abs());
299                energy += f64::from(s) * f64::from(s);
300            }
301            let rms = if self.bus.is_empty() {
302                0.0
303            } else {
304                (energy / self.bus.len() as f64).sqrt() as f32
305            };
306            meter.peak_out.fetch_max(peak.to_bits(), Ordering::Relaxed);
307            let prev = f32::from_bits(meter.rms_out.load(Ordering::Relaxed));
308            let ewma = if meter.blocks.load(Ordering::Relaxed) == 0 {
309                rms
310            } else {
311                prev * 0.9 + rms * 0.1
312            };
313            meter.rms_out.store(ewma.to_bits(), Ordering::Relaxed);
314            meter.blocks.fetch_add(1, Ordering::Relaxed);
315        }
316
317        // Exit: one saturating float -> int conversion, clipping counted.
318        let mut clipped = 0u64;
319        frame.clear();
320        frame.extend(self.bus.iter().map(|&s| {
321            let scaled = (s * 32768.0).round();
322            if !(-32768.0..=32767.0).contains(&scaled) {
323                clipped += 1;
324            }
325            scaled.clamp(-32768.0, 32767.0) as i16
326        }));
327        if clipped > 0 {
328            metrics.exit_clipped.fetch_add(clipped, Ordering::Relaxed);
329        }
330        metrics.blocks.fetch_add(1, Ordering::Relaxed);
331    }
332}
333
334#[cfg(test)]
335mod tests {
336    use super::*;
337
338    struct Gain(f32);
339    impl DspStage for Gain {
340        fn name(&self) -> &'static str {
341            "gain"
342        }
343        fn process(&mut self, bus: &mut AudioBus) {
344            for s in bus.samples.iter_mut() {
345                *s *= self.0;
346            }
347        }
348    }
349
350    #[cfg(feature = "dsp")]
351    #[test]
352    fn rate_change_does_not_persist_across_frames() {
353        // A mid-chain resampler rewrites the bus rate for LATER stages in
354        // the SAME frame only. The next frame's fresh input PCM arrives at
355        // the chain's input rate again — persisting the output rate fed
356        // 16 kHz-labeled audio to a resampler built for 48 kHz (debug
357        // panic; silent rate corruption in release).
358        let mut chain = DspChain::new(48_000).stage(
359            crate::voice::dsp::resample::SincResampler::new(48_000, 16_000),
360        );
361        let mut frame = vec![1000i16; 960]; // 20 ms at 48 kHz
362        chain.process_frame(&mut frame);
363        let mut frame2 = vec![1000i16; 960];
364        chain.process_frame(&mut frame2); // panicked before the fix
365    }
366
367    #[test]
368    fn empty_chain_is_bit_transparent() {
369        let mut chain = DspChain::new(16_000);
370        let original: Vec<i16> = (-40..40).map(|i| (i * 400) as i16).collect();
371        let mut frame = original.clone();
372        chain.process_frame(&mut frame);
373        assert_eq!(frame, original);
374    }
375
376    #[test]
377    fn exit_clipping_is_counted_not_silent() {
378        let mut chain = DspChain::new(16_000).stage(Gain(4.0));
379        let metrics = chain.metrics();
380        let mut frame = vec![20_000i16; 160];
381        chain.process_frame(&mut frame);
382        assert!(frame.iter().all(|&s| s == i16::MAX || s == i16::MIN));
383        let snap = metrics.snapshot();
384        assert_eq!(snap.exit_clipped, 160);
385        assert!(snap.stages[0].peak_out > 2.0);
386    }
387
388    #[test]
389    fn int_adapter_round_trips_a_passthrough() {
390        struct Nop;
391        impl InputAudioProcessor for Nop {
392            fn process_frame(&mut self, _frame: &mut Vec<i16>) {}
393        }
394        let mut chain = DspChain::new(16_000).stage(IntStage::named(Nop, "nop"));
395        let original: Vec<i16> = (0..160).map(|i| (i * 100 - 8000) as i16).collect();
396        let mut frame = original.clone();
397        chain.process_frame(&mut frame);
398        // Symmetric 1/32768 scaling makes the boundary exact for i16 values.
399        assert_eq!(frame, original);
400    }
401
402    #[test]
403    fn metrics_report_stage_names_and_latency() {
404        struct Delayed;
405        impl DspStage for Delayed {
406            fn name(&self) -> &'static str {
407                "delayed"
408            }
409            fn process(&mut self, _bus: &mut AudioBus) {}
410            fn latency_samples(&self) -> usize {
411                80
412            }
413        }
414        let mut chain = DspChain::new(16_000).stage(Delayed).stage(Gain(1.0));
415        let snap = chain.metrics().snapshot();
416        assert_eq!(snap.stages.len(), 2);
417        assert_eq!(snap.stages[0].name, "delayed");
418        assert_eq!(snap.total_latency_samples, 80);
419        assert_eq!(chain.total_latency_samples(), 80);
420    }
421}
gemini_adk_fluent_rs/voice/dsp/mod.rs

gemini_adk_fluent_rs/voice/dsp/
mod.rs
