alpha2_audio/
stretch.rs

1//! Speed change that keeps pitch: waveform-similarity overlap-add (WSOLA).
2//!
3//! Grain `k` is written at output sample `k * hop` through a periodic Hann window twice the hop,
4//! so neighbours sum to one. It is read at the caller's source position for `k * hop`, shifted
5//! within the tolerance to best continue the previous grain. The search runs on the channel mean,
6//! coarse on sums of four samples, then at full rate.
7//!
8//! Offsets chain from grain to grain, so chains are cut into segments of
9//! [`STRETCH_SEGMENT_GRAINS`], each restarting from offset zero [`STRETCH_WARMUP_GRAINS`] grains
10//! early. Every output sample therefore depends only on its index, never on blocks or seeks.
11//! Sums run in a fixed order without FMA, so native and wasm agree bit for bit.
12
13use std::f32::consts::PI;
14
15/// Grains whose offsets one chain decides.
16pub const STRETCH_SEGMENT_GRAINS: i64 = 32;
17
18/// Grains a segment's chain runs before its first, from offset zero.
19pub const STRETCH_WARMUP_GRAINS: i64 = 8;
20
21/// `samples` from index `start`; silent elsewhere and wherever a sample is not finite.
22#[derive(Clone, Copy, Debug)]
23pub struct StretchSignal<'a> {
24    /// Index of the first sample.
25    pub start: i64,
26    /// The samples.
27    pub samples: &'a [f32],
28}
29
30impl StretchSignal<'_> {
31    fn at(&self, index: i64) -> f32 {
32        let value = index
33            .checked_sub(self.start)
34            .and_then(|at| usize::try_from(at).ok())
35            .and_then(|at| self.samples.get(at).copied())
36            .unwrap_or(0.0);
37        if value.is_finite() { value } else { 0.0 }
38    }
39}
40
41/// WSOLA's grid at one sample rate.
42#[derive(Clone, Debug)]
43pub struct Stretch {
44    hop: i64,
45    tolerance: i64,
46    window: Vec<f32>,
47}
48
49impl Stretch {
50    /// A 20 ms hop, 40 ms window and 10 ms tolerance at `sample_rate`.
51    #[must_use]
52    pub fn new(sample_rate: u32) -> Self {
53        let rate = f64::from(sample_rate.max(1));
54        #[allow(clippy::cast_possible_truncation)]
55        let hop = ((rate * 0.020).round() as i64).max(1);
56        #[allow(clippy::cast_possible_truncation)]
57        let tolerance = (rate * 0.010).round() as i64;
58        let length = 2 * hop;
59        #[allow(clippy::cast_precision_loss)]
60        let window = (0..length)
61            .map(|index| {
62                let value = libm::sinf(PI * index as f32 / length as f32);
63                value * value
64            })
65            .collect();
66        Self {
67            hop,
68            tolerance,
69            window,
70        }
71    }
72
73    /// Output samples between grains.
74    #[must_use]
75    pub const fn hop(&self) -> i64 {
76        self.hop
77    }
78
79    /// How far, in samples, a grain may move from its position.
80    #[must_use]
81    pub const fn tolerance(&self) -> i64 {
82        self.tolerance
83    }
84
85    /// Samples one grain spans.
86    #[must_use]
87    pub const fn window(&self) -> i64 {
88        2 * self.hop
89    }
90
91    /// The first and last grains whose windows touch output samples `[first, end)`.
92    #[must_use]
93    pub fn grains(&self, first: i64, end: i64) -> (i64, i64) {
94        (
95            first.div_euclid(self.hop) - 1,
96            (end - 1).div_euclid(self.hop),
97        )
98    }
99
100    /// The segment whose chain decides grain `grain`'s offset.
101    #[must_use]
102    pub const fn segment(grain: i64) -> i64 {
103        grain.div_euclid(STRETCH_SEGMENT_GRAINS)
104    }
105
106    /// The inclusive grain range `segment`'s chain runs over, warm-up included.
107    #[must_use]
108    pub const fn chain(segment: i64) -> (i64, i64) {
109        let first = segment * STRETCH_SEGMENT_GRAINS;
110        (
111            first - STRETCH_WARMUP_GRAINS,
112            first + STRETCH_SEGMENT_GRAINS - 1,
113        )
114    }
115
116    /// The signal samples `[from, to)` that deciding and reading grains `first ..= last` can
117    /// touch, with grain `k` at `position(k)`.
118    pub fn reach(&self, first: i64, last: i64, position: impl Fn(i64) -> i64) -> (i64, i64) {
119        let span = self.window() + self.tolerance;
120        let mut from = i64::MAX;
121        let mut to = i64::MIN;
122        let mut previous = None;
123        for grain in first..=last {
124            let at = position(grain);
125            from = from.min(at - self.tolerance);
126            to = to.max(at + span);
127            if let Some(previous) = previous {
128                to = to.max(previous + self.hop + span);
129            }
130            previous = Some(at);
131        }
132        (from, to)
133    }
134
135    /// The signal samples `[from, to)` that reading output `[first, end)` can touch, including
136    /// the full chains of their segments.
137    pub fn reach_output(&self, first: i64, end: i64, position: impl Fn(i64) -> i64) -> (i64, i64) {
138        let (low, high) = self.grains(first, end);
139        let mut reach = (i64::MAX, i64::MIN);
140        for segment in Self::segment(low)..=Self::segment(high) {
141            let (from, to) = Self::chain(segment);
142            let (from, to) = self.reach(from, to, &position);
143            reach = (reach.0.min(from), reach.1.max(to));
144        }
145        reach
146    }
147
148    /// The offsets `segment`'s chain gives its grains, in order, searched on `mono`.
149    pub fn offsets(
150        &self,
151        segment: i64,
152        position: impl Fn(i64) -> i64,
153        mono: StretchSignal<'_>,
154    ) -> Vec<i64> {
155        let (first, last) = Self::chain(segment);
156        let own = segment * STRETCH_SEGMENT_GRAINS;
157        let (from, to) = self.reach(first, last, &position);
158        // Sums of four samples from each index, for the coarse search.
159        let boxes: Vec<f32> = (from..to)
160            .map(|index| {
161                ((mono.at(index) + mono.at(index + 1)) + mono.at(index + 2)) + mono.at(index + 3)
162            })
163            .collect();
164        let boxes = StretchSignal {
165            start: from,
166            samples: &boxes,
167        };
168        let mut offsets = Vec::with_capacity(usize::try_from(STRETCH_SEGMENT_GRAINS).unwrap_or(0));
169        let mut read = position(first);
170        for grain in first + 1..=last {
171            let at = position(grain);
172            let offset = self.search(read + self.hop, at, mono, boxes);
173            read = at + offset;
174            if grain >= own {
175                offsets.push(offset);
176            }
177        }
178        offsets
179    }
180
181    /// The offset within the tolerance with the largest normalized correlation between `at` and
182    /// `reference`; ties go to the smallest magnitude, then the negative.
183    fn search(
184        &self,
185        reference: i64,
186        at: i64,
187        mono: StretchSignal<'_>,
188        boxes: StretchSignal<'_>,
189    ) -> i64 {
190        let length = self.window();
191        let coarse = self.tolerance / 4;
192        let mut best = (0, f64::NEG_INFINITY);
193        for step in nearest_first(-coarse, coarse) {
194            let score = correlation(boxes, reference, at + 4 * step, length / 4, 4);
195            if score > best.1 {
196                best = (step, score);
197            }
198        }
199        let centre = 4 * best.0;
200        let mut best = (0, f64::NEG_INFINITY);
201        let low = (centre - 4).max(-self.tolerance);
202        let high = (centre + 4).min(self.tolerance);
203        for offset in nearest_first(low, high) {
204            let score = correlation(mono, reference, at + offset, length, 1);
205            if score > best.1 {
206                best = (offset, score);
207            }
208        }
209        best.0
210    }
211
212    /// Adds each windowed grain `(k, read)` from `signal` into `output`, which starts at output
213    /// sample `first`.
214    pub fn overlap_add(
215        &self,
216        first: i64,
217        output: &mut [f32],
218        grains: impl IntoIterator<Item = (i64, i64)>,
219        signal: StretchSignal<'_>,
220    ) {
221        let end = first + i64::try_from(output.len()).unwrap_or(i64::MAX);
222        for (grain, read) in grains {
223            let start = grain * self.hop;
224            let from = start.max(first);
225            let to = (start + self.window()).min(end);
226            for sample in from..to {
227                let index = usize::try_from(sample - start).unwrap_or(0);
228                let out = usize::try_from(sample - first).unwrap_or(0);
229                output[out] += self.window[index] * signal.at(read + (sample - start));
230            }
231        }
232    }
233}
234
235/// A [`Stretch`] that caches the last segments' offsets, so sequential reads search each once.
236#[derive(Clone, Debug)]
237pub struct StretchReader {
238    stretch: Stretch,
239    segments: Vec<(i64, Vec<i64>)>,
240}
241
242impl StretchReader {
243    /// An empty reader at `sample_rate`.
244    #[must_use]
245    pub fn new(sample_rate: u32) -> Self {
246        Self {
247            stretch: Stretch::new(sample_rate),
248            segments: Vec::new(),
249        }
250    }
251
252    /// The grid.
253    #[must_use]
254    pub const fn stretch(&self) -> &Stretch {
255        &self.stretch
256    }
257
258    /// Adds output samples from `first` into each of `outputs` from the matching `channels`
259    /// signal, searching on `mono`. Outputs should start silent.
260    pub fn render(
261        &mut self,
262        first: i64,
263        outputs: &mut [&mut [f32]],
264        position: impl Fn(i64) -> i64,
265        channels: &[StretchSignal<'_>],
266        mono: StretchSignal<'_>,
267    ) {
268        let Some(length) = outputs.first().map(|output| output.len()) else {
269            return;
270        };
271        if length == 0 {
272            return;
273        }
274        let end = first + i64::try_from(length).unwrap_or(i64::MAX);
275        let (low, high) = self.stretch.grains(first, end);
276        let segments = Stretch::segment(low)..=Stretch::segment(high);
277        self.segments
278            .retain(|(segment, _)| segments.contains(segment));
279        for segment in segments.clone() {
280            if !self.segments.iter().any(|(kept, _)| *kept == segment) {
281                let offsets = self.stretch.offsets(segment, &position, mono);
282                self.segments.push((segment, offsets));
283            }
284        }
285        let offset = |grain: i64| {
286            let segment = Stretch::segment(grain);
287            let index = grain - segment * STRETCH_SEGMENT_GRAINS;
288            self.segments
289                .iter()
290                .find(|(kept, _)| *kept == segment)
291                .and_then(|(_, offsets)| offsets.get(usize::try_from(index).ok()?).copied())
292                .unwrap_or(0)
293        };
294        let reads: Vec<(i64, i64)> = (low..=high)
295            .map(|grain| (grain, position(grain) + offset(grain)))
296            .collect();
297        for (output, signal) in outputs.iter_mut().zip(channels) {
298            self.stretch
299                .overlap_add(first, output, reads.iter().copied(), *signal);
300        }
301    }
302}
303
304/// `low ..= high` in order of magnitude, negatives first.
305fn nearest_first(low: i64, high: i64) -> impl Iterator<Item = i64> {
306    let mut values: Vec<i64> = (low..=high).collect();
307    values.sort_by_key(|&value| (value.abs(), value));
308    values.into_iter()
309}
310
311/// Correlation of `count` samples `stride` apart from `a` and `b`, normalized by the root energy
312/// at `b`; zero where silent.
313fn correlation(signal: StretchSignal<'_>, a: i64, b: i64, count: i64, stride: i64) -> f64 {
314    let mut dot = 0.0_f64;
315    let mut energy = 0.0_f64;
316    for index in 0..count {
317        let x = f64::from(signal.at(a + index * stride));
318        let y = f64::from(signal.at(b + index * stride));
319        dot += x * y;
320        energy += y * y;
321    }
322    if energy > 0.0 {
323        dot / energy.sqrt()
324    } else {
325        0.0
326    }
327}
328
329#[cfg(test)]
330#[path = "stretch_tests.rs"]
331mod tests;
alpha2_audio/stretch.rs

alpha2_audio/
stretch.rs
