Skip to main content

chromaprint/audio/
processor.rs

1use crate::audio::resample::AvResampleContext;
2use crate::error::{Error, Result};
3
4const MIN_SAMPLE_RATE: u32 = 1000;
5const MAX_BUFFER_SIZE: usize = 1024 * 32;
6
7// Resampler configuration matching the C implementation
8const RESAMPLE_FILTER_LENGTH: i32 = 16;
9const RESAMPLE_PHASE_SHIFT: i32 = 8;
10const RESAMPLE_LINEAR: bool = false;
11const RESAMPLE_CUTOFF: f64 = 0.8;
12
13/// Audio preprocessor that handles channel mixing to mono and resampling.
14pub struct AudioProcessor {
15    target_sample_rate: u32,
16    num_channels: u16,
17    buffer: Vec<i16>,
18    buffer_offset: usize,
19    resample_buffer: Vec<i16>,
20    resample_ctx: Option<AvResampleContext>,
21}
22
23impl AudioProcessor {
24    pub fn new(target_sample_rate: u32) -> Self {
25        Self {
26            target_sample_rate,
27            num_channels: 0,
28            buffer: vec![0i16; MAX_BUFFER_SIZE],
29            buffer_offset: 0,
30            resample_buffer: vec![0i16; MAX_BUFFER_SIZE],
31            resample_ctx: None,
32        }
33    }
34
35    /// Reset the processor for a new audio stream.
36    pub fn reset(&mut self, sample_rate: u32, num_channels: u16) -> Result<()> {
37        if num_channels == 0 {
38            return Err(Error::InvalidChannelCount(num_channels));
39        }
40        if sample_rate <= MIN_SAMPLE_RATE {
41            return Err(Error::InvalidSampleRate(sample_rate));
42        }
43
44        self.num_channels = num_channels;
45        self.buffer_offset = 0;
46
47        if sample_rate != self.target_sample_rate {
48            self.resample_ctx = Some(AvResampleContext::new(
49                self.target_sample_rate as i32,
50                sample_rate as i32,
51                RESAMPLE_FILTER_LENGTH,
52                RESAMPLE_PHASE_SHIFT,
53                RESAMPLE_LINEAR,
54                RESAMPLE_CUTOFF,
55            ));
56        } else {
57            self.resample_ctx = None;
58        }
59
60        Ok(())
61    }
62
63    /// Process interleaved multi-channel i16 samples.
64    /// Mixes to mono, optionally resamples, and calls the callback with mono samples.
65    pub fn consume<F>(&mut self, input: &[i16], mut callback: F)
66    where
67        F: FnMut(&[i16]),
68    {
69        let num_channels = self.num_channels as usize;
70        debug_assert!(input.len() % num_channels == 0);
71        let num_frames = input.len() / num_channels;
72
73        let mut pos = 0;
74        let mut remaining = num_frames;
75
76        while remaining > 0 {
77            let space = self.buffer.len() - self.buffer_offset;
78            let to_load = remaining.min(space);
79
80            match num_channels {
81                1 => {
82                    self.buffer[self.buffer_offset..self.buffer_offset + to_load]
83                        .copy_from_slice(&input[pos..pos + to_load]);
84                }
85                2 => {
86                    for i in 0..to_load {
87                        let idx = (pos + i) * 2;
88                        self.buffer[self.buffer_offset + i] =
89                            ((input[idx] as i32 + input[idx + 1] as i32) / 2) as i16;
90                    }
91                }
92                _ => {
93                    for i in 0..to_load {
94                        let idx = (pos + i) * num_channels;
95                        let mut sum: i32 = 0;
96                        for ch in 0..num_channels {
97                            sum += input[idx + ch] as i32;
98                        }
99                        self.buffer[self.buffer_offset + i] = (sum / num_channels as i32) as i16;
100                    }
101                }
102            }
103
104            self.buffer_offset += to_load;
105            pos += to_load;
106            remaining -= to_load;
107
108            if self.buffer_offset == self.buffer.len() {
109                self.resample_and_emit(&mut callback);
110                if self.buffer_offset == self.buffer.len() {
111                    // Resample didn't consume anything — shouldn't happen, but prevent infinite loop
112                    return;
113                }
114            }
115        }
116    }
117
118    /// Flush any remaining buffered samples.
119    pub fn flush<F>(&mut self, mut callback: F)
120    where
121        F: FnMut(&[i16]),
122    {
123        if self.buffer_offset > 0 {
124            self.resample_and_emit(&mut callback);
125        }
126    }
127
128    fn resample_and_emit<F>(&mut self, callback: &mut F)
129    where
130        F: FnMut(&[i16]),
131    {
132        if let Some(ref mut ctx) = self.resample_ctx {
133            let mut consumed: usize = 0;
134            let length = ctx.resample(
135                &mut self.resample_buffer,
136                &self.buffer[..self.buffer_offset],
137                &mut consumed,
138                MAX_BUFFER_SIZE,
139            );
140            if length > MAX_BUFFER_SIZE {
141                // Shouldn't happen, but match C behavior of clamping
142                callback(&self.resample_buffer[..MAX_BUFFER_SIZE]);
143            } else {
144                callback(&self.resample_buffer[..length]);
145            }
146            let remaining = self.buffer_offset as isize - consumed as isize;
147            if remaining > 0 {
148                self.buffer.copy_within(consumed..self.buffer_offset, 0);
149                self.buffer_offset = remaining as usize;
150            } else {
151                self.buffer_offset = 0;
152            }
153        } else {
154            callback(&self.buffer[..self.buffer_offset]);
155            self.buffer_offset = 0;
156        }
157    }
158}