Skip to main content

pito_codecs/filter/
resample.rs

1use crate::media::{AudioBuffer, Error, Result};
2use rayon::prelude::*;
3
4const CROSSINGS: f64 = 48.0;
5const ROLLOFF: f64 = 0.93;
6const BETA: f64 = 9.0;
7const MAX_PHASES: u64 = 2048;
8const MAX_RATE: u32 = 1_536_000;
9const MAX_RATIO: u32 = 256;
10
11fn gcd(a: u64, b: u64) -> u64 {
12    if b == 0 { a } else { gcd(b, a % b) }
13}
14
15pub(crate) fn kaiser(r: f64, beta: f64, i0: f64) -> f64 {
16    if r.abs() >= 1.0 {
17        0.0
18    } else {
19        bessel_i0(beta * (1.0 - r * r).sqrt()) / i0
20    }
21}
22
23pub(crate) fn sinc(x: f64) -> f64 {
24    if x.abs() < 1e-12 {
25        1.0
26    } else {
27        (std::f64::consts::PI * x).sin() / (std::f64::consts::PI * x)
28    }
29}
30
31pub(crate) fn bessel_i0(x: f64) -> f64 {
32    let mut sum = 1.0;
33    let mut term = 1.0;
34    let q = x * x / 4.0;
35    for k in 1..64 {
36        term *= q / (k * k) as f64;
37        sum += term;
38        if term < sum * 1e-17 {
39            break;
40        }
41    }
42    sum
43}
44
45#[derive(Debug, Clone)]
46pub struct Resampler {
47    from: u32,
48    to: u32,
49    channels: usize,
50    up: u64,
51    down: u64,
52    phases: u64,
53    taps: usize,
54    table: Vec<f32>,
55    history: Vec<Vec<f32>>,
56    origin: i64,
57    index: i64,
58    phase: u64,
59    consumed: i64,
60    emitted: u64,
61}
62
63impl Resampler {
64    pub fn new(from: u32, to: u32, channels: u16) -> Result<Self> {
65        if from == 0 || to == 0 || channels == 0 {
66            return Err(Error::Invalid(format!(
67                "resample {from} Hz to {to} Hz with {channels} channels"
68            )));
69        }
70        if from.max(to) > MAX_RATE || from.max(to) / from.min(to) > MAX_RATIO {
71            return Err(Error::Invalid(format!(
72                "resample {from} Hz to {to} Hz: rates go up to {MAX_RATE} Hz and {MAX_RATIO} to 1"
73            )));
74        }
75        let g = gcd(from as u64, to as u64);
76        let (up, down) = (to as u64 / g, from as u64 / g);
77        let phases = up.min(MAX_PHASES);
78        let f = (up as f64 / down as f64).min(1.0) * ROLLOFF;
79        let half = (CROSSINGS / f).ceil() as usize;
80        let taps = 2 * half;
81        let i0 = bessel_i0(BETA);
82        let mut table = vec![0.0f32; (phases as usize + 1) * taps];
83        for p in 0..=phases as usize {
84            let frac = p as f64 / phases as f64;
85            let row = &mut table[p * taps..(p + 1) * taps];
86            let mut w = vec![0.0f64; taps];
87            for (m, v) in w.iter_mut().enumerate() {
88                let t = frac + (half as f64 - 1.0) - m as f64;
89                *v = f * sinc(f * t) * kaiser(t / half as f64, BETA, i0);
90            }
91            let sum: f64 = w.iter().sum();
92            for (o, v) in row.iter_mut().zip(&w) {
93                *o = (v / sum) as f32;
94            }
95        }
96        let channels = channels as usize;
97        Ok(Resampler {
98            from,
99            to,
100            channels,
101            up,
102            down,
103            phases,
104            taps,
105            table,
106            history: vec![vec![0.0; half - 1]; channels],
107            origin: -(half as i64 - 1),
108            index: 0,
109            phase: 0,
110            consumed: 0,
111            emitted: 0,
112        })
113    }
114
115    pub fn rates(&self) -> (u32, u32) {
116        (self.from, self.to)
117    }
118
119    pub fn process(&mut self, input: &AudioBuffer, out: &mut AudioBuffer) -> Result<()> {
120        if input.rate != self.from || input.channels as usize != self.channels {
121            return Err(Error::Invalid(format!(
122                "resampler expects {} Hz {} channels, got {} Hz {} channels",
123                self.from, self.channels, input.rate, input.channels
124            )));
125        }
126        out.rate = self.to;
127        out.channels = self.channels as u16;
128        out.samples.clear();
129        let frames = input.frames();
130        for (c, h) in self.history.iter_mut().enumerate() {
131            h.extend(
132                input
133                    .samples
134                    .iter()
135                    .skip(c)
136                    .step_by(self.channels)
137                    .take(frames),
138            );
139        }
140        self.consumed += frames as i64;
141        self.run(self.consumed, u64::MAX, out);
142        Ok(())
143    }
144
145    pub fn flush(&mut self, out: &mut AudioBuffer) {
146        out.rate = self.to;
147        out.channels = self.channels as u16;
148        out.samples.clear();
149        let total = (self.consumed as u128 * self.up as u128).div_ceil(self.down as u128) as u64;
150        let pad = self.taps / 2;
151        for h in self.history.iter_mut() {
152            h.extend(std::iter::repeat_n(0.0, pad));
153        }
154        self.run(self.consumed + pad as i64, total, out);
155    }
156
157    fn run(&mut self, available: i64, limit: u64, out: &mut AudioBuffer) {
158        #[cfg(target_arch = "x86_64")]
159        if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma")
160        {
161            unsafe { self.run_fma(available, limit, out) };
162            return;
163        }
164        self.run_plain(available, limit, out);
165    }
166
167    #[cfg(target_arch = "x86_64")]
168    #[target_feature(enable = "avx2,fma")]
169    unsafe fn run_fma(&mut self, available: i64, limit: u64, out: &mut AudioBuffer) {
170        self.run_plain(available, limit, out);
171    }
172
173    #[inline(always)]
174    fn tap(&self, phase: u64, h: &[f32]) -> f32 {
175        if self.phases == self.up {
176            let p = phase as usize;
177            dot(&self.table[p * self.taps..(p + 1) * self.taps], h)
178        } else {
179            let pos = phase as f64 * self.phases as f64 / self.up as f64;
180            let p = pos.floor() as usize;
181            let t = (pos - p as f64) as f32;
182            let a = &self.table[p * self.taps..(p + 1) * self.taps];
183            let b = &self.table[(p + 1) * self.taps..(p + 2) * self.taps];
184            dot(a, h) * (1.0 - t) + dot(b, h) * t
185        }
186    }
187
188    #[inline(always)]
189    fn run_plain(&mut self, available: i64, limit: u64, out: &mut AudioBuffer) {
190        let half = (self.taps / 2) as i64;
191        while self.index + half < available && self.emitted < limit {
192            let at = (self.index - (half - 1) - self.origin) as usize;
193            for h in &self.history {
194                out.samples
195                    .push(self.tap(self.phase, &h[at..at + self.taps]));
196            }
197            self.emitted += 1;
198            self.phase += self.down;
199            self.index += (self.phase / self.up) as i64;
200            self.phase %= self.up;
201        }
202        let keep_from = (self.index - (half - 1) - self.origin).max(0) as usize;
203        let keep_from = keep_from.min(self.history.first().map_or(0, Vec::len));
204        if keep_from > 0 {
205            for h in self.history.iter_mut() {
206                h.drain(..keep_from);
207            }
208            self.origin += keep_from as i64;
209        }
210    }
211
212    fn render(&self, padded: &[Vec<f32>], first: usize, out: &mut [f32]) {
213        #[cfg(target_arch = "x86_64")]
214        if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma")
215        {
216            unsafe { self.render_fma(padded, first, out) };
217            return;
218        }
219        self.render_plain(padded, first, out);
220    }
221
222    #[cfg(target_arch = "x86_64")]
223    #[target_feature(enable = "avx2,fma")]
224    unsafe fn render_fma(&self, padded: &[Vec<f32>], first: usize, out: &mut [f32]) {
225        self.render_plain(padded, first, out);
226    }
227
228    #[inline(always)]
229    fn render_plain(&self, padded: &[Vec<f32>], first: usize, out: &mut [f32]) {
230        for (k, frame) in out.chunks_exact_mut(self.channels).enumerate() {
231            let pos = (first + k) as u64 * self.down;
232            let (index, phase) = ((pos / self.up) as usize, pos % self.up);
233            for (o, h) in frame.iter_mut().zip(padded) {
234                *o = self.tap(phase, &h[index..index + self.taps]);
235            }
236        }
237    }
238}
239
240#[inline(always)]
241pub(crate) fn dot(a: &[f32], b: &[f32]) -> f32 {
242    let mut acc = [0.0f32; 16];
243    let mut ca = a.chunks_exact(16);
244    let mut cb = b.chunks_exact(16);
245    for (x, y) in (&mut ca).zip(&mut cb) {
246        for i in 0..16 {
247            acc[i] += x[i] * y[i];
248        }
249    }
250    let mut s: f32 = ca
251        .remainder()
252        .iter()
253        .zip(cb.remainder())
254        .map(|(x, y)| x * y)
255        .sum();
256    for v in acc {
257        s += v;
258    }
259    s
260}
261
262const BLOCK: usize = 2048;
263
264pub fn resample(input: &AudioBuffer, rate: u32) -> Result<AudioBuffer> {
265    if input.rate == rate {
266        return Ok(input.clone());
267    }
268    let r = Resampler::new(input.rate, rate, input.channels)?;
269    let ch = r.channels;
270    let frames = input.frames();
271    let half = r.taps / 2;
272    let total = (frames as u128 * r.up as u128).div_ceil(r.down as u128) as usize;
273    let padded: Vec<Vec<f32>> = (0..ch)
274        .into_par_iter()
275        .map(|c| {
276            let mut h = Vec::with_capacity(frames + 2 * half);
277            h.resize(half - 1, 0.0);
278            h.extend(input.samples.iter().skip(c).step_by(ch).take(frames));
279            h.resize(frames + 2 * half - 1, 0.0);
280            h
281        })
282        .collect();
283    let mut out = AudioBuffer::new(rate, input.channels);
284    out.samples = vec![0.0; total * ch];
285    out.samples
286        .par_chunks_mut(BLOCK * ch)
287        .enumerate()
288        .for_each(|(b, chunk)| r.render(&padded, b * BLOCK, chunk));
289    Ok(out)
290}