servo_media_audio/
convolver_node.rs1use std::sync::Arc;
6
7use fft_convolver::FFTConvolver;
8use log::error;
9use malloc_size_of_derive::MallocSizeOf;
10
11use crate::audio_node::{
12 AudioNodeEngine, AudioNodeType, BlockInfo, ChannelInfo, ChannelInterpretation,
13};
14use crate::block::{Block, Chunk, FRAMES_PER_BLOCK_USIZE};
15use crate::buffer_source_node::AudioBuffer;
16
17#[derive(Clone, Debug, MallocSizeOf)]
18pub struct ConvolverNodeOptions {
19 pub buffer: Option<AudioBuffer>,
20 pub normalize: bool,
21}
22
23#[derive(Clone, Debug, MallocSizeOf)]
24pub enum ConvolverNodeMessage {
25 SetBuffer(#[conditional_malloc_size_of] Option<Arc<AudioBuffer>>),
26 SetNormalize(bool),
27}
28
29#[derive(AudioNodeCommon)]
30pub(crate) struct ConvolverNode {
31 channel_info: ChannelInfo,
32 buffer: Option<Arc<AudioBuffer>>,
33 normalize: bool,
34 convolvers: Option<Vec<FFTConvolver<f32>>>,
35}
36
37fn calculate_normalization_scale(buffer: &AudioBuffer) -> f64 {
39 let gain_calibration = 0.00125_f64;
40 let gain_calibration_sample_rate = 44100_f64;
41 let min_power = 0.000125_f64;
42 let number_of_channels = buffer.chans() as f64;
44 let buffer_length = buffer.len() as f64;
45
46 let mut power = buffer.buffers.iter().fold(0_f64, |power, channel| {
47 power +
48 channel.iter().fold(0_f64, |channel_power, sample| {
49 channel_power + sample.powi(2) as f64
50 })
51 });
52
53 power = (power / (number_of_channels * buffer_length)).sqrt();
54 if power.is_infinite() {
55 power = min_power;
56 }
57 power = power.max(min_power);
58 let mut scale = 1.0 / power;
59 scale *= gain_calibration;
61 scale *= gain_calibration_sample_rate / buffer.sample_rate as f64;
63 if number_of_channels == 4.0 {
65 scale *= 0.5;
66 }
67 scale
68}
69
70fn initialize_convolvers(buffer: Arc<AudioBuffer>) -> Option<Vec<FFTConvolver<f32>>> {
71 let initialize_convolver = |impulse_response: &Vec<f32>| {
72 let mut convolver = FFTConvolver::<f32>::default();
73 convolver
74 .init(FRAMES_PER_BLOCK_USIZE * 8, impulse_response.as_slice())
75 .inspect_err(|error| error!("Failed to initialize convolver {}", error))
76 .ok()?;
77 Some(convolver)
78 };
79 let mut convolvers = buffer
80 .buffers
81 .iter()
82 .map(initialize_convolver)
83 .collect::<Option<Vec<_>>>();
84 if let Some(convolvers) = &mut convolvers &&
88 convolvers.len() == 1
89 {
90 let impulse_response = buffer.buffers.first()?;
91 let convolver = initialize_convolver(impulse_response)?;
92 convolvers.push(convolver);
93 }
94 convolvers
95}
96
97fn linear_convolution(input: &[f32], convolver: &mut FFTConvolver<f32>) -> Option<Vec<f32>> {
101 let mut output = vec![0.0; FRAMES_PER_BLOCK_USIZE];
102 convolver
103 .process(input, output.as_mut_slice())
104 .inspect_err(|e| {
105 error!(
106 "Linear convolution of input with impulse response failed {}",
107 e
108 )
109 })
110 .ok()?;
111 Some(output)
112}
113
114fn downmix_four_channel_output_to_stereo(outputs: Vec<Vec<f32>>) -> Vec<f32> {
115 assert_eq!(outputs.len(), 4, "Output does not have four channels.");
116 let mut output = Vec::with_capacity(outputs[0].len() * 2);
117 let mixed_output_0 = outputs[0]
118 .iter()
119 .zip(outputs[2].iter())
120 .map(|(x, y)| x + y)
121 .collect::<Vec<_>>();
122 let mixed_output_1 = outputs[1]
123 .iter()
124 .zip(outputs[3].iter())
125 .map(|(x, y)| x + y)
126 .collect::<Vec<_>>();
127 output.extend(mixed_output_0);
128 output.extend(mixed_output_1);
129 output
130}
131
132impl ConvolverNode {
133 pub fn new(options: ConvolverNodeOptions, channel_info: ChannelInfo) -> Self {
134 let buffer = options.buffer.map(|mut buffer| {
135 if options.normalize {
136 let normalization_scale = calculate_normalization_scale(&buffer);
137 buffer.scale(normalization_scale);
138 }
139 Arc::new(buffer)
140 });
141
142 let convolvers = buffer
143 .as_ref()
144 .and_then(|buffer| initialize_convolvers(buffer.clone()));
145 Self {
146 channel_info,
147 buffer,
148 normalize: options.normalize,
149 convolvers,
150 }
151 }
152
153 fn handle_convolver_message(&mut self, message: ConvolverNodeMessage, _sample_rate: f32) {
154 match message {
155 ConvolverNodeMessage::SetBuffer(maybe_buffer) => {
156 let mut buffer = maybe_buffer;
157 if let Some(input_buffer) = buffer.as_ref() &&
160 self.normalize
161 {
162 let mut normalized_buffer = (**input_buffer).clone();
163 let normalization_scale = calculate_normalization_scale(&normalized_buffer);
164 normalized_buffer.scale(normalization_scale);
165 buffer = Some(Arc::new(normalized_buffer));
166 }
167 self.buffer = buffer;
168 self.convolvers = self
170 .buffer
171 .as_ref()
172 .and_then(|buffer| initialize_convolvers(buffer.clone()));
173 },
174 ConvolverNodeMessage::SetNormalize(normalize) => {
175 self.normalize = normalize;
176 },
177 }
178 }
179}
180
181impl AudioNodeEngine for ConvolverNode {
182 fn node_type(&self) -> AudioNodeType {
183 AudioNodeType::ConvolverNode
184 }
185
186 fn process(&mut self, inputs: Chunk, _info: &BlockInfo) -> Chunk {
187 debug_assert!(inputs.len() == 1);
188
189 let Some(buffer) = self.buffer.as_ref() else {
190 return Chunk::explicit_silence();
191 };
192
193 let Some(convolvers) = &mut self.convolvers else {
194 return inputs;
195 };
196
197 let input_block = &inputs.blocks[0];
198
199 let convolution_output = match (input_block.chan_count(), buffer.chans()) {
201 (1, 1) => {
212 let output = linear_convolution(input_block.data_chan(0), &mut convolvers[0]);
213 let input = match self.channel_info.interpretation {
214 ChannelInterpretation::Discrete => &[0.; FRAMES_PER_BLOCK_USIZE],
215 ChannelInterpretation::Speakers => input_block.data_chan(0),
216 };
217 let _ = linear_convolution(input, &mut convolvers[1]);
218 output
219 },
220 (2, 1) => {
222 let outputs = vec![
223 linear_convolution(input_block.data_chan(0), &mut convolvers[0]),
224 linear_convolution(input_block.data_chan(1), &mut convolvers[1]),
225 ]
226 .into_iter()
227 .collect::<Option<Vec<Vec<_>>>>();
228 outputs.map(|outputs| outputs.into_iter().flatten().collect())
229 },
230 (1, 2) => {
232 let outputs = vec![
233 linear_convolution(input_block.data_chan(0), &mut convolvers[0]),
234 linear_convolution(input_block.data_chan(0), &mut convolvers[1]),
235 ]
236 .into_iter()
237 .collect::<Option<Vec<Vec<_>>>>();
238 outputs.map(|outputs| outputs.into_iter().flatten().collect())
239 },
240 (2, 2) => {
242 let outputs = vec![
243 linear_convolution(input_block.data_chan(0), &mut convolvers[0]),
244 linear_convolution(input_block.data_chan(1), &mut convolvers[1]),
245 ]
246 .into_iter()
247 .collect::<Option<Vec<Vec<_>>>>();
248 outputs.map(|outputs| outputs.into_iter().flatten().collect())
249 },
250 (2, 4) => {
252 let outputs = vec![
253 linear_convolution(input_block.data_chan(0), &mut convolvers[0]),
254 linear_convolution(input_block.data_chan(0), &mut convolvers[1]),
255 linear_convolution(input_block.data_chan(1), &mut convolvers[2]),
256 linear_convolution(input_block.data_chan(1), &mut convolvers[3]),
257 ]
258 .into_iter()
259 .collect::<Option<Vec<Vec<_>>>>();
260 outputs.map(downmix_four_channel_output_to_stereo)
261 },
262 (1, 4) => {
264 let outputs = convolvers
265 .iter_mut()
266 .map(|convolver| linear_convolution(input_block.data_chan(0), convolver))
267 .collect::<Option<Vec<Vec<_>>>>();
268 outputs.map(downmix_four_channel_output_to_stereo)
269 },
270 _ => {
271 error!(
272 "Invalid channel configuration. Input channels: {}, Impulse Response channels: {}",
273 input_block.chan_count(),
274 buffer.chans()
275 );
276 return inputs;
277 },
278 };
279
280 let Some(output) = convolution_output else {
281 error!("Failed to calculate convolution.");
282 return Chunk::explicit_silence();
283 };
284 let block = Block::for_vec(output);
285 let mut chunk = Chunk::default();
286 chunk.blocks.push(block);
287 chunk
288 }
289 make_message_handler!(
290 ConvolverNode: handle_convolver_message
291 );
292}