Add a dry-wet mixer to Spectral Compressor

plugins/spectral_compressor/src/dry_wet_mixer.rs

@@ -0,0 +1,168 @@
+// Spectral Compressor: an FFT based compressor
+// Copyright (C) 2021-2022 Robbert van der Helm
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU General Public License for more details.
+//
+// You should have received a copy of the GNU General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+use nih_plug::prelude::Buffer;
+
+/// A simple dry-wet mixer with latency compensation that operates on entire buffers.
+pub struct DryWetMixer {
+    /// The delay line for the latency compensation. This is indexed by `[channel_idx][sample_idx]`,
+    /// with the size set to the maximum latency plus the maximum block size rounded up to the next
+    /// power of two.
+    delay_line: Vec<Vec<f32>>,
+    /// The position in the inner delay line buffer where the next samples should be written from.
+    /// This is incremented after writing. When reading the data for mixing the dry signal back in,
+    /// the starting read position is determined by subtracting the buffer's length from this
+    /// position and then subtracting the latency.
+    next_write_position: usize,
+}
+
+/// The mixing style for the [`DryWetMixer`].
+#[derive(Debug, Clone, Copy)]
+#[allow(unused)]
+pub enum MixingStyle {
+    Linear,
+    EqualPower,
+}
+
+impl DryWetMixer {
+    /// Set up the mixer for the given parameters.
+    pub fn new(num_channels: usize, max_block_size: usize, max_latency: usize) -> Self {
+        // TODO: This could be more efficient if we don't use the entire buffer when the actual
+        //       latency is lower than the maximum latency, but that's an optimization for later
+        let delay_line_len = (max_block_size + max_latency).next_power_of_two();
+
+        DryWetMixer {
+            delay_line: vec![vec![0.0; delay_line_len]; num_channels],
+            next_write_position: 0,
+        }
+    }
+
+    /// Resize the itnernal buffers to fit new parameters.
+    pub fn resize(&mut self, num_channels: usize, max_block_size: usize, max_latency: usize) {
+        let delay_line_len = (max_block_size + max_latency).next_power_of_two();
+
+        self.delay_line.resize_with(num_channels, Vec::new);
+        for buffer in &mut self.delay_line {
+            buffer.resize(delay_line_len, 0.0);
+            buffer.fill(0.0);
+        }
+        self.next_write_position = 0;
+    }
+
+    /// Clear out the buffers.
+    pub fn reset(&mut self) {
+        for buffer in &mut self.delay_line {
+            buffer.fill(0.0);
+        }
+        self.next_write_position = 0;
+    }
+
+    /// Write the dry signal into the buffer. This should be called at the start of the process
+    /// function.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the buffer is larger than the maximum block size or if the channel counts don't
+    /// match.
+    pub fn write_dry(&mut self, buffer: &Buffer) {
+        if buffer.channels() == 0 {
+            return;
+        }
+
+        assert_eq!(buffer.channels(), self.delay_line.len());
+        let delay_line_len = self.delay_line[0].len();
+        assert!(buffer.len() <= delay_line_len);
+
+        let num_samples_before_wrap = buffer.len().min(delay_line_len - self.next_write_position);
+        let num_samples_after_wrap = buffer.len() - num_samples_before_wrap;
+
+        for (buffer_channel, delay_line) in buffer
+            .as_slice_immutable()
+            .iter()
+            .zip(self.delay_line.iter_mut())
+        {
+            delay_line
+                [self.next_write_position..self.next_write_position + num_samples_before_wrap]
+                .copy_from_slice(&buffer_channel[..num_samples_before_wrap]);
+            delay_line[..num_samples_after_wrap]
+                .copy_from_slice(&buffer_channel[num_samples_before_wrap..]);
+        }
+
+        self.next_write_position = (self.next_write_position + buffer.len()) % delay_line_len;
+    }
+
+    /// Mix the dry signal into the buffer. The ratio is a `[0, 1]` integer where 0 results in an
+    /// all-dry signal, and 1 results in an all-wet signal. This should be called at the start of
+    /// the process function.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the buffer is larger than the maximum block size, if the latency is larger than
+    /// the maximum latency, or if the channel counts don't match.
+    pub fn mix_in_dry(
+        &mut self,
+        buffer: &mut Buffer,
+        ratio: f32,
+        style: MixingStyle,
+        latency: usize,
+    ) {
+        if buffer.channels() == 0 {
+            return;
+        }
+
+        let ratio = ratio.clamp(0.0, 1.0);
+        if ratio == 1.0 {
+            return;
+        }
+        let (wet_t, dry_t) = match style {
+            MixingStyle::Linear => (ratio, 1.0 - ratio),
+            MixingStyle::EqualPower => (ratio.sqrt(), (1.0 - ratio).sqrt()),
+        };
+
+        assert_eq!(buffer.channels(), self.delay_line.len());
+        let delay_line_len = self.delay_line[0].len();
+        assert!(buffer.len() + latency <= delay_line_len);
+
+        let read_position =
+            (self.next_write_position + delay_line_len - buffer.len() - latency) % delay_line_len;
+        let num_samples_before_wrap = buffer.len().min(delay_line_len - read_position);
+        let num_samples_after_wrap = buffer.len() - num_samples_before_wrap;
+
+        for (buffer_channel, delay_line) in buffer.as_slice().iter_mut().zip(self.delay_line.iter())
+        {
+            if ratio == 0.0 {
+                buffer_channel[..num_samples_before_wrap].copy_from_slice(
+                    &delay_line[read_position..read_position + num_samples_before_wrap],
+                );
+                buffer_channel[num_samples_before_wrap..]
+                    .copy_from_slice(&delay_line[..num_samples_after_wrap]);
+            } else {
+                for (buffer_sample, delay_sample) in buffer_channel[..num_samples_before_wrap]
+                    .iter_mut()
+                    .zip(&delay_line[read_position..read_position + num_samples_before_wrap])
+                {
+                    *buffer_sample = (*buffer_sample * wet_t) + (delay_sample * dry_t);
+                }
+                for (buffer_sample, delay_sample) in buffer_channel[num_samples_before_wrap..]
+                    .iter_mut()
+                    .zip(&delay_line[..num_samples_after_wrap])
+                {
+                    *buffer_sample = (*buffer_sample * wet_t) + (delay_sample * dry_t);
+                }
+            }
+        }
+    }
+}

