Skip to main content

servo_media_audio/
convolver_node.rs

1/* This Source Code Form is subject to the terms of the Mozilla Public
2 * License, v. 2.0. If a copy of the MPL was not distributed with this
3 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
4
5use 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
37// <https://webaudio.github.io/web-audio-api/#dom-convolvernode-normalize>
38fn 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    // Normalize by RMS power.
43    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    // Calibrate to make perceived volume same as unprocessed.
60    scale *= gain_calibration;
61    // Scale depends on sample-rate.
62    scale *= gain_calibration_sample_rate / buffer.sample_rate as f64;
63    // True-stereo compensation.
64    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 we have a mono IR, need to create two convolvers.
85    // This is to ensure we can handle the stereo input case, since the FFT Convolver
86    // assumes inputs are blocks of a long-running sample.
87    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
97/// Calculates the convolution between the input and the impulse response of the convolver.
98/// FFTConvolver uses the overlap-add algorithm to calculate the convolution.
99// TODO: Switch to ThreadedFFTConvolver when this feature is stable.
100fn 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                // <https://webaudio.github.io/web-audio-api/#dom-convolvernode-normalize>
158                // > Changes to this value do not take effect until the next time the buffer attribute is set.
159                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                // Precompute buffer FFTs.
169                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        // <https://webaudio.github.io/web-audio-api/#Convolution-channel-configurations>
200        let convolution_output = match (input_block.chan_count(), buffer.chans()) {
201            // If we have a mono impulse response, we need to calculate convolution for the
202            // second convolver, which in the mono case uses the same impulse response as the first
203            // convolver.
204            // This is because fast convolution algorithms such as overlap-add used by the FFT Convolver
205            // assume inputs are blocks of a long-running sample.
206            // Results from previous blocks are used when calculating for subsequent blocks.
207            // If the input increases channel count, the convolver will need the results from
208            // previous blocks as if it had always had the upmixed input.
209            //
210            // Mono with Mono Response
211            (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            // Stereo with Mono Response
221            (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            // Mono with Stereo Response
231            (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            // Stereo with Stereo Response
241            (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            // Stereo with "true" Stereo Matrix Response
251            (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            // Mono with Stereo Matrix Response
263            (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}