Skip to main content

aurum_core/
window.rs

1//! Rolling-window helpers for host-driven “partial-like” decode loops.
2//!
3//! Aurum does not run a background streaming decoder. Hosts (e.g. dictation apps)
4//! can push PCM into a [`crate::pcm::PcmBuffer`], then use these policies to decide
5//! when to call [`crate::LocalWhisperProvider::transcribe_pcm`] on a slice.
6//!
7//! Used as the energy/interval engine inside [`crate::partial::PartialSession`]
8//! and for simple host-driven partial loops.
9
10use crate::audio::WHISPER_SAMPLE_RATE;
11
12/// Policy aligned with common hold-to-talk UX.
13#[derive(Debug, Clone, Copy, PartialEq, Eq)]
14pub struct PartialWindowPolicy {
15    /// Minimum samples before a partial is worth running (~1 s default).
16    pub min_partial_samples: usize,
17    /// Max samples included in a partial slice (~15 s default).
18    pub window_samples: usize,
19    /// Suggested spacing between partial attempts (nanoseconds).
20    pub interval_nanos: u64,
21    /// RMS below this → skip partial (near-silence).
22    pub min_rms_bits: u32, // store as fixed-point * 1e6 for Eq; use helpers
23}
24
25impl PartialWindowPolicy {
26    /// Defaults suited to progressive hold-to-talk decode.
27    pub fn dictation() -> Self {
28        Self {
29            min_partial_samples: WHISPER_SAMPLE_RATE as usize, // 1 s
30            window_samples: WHISPER_SAMPLE_RATE as usize * 15, // 15 s
31            interval_nanos: 1_200_000_000,                     // 1.2 s
32            min_rms_bits: 500,                                 // 0.0005
33        }
34    }
35
36    pub fn min_rms(&self) -> f32 {
37        self.min_rms_bits as f32 / 1_000_000.0
38    }
39
40    pub fn with_min_rms(mut self, rms: f32) -> Self {
41        self.min_rms_bits = (rms.clamp(0.0, 1.0) * 1_000_000.0).round() as u32;
42        self
43    }
44
45    pub fn min_partial_secs(&self) -> f64 {
46        self.min_partial_samples as f64 / f64::from(WHISPER_SAMPLE_RATE)
47    }
48
49    pub fn window_secs(&self) -> f64 {
50        self.window_samples as f64 / f64::from(WHISPER_SAMPLE_RATE)
51    }
52
53    /// Enough audio accumulated for a partial attempt?
54    pub fn can_run_partial(&self, sample_count: usize) -> bool {
55        sample_count >= self.min_partial_samples
56    }
57
58    /// Most recent `window_samples` (or all if shorter).
59    pub fn slice_for_partial<'a>(&self, samples: &'a [f32]) -> &'a [f32] {
60        if samples.len() > self.window_samples {
61            &samples[samples.len() - self.window_samples..]
62        } else {
63            samples
64        }
65    }
66
67    /// Energy gate: skip near-silent buffers.
68    pub fn passes_energy_gate(&self, samples: &[f32]) -> bool {
69        if samples.is_empty() {
70            return false;
71        }
72        let sum: f32 = samples.iter().map(|s| s * s).sum();
73        let rms = (sum / samples.len() as f32).sqrt();
74        rms >= self.min_rms()
75    }
76
77    /// Combined: enough samples + energy.
78    pub fn should_run_partial(&self, samples: &[f32]) -> bool {
79        self.can_run_partial(samples.len()) && self.passes_energy_gate(samples)
80    }
81}
82
83/// Tracks wall-clock spacing between partial attempts.
84#[derive(Debug, Clone)]
85pub struct PartialClock {
86    policy: PartialWindowPolicy,
87    last_partial_at: Option<std::time::Instant>,
88}
89
90impl PartialClock {
91    pub fn new(policy: PartialWindowPolicy) -> Self {
92        Self {
93            policy,
94            last_partial_at: None,
95        }
96    }
97
98    pub fn policy(&self) -> &PartialWindowPolicy {
99        &self.policy
100    }
101
102    /// True if enough time has elapsed since the last partial (or never ran).
103    pub fn interval_elapsed(&self) -> bool {
104        match self.last_partial_at {
105            None => true,
106            Some(t) => t.elapsed().as_nanos() as u64 >= self.policy.interval_nanos,
107        }
108    }
109
110    /// Mark that a partial was started/completed now.
111    pub fn mark(&mut self) {
112        self.last_partial_at = Some(std::time::Instant::now());
113    }
114
115    /// Ready for another partial given full buffer samples.
116    /// Energy is evaluated on the same rolling window that would be decoded.
117    pub fn ready(&self, samples: &[f32]) -> bool {
118        if !self.interval_elapsed() || !self.policy.can_run_partial(samples.len()) {
119            return false;
120        }
121        let window = self.policy.slice_for_partial(samples);
122        self.policy.passes_energy_gate(window)
123    }
124
125    /// Slice + mark if ready; returns `None` if not ready.
126    pub fn take_partial_slice<'a>(&mut self, samples: &'a [f32]) -> Option<&'a [f32]> {
127        if !self.ready(samples) {
128            return None;
129        }
130        self.mark();
131        Some(self.policy.slice_for_partial(samples))
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138
139    #[test]
140    fn slice_takes_tail() {
141        let p = PartialWindowPolicy {
142            window_samples: 4,
143            min_partial_samples: 2,
144            interval_nanos: 0,
145            min_rms_bits: 0,
146        };
147        let s = [1., 2., 3., 4., 5., 6.];
148        assert_eq!(p.slice_for_partial(&s), &[3., 4., 5., 6.]);
149    }
150
151    #[test]
152    fn energy_gate() {
153        let p = PartialWindowPolicy::dictation().with_min_rms(0.01);
154        assert!(!p.passes_energy_gate(&[0.0; 1600]));
155        assert!(p.passes_energy_gate(&[0.5; 1600]));
156    }
157
158    #[test]
159    fn clock_interval() {
160        let mut c = PartialClock::new(PartialWindowPolicy {
161            min_partial_samples: 1,
162            window_samples: 100,
163            interval_nanos: 10_000_000_000, // 10s
164            min_rms_bits: 0,
165        });
166        let s = [0.1; 10];
167        assert!(c.ready(&s));
168        c.mark();
169        assert!(!c.ready(&s));
170    }
171}