plugins/spectral_compressor/src/lib.rs

@@ -20,6 +20,7 @@ use realfft::num_complex::Complex32;
 use realfft::{ComplexToReal, RealFftPlanner, RealToComplex};
 use std::sync::Arc;
 
+mod dry_wet_mixer;
 mod editor;
 
 const MIN_WINDOW_ORDER: usize = 6;
@@ -51,6 +52,8 @@ struct SpectralCompressor {
     /// Contains a Hann window function of the current window length, passed to the overlap-add
     /// helper. Allocated with a `MAX_WINDOW_SIZE` initial capacity.
     window_function: Vec<f32>,
+    /// A mixer to mix the dry signal back into the processed signal with latency compensation.
+    dry_wet_mixer: dry_wet_mixer::DryWetMixer,
 
     /// The algorithms for the FFT and IFFT operations, for each supported order so we can switch
     /// between them without replanning or allocations. Initialized during `initialize()`.
@@ -96,9 +99,10 @@ impl Default for SpectralCompressor {
             params: Arc::new(SpectralCompressorParams::default()),
             editor_state: editor::default_state(),
 
-            // These two will be set to the correct values in the initialize function
+            // These three will be set to the correct values in the initialize function
             stft: util::StftHelper::new(Self::DEFAULT_NUM_OUTPUTS as usize, MAX_WINDOW_SIZE, 0),
             window_function: Vec::with_capacity(MAX_WINDOW_SIZE),
+            dry_wet_mixer: dry_wet_mixer::DryWetMixer::new(0, 0, 0),
 
             // This is initialized later since we don't want to do non-trivial computations before
             // the plugin is initialized
@@ -136,6 +140,7 @@ impl Default for SpectralCompressorParams {
             auto_makeup_gain: BoolParam::new("Auto Makeup Gain", true),
             dry_wet_ratio: FloatParam::new("Mix", 1.0, FloatRange::Linear { min: 0.0, max: 1.0 })
                 .with_unit("%")
+                .with_smoother(SmoothingStyle::Linear(15.0))
                 .with_value_to_string(formatters::v2s_f32_percentage(0))
                 .with_string_to_value(formatters::s2v_f32_percentage()),
             dc_filter: BoolParam::new("DC Filter", true),
@@ -172,7 +177,7 @@ impl Plugin for SpectralCompressor {
     fn initialize(
         &mut self,
         bus_config: &BusConfig,
-        _buffer_config: &BufferConfig,
+        buffer_config: &BufferConfig,
         context: &mut impl InitContext,
     ) -> bool {
         // This plugin can accept any number of channels, so we need to resize channel-dependent
@@ -181,6 +186,12 @@ impl Plugin for SpectralCompressor {
             self.stft = util::StftHelper::new(self.stft.num_channels(), MAX_WINDOW_SIZE, 0);
         }
 
+        self.dry_wet_mixer.resize(
+            bus_config.num_output_channels as usize,
+            buffer_config.max_buffer_size as usize,
+            MAX_WINDOW_SIZE,
+        );
+
         // Planning with RustFFT is very fast, but it will still allocate we we'll plan all of the
         // FFTs we might need in advance
         if self.plan_for_order.is_none() {
@@ -206,6 +217,10 @@ impl Plugin for SpectralCompressor {
         true
     }
 
+    fn reset(&mut self) {
+        self.dry_wet_mixer.reset();
+    }
+
     fn process(
         &mut self,
         buffer: &mut Buffer,
@@ -244,7 +259,10 @@ impl Plugin for SpectralCompressor {
             util::db_to_gain(self.params.input_gain_db.value) * gain_compensation.sqrt();
         let output_gain =
             util::db_to_gain(self.params.output_gain_db.value) * gain_compensation.sqrt();
-        // TODO: Mix in the dry signal
+        // TODO: Auto makeup gain
+
+        // This is mixed in later with latency compensation applied
+        self.dry_wet_mixer.write_dry(buffer);
 
         self.stft
             .process_overlap_add(buffer, overlap_times, |_channel_idx, real_fft_buffer| {
@@ -291,6 +309,17 @@ impl Plugin for SpectralCompressor {
                 }
             });
 
+        self.dry_wet_mixer.mix_in_dry(
+            buffer,
+            self.params
+                .dry_wet_ratio
+                .smoothed
+                .next_step(buffer.len() as u32),
+            // The dry and wet signals are in phase, so we can do a linear mix
+            dry_wet_mixer::MixingStyle::Linear,
+            self.stft.latency_samples() as usize,
+        );
+
         ProcessStatus::Normal
     }
 }