account for latency

8 files changed, 145 insertions, 97 deletions

diff --git a/Cargo.toml b/Cargo.toml
index 8fea9d7..6bbf43f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,8 +26,8 @@ egui-baseview = { git = "https://codeberg.org/BillyDM/egui-baseview.git", versio
 prost = "0.14.3"
 prost-build = "0.14.3"
 realfft = { version = "3.5.0", features = ["avx", "neon", "sse"] }
-hashbrown = "0.17.1"
 num-complex = "0.4.6"
+rustc-hash = "2.1.2"
 
 [build-dependencies]
 prost-build = "0.14.3"
diff --git a/src/audio.rs b/src/audio.rs
index 4858346..71be7d3 100644
--- a/src/audio.rs
+++ b/src/audio.rs
@@ -31,7 +31,8 @@ pub struct RetainPluginAudioProcessor<'a> {
     host: HostAudioProcessorHandle<'a>,
     /// Fft for the left channel
     fft_left: WindowedRealFft,
-	fft_right: WindowedRealFft,
+    /// Fft for the right channel
+    fft_right: WindowedRealFft,
 }
 
 impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread<'a>>
@@ -39,24 +40,22 @@ impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread
 {
     fn activate(
         host: HostAudioProcessorHandle<'a>,
-        main_thread: &mut RetainPluginMainThread,
+        _main_thread: &mut RetainPluginMainThread,
         shared: &'a RetainPluginShared<'a>,
         _audio_config: PluginAudioConfiguration,
     ) -> Result<Self, PluginError> {
-        if let Some(latency) = shared.host.get_extension::<HostLatency>() {
-            latency.changed(&mut main_thread.host_main_thread);
-        }
+        let params = RetainParamsLocal::new(&shared.params);
 
-        let fft_left = WindowedRealFft::new(32768);
-		let fft_right = WindowedRealFft::new(32768);
+        let fft_left = WindowedRealFft::new(params.get_window_size().into());
+        let fft_right = WindowedRealFft::new(params.get_window_size().into());
 
         // This is where we would allocate intermediate buffers and such if we needed them.
         Ok(Self {
-            params: RetainParamsLocal::new(&shared.params),
+            params,
             shared,
             host,
             fft_left,
-			fft_right,
+            fft_right,
         })
     }
 
@@ -93,9 +92,27 @@ impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread
             }
         }
 
+        let prev_window_size = self.params.get_window_size();
+
         // Receive any param updates from the main thread and/or the GUI.
         let has_ui_param_updates = self.params.fetch_updates(&self.shared.params);
 
+        // update window size and latency if it has changed
+        // updates fft window sizes only if necessary
+        let window_size = self.params.get_window_size();
+        if prev_window_size != window_size {
+            self.fft_left.window_size(window_size.into());
+            self.fft_right.window_size(window_size.into());
+
+            if let Some(latency) = self.shared.host.get_extension::<HostLatency>() {
+                // should be safe
+                let mut main = unsafe { self.shared.host.as_main_thread_unchecked() };
+
+                latency.changed(&mut main);
+                // self.shared.host.request_restart();
+            }
+        }
+
         // Now let's process the audio, while splitting the processing in batches between each
         // sample-accurate event.
         for event_batch in events.input.batch() {
@@ -107,10 +124,20 @@ impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread
             // Get the parameters after all changes have been handled.
             let order = self.params.get_order();
 
-            // process in place here
+            if order == 0 {
+                for channel in channel_buffers.iter_mut().flatten() {
+                    for sample in channel.iter_mut() {
+                        *sample = 0.0;
+                    }
+                }
+
+                continue;
+            }
+
+            // process samples in place here
             if let [Some(left), Some(right)] = &mut channel_buffers {
                 for sample in left.iter_mut() {
-                    if self.fft_left.push_front_input(*sample) {
+                    if self.fft_left.push_back_input(*sample) {
                         self.fft_left.forward();
 
                         retain_top_n_magnitudes(self.fft_left.get_spectrum(), order);
@@ -119,11 +146,11 @@ impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread
                         self.fft_left.clear_input();
                     }
 
-                    *sample = self.fft_left.pop_back_output();
+                    *sample = self.fft_left.pop_front_output();
                 }
 
                 for sample in right.iter_mut() {
-					if self.fft_right.push_front_input(*sample) {
+                    if self.fft_right.push_back_input(*sample) {
                         self.fft_right.forward();
 
                         retain_top_n_magnitudes(self.fft_right.get_spectrum(), order);
@@ -132,8 +159,8 @@ impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread
                         self.fft_right.clear_input();
                     }
 
-                    *sample = self.fft_right.pop_back_output();
-				}
+                    *sample = self.fft_right.pop_front_output();
+                }
             }
         }
 
diff --git a/src/gui.rs b/src/gui.rs
index b10705b..3d2dfb2 100644
--- a/src/gui.rs
+++ b/src/gui.rs
@@ -66,9 +66,10 @@ impl RetainPluginGui {
                 egui::CentralPanel::default().show(egui_ctx, |ui| {
                     ui.heading("Retain");
                     let mut order = state.local_params.get_order();
+                    let window_size = state.local_params.get_window_size();
 
                     let slider = ui.add(
-                        Slider::new(&mut order, 0..=100_000)
+                        Slider::new(&mut order, 0..=window_size)
                             .text("Order")
                             .logarithmic(true),
                     );
@@ -82,20 +83,18 @@ impl RetainPluginGui {
                     state.local_params.push_gesture(&state.shared_params);
 
                     let mut window_size = WindowSize::from(state.local_params.get_window_size());
-
                     ComboBox::from_label("Window Size")
                         .selected_text(window_size.as_str())
                         .show_ui(ui, |ui| {
-                            for size in WindowSize::into_iter() {
+                            for size in WindowSize::iter() {
                                 let display = size.as_str();
+                                let inner = size.inner();
 
                                 if ui
                                     .selectable_value(&mut window_size, size, display)
                                     .changed()
                                 {
-                                    let u = window_size.value();
-
-                                    state.local_params.set_window_size(u);
+                                    state.local_params.set_window_size(inner);
                                     state.local_params.push_updates(&state.shared_params);
                                 }
                             }
diff --git a/src/lib.rs b/src/lib.rs
index e14a81e..3288088 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -59,13 +59,12 @@ impl DefaultPluginFactory for RetainPlugin {
     }
 
     fn new_main_thread<'a>(
-        host: HostMainThreadHandle<'a>,
+        _host: HostMainThreadHandle<'a>,
         shared: &'a Self::Shared<'a>,
     ) -> Result<Self::MainThread<'a>, PluginError> {
         Ok(Self::MainThread {
             shared,
             params: RetainParamsLocal::new(&shared.params),
-            host_main_thread: host,
             gui: None,
         })
     }
@@ -75,15 +74,14 @@ impl DefaultPluginFactory for RetainPlugin {
 pub struct RetainPluginShared<'a> {
     /// The plugin's parameter values.
     params: Arc<RetainParamsShared>,
-    /// A handle to the host
-    host: HostSharedHandle<'a>,
+    host: Arc<HostSharedHandle<'a>>,
 }
 
 impl<'a> RetainPluginShared<'a> {
     fn new(host: HostSharedHandle<'a>) -> Self {
         RetainPluginShared {
             params: Arc::new(RetainParamsShared::new()),
-            host,
+            host: Arc::new(host),
         }
     }
 }
@@ -96,7 +94,6 @@ pub struct RetainPluginMainThread<'a> {
     params: RetainParamsLocal,
     /// A reference to the plugin's shared data.
     shared: &'a RetainPluginShared<'a>,
-    host_main_thread: HostMainThreadHandle<'a>,
     /// The plugin's GUI state and context
     gui: Option<RetainPluginGui>,
 }
@@ -111,7 +108,7 @@ impl<'a> PluginMainThread<'a, RetainPluginShared<'a>> for RetainPluginMainThread
 
 impl PluginLatencyImpl for RetainPluginMainThread<'_> {
     fn get(&mut self) -> u32 {
-        self.params.get_window_size() as u32
+        self.shared.params.get_window_size() as u32
     }
 }
 
diff --git a/src/params.rs b/src/params.rs
index c895bca..041f838 100644
--- a/src/params.rs
+++ b/src/params.rs
@@ -22,8 +22,8 @@ use std::{
 };
 
 /// The default value of the order parameter.
-const DEFAULT_ORDER: usize = 1000;
-pub const DEFAULT_WINDOW_SIZE: WindowSize = WindowSize::Size32768;
+const DEFAULT_ORDER: usize = 1;
+pub const DEFAULT_WINDOW_SIZE: WindowSize = WindowSize::Size4096;
 
 /// A struct that manages the parameters for our plugin.
 ///
@@ -48,10 +48,14 @@ impl RetainParamsShared {
     pub fn new() -> Self {
         Self {
             order: AtomicUsize::new(DEFAULT_ORDER),
-            window_size: AtomicUsize::new(DEFAULT_WINDOW_SIZE.value()),
+            window_size: AtomicUsize::new(DEFAULT_WINDOW_SIZE.into_inner()),
             has_gesture: AtomicBool::new(false),
         }
     }
+
+    pub fn get_window_size(&self) -> usize {
+        self.window_size.load(Ordering::Relaxed)
+    }
 }
 
 /// A param gesture change type
@@ -102,7 +106,7 @@ impl RetainParamsLocal {
     /// It is clamped to the range `0..=100_000`.
     #[inline]
     pub fn set_order(&mut self, value: usize) {
-        self.order = value.clamp(0, 100_000);
+        self.order = value;
     }
 
     /// Returns the current window size value.
@@ -297,7 +301,7 @@ impl PluginMainThreadParams for RetainPluginMainThread<'_> {
             name: b"Order",
             module: b"",
             min_value: 0.0,
-            max_value: 100_000.0,
+            max_value: f64::MAX,
             default_value: DEFAULT_ORDER as f64,
         });
 
@@ -309,7 +313,7 @@ impl PluginMainThreadParams for RetainPluginMainThread<'_> {
             module: b"",
             min_value: 128.0,
             max_value: usize::MAX as f64,
-            default_value: DEFAULT_WINDOW_SIZE.value() as f64,
+            default_value: DEFAULT_WINDOW_SIZE.into_inner() as f64,
         });
     }
 
diff --git a/src/retain.rs b/src/retain.rs
index bc3934a..49730e1 100644
--- a/src/retain.rs
+++ b/src/retain.rs
@@ -1,11 +1,18 @@
-use hashbrown::HashSet;
 use num_complex::Complex;
+use rustc_hash::FxHashSet;
 
 #[inline(always)]
 pub fn retain_top_n_magnitudes(fft_real_signal: &mut [Complex<f32>], n: usize) {
     // you'd be keeping the entire signal
+    // or keeping nothing at all
     if n >= fft_real_signal.len() {
         return;
+    } else if n == 0 {
+        for phasor in fft_real_signal {
+            *phasor = Complex::ZERO;
+        }
+
+        return;
     }
 
     let mut indexed = fft_real_signal
@@ -20,11 +27,11 @@ pub fn retain_top_n_magnitudes(fft_real_signal: &mut [Complex<f32>], n: usize) {
     let top_indices = indexed[target..]
         .iter()
         .map(|&(i, _)| i)
-        .collect::<HashSet<usize>>();
+        .collect::<FxHashSet<usize>>();
 
-    for (i, val) in fft_real_signal.iter_mut().enumerate() {
+    for (i, phasor) in fft_real_signal.iter_mut().enumerate() {
         if !top_indices.contains(&i) {
-            *val = Complex::ZERO;
+            *phasor = Complex::ZERO;
         }
     }
 }
diff --git a/src/window_size.rs b/src/window_size.rs
index 0c1ef08..9d3f66b 100644
--- a/src/window_size.rs
+++ b/src/window_size.rs
@@ -1,3 +1,5 @@
+/// The max value of the custom window size is u32::MAX.
+/// That's because it's the maximum latency one can report.
 #[derive(Debug, PartialEq)]
 pub enum WindowSize {
     Size128,
@@ -14,6 +16,8 @@ pub enum WindowSize {
 
 impl From<usize> for WindowSize {
     fn from(item: usize) -> Self {
+        let item = item.clamp(0, u32::MAX as usize);
+
         match item {
             128 => Self::Size128,
             256 => Self::Size256,
@@ -45,7 +49,22 @@ impl WindowSize {
         }
     }
 
-    pub fn value(&self) -> usize {
+    pub fn into_inner(self) -> usize {
+        match self {
+            Self::Size128 => 128,
+            Self::Size256 => 256,
+            Self::Size512 => 512,
+            Self::Size1024 => 1024,
+            Self::Size2048 => 2048,
+            Self::Size4096 => 4096,
+            Self::Size8192 => 8192,
+            Self::Size16384 => 16384,
+            Self::Size32768 => 32768,
+            Self::Custom(x) => x,
+        }
+    }
+
+    pub fn inner(&self) -> usize {
         match self {
             Self::Size128 => 128,
             Self::Size256 => 256,
@@ -60,7 +79,7 @@ impl WindowSize {
         }
     }
 
-    pub fn into_iter() -> impl Iterator<Item = Self> {
+    pub fn iter() -> impl Iterator<Item = Self> {
         [
             Self::Size128,
             Self::Size256,
diff --git a/src/windowed_fft.rs b/src/windowed_fft.rs
index 6d11564..709a68b 100644
--- a/src/windowed_fft.rs
+++ b/src/windowed_fft.rs
@@ -4,13 +4,16 @@
 /*use crate::window::{
     BlackmanHarrisWindow, HammingWindow, HannWindow, RectangularWindow, WindowFunction,
 };*/
-use crate::window_function::{RectangularWindow, WindowFunction};
+use crate::{
+    window_function::{RectangularWindow, WindowFunction},
+    window_size::WindowSize,
+};
 use num_complex::Complex;
 use realfft::{ComplexToReal, RealFftPlanner, RealToComplex};
 use std::{collections::VecDeque, sync::Arc};
 
 pub struct WindowedRealFft {
-    fft_size: usize,
+    window_size: usize,
     planner: RealFftPlanner<f32>,
     forward: Arc<dyn RealToComplex<f32>>,
     inverse: Arc<dyn ComplexToReal<f32>>,
@@ -22,27 +25,25 @@ pub struct WindowedRealFft {
 }
 
 impl WindowedRealFft {
-    pub fn new(mut fft_size: usize) -> Self {
-		if fft_size == 0 {
-			fft_size = 1;
-		}
-		
+    pub fn new(window_size: WindowSize) -> Self {
+        let window_size = window_size.into_inner();
+
         let mut planner = RealFftPlanner::new();
-        let forward = planner.plan_fft_forward(fft_size);
-        let inverse = planner.plan_fft_inverse(fft_size);
+        let forward = planner.plan_fft_forward(window_size);
+        let inverse = planner.plan_fft_inverse(window_size);
 
-        let window_function = Box::new(RectangularWindow::new(fft_size));
+        let window_function = Box::new(RectangularWindow::new(window_size));
 
-        let input = VecDeque::with_capacity(fft_size);
-        let output = VecDeque::with_capacity(fft_size);
+        let input = VecDeque::with_capacity(window_size);
+        let output = VecDeque::with_capacity(window_size);
 
-        let spectrum = vec![Complex::ZERO; (fft_size / 2) + 1];
+        let spectrum = vec![Complex::ZERO; (window_size / 2) + 1];
 
-        let max_scratch_len = forward.get_scratch_len().max(inverse.get_scratch_len());
-        let scratch = vec![Complex::ZERO; max_scratch_len];
+        // scratches should be the same length for both left and right channels
+        let scratch = forward.make_scratch_vec();
 
         Self {
-            fft_size,
+            window_size,
             planner,
             forward,
             inverse,
@@ -54,45 +55,41 @@ impl WindowedRealFft {
         }
     }
 
-    pub fn fft_size(mut self, mut value: usize) -> Self {
-		if value == 0 {
-			value = 1;
-		}
-		
-        self.fft_size = value;
+    pub fn window_size(&mut self, window_size: WindowSize) {
+        if window_size == self.window_size.into() {
+            return;
+        }
+
+        self.window_size = window_size.into_inner();
 
-        self.forward = self.planner.plan_fft_forward(self.fft_size);
-        self.inverse = self.planner.plan_fft_inverse(self.fft_size);
+        self.forward = self.planner.plan_fft_forward(self.window_size);
+        self.inverse = self.planner.plan_fft_inverse(self.window_size);
 
-        self.input.reserve_exact(self.fft_size);
-        self.output.reserve_exact(self.fft_size);
+        self.input.reserve_exact(self.window_size);
+        self.output.reserve_exact(self.window_size);
 
         self.input.clear();
         self.output.clear();
 
-        self.spectrum.resize((self.fft_size / 2) + 1, Complex::ZERO);
+        self.spectrum
+            .resize((self.window_size / 2) + 1, Complex::ZERO);
 
-        let max_scratch_len = self
-            .forward
-            .get_scratch_len()
-            .max(self.inverse.get_scratch_len());
-        self.scratch.resize(max_scratch_len, Complex::ZERO);
-
-        self
+        self.scratch
+            .resize(self.forward.get_scratch_len(), Complex::ZERO);
     }
 
     pub fn clear_input(&mut self) {
         self.input.clear();
     }
 
-    pub fn push_front_input(&mut self, value: f32) -> bool {
-        self.input.push_front(value);
+    pub fn push_back_input(&mut self, value: f32) -> bool {
+        self.input.push_back(value);
 
-        self.input.len() >= self.fft_size
+        self.input.len() >= self.window_size
     }
 
-    pub fn pop_back_output(&mut self) -> f32 {
-        self.output.pop_back().unwrap_or(0.0)
+    pub fn pop_front_output(&mut self) -> f32 {
+        self.output.pop_front().unwrap_or(0.0)
     }
 
     pub fn get_spectrum(&mut self) -> &mut [Complex<f32>] {
@@ -102,27 +99,25 @@ impl WindowedRealFft {
     pub fn forward(&mut self) {
         self.window_function.apply(self.input.make_contiguous());
 
-        self.forward
-            .process_with_scratch(
-                self.input.make_contiguous(),
-                &mut self.spectrum,
-                &mut self.scratch,
-            );
+        let _ = self.forward.process_with_scratch(
+            self.input.make_contiguous(),
+            &mut self.spectrum,
+            &mut self.scratch,
+        );
     }
 
     pub fn inverse(&mut self) {
-		self.output.resize(self.fft_size, 0.0);
-		
-        self.inverse
-            .process_with_scratch(
-                &mut self.spectrum,
-                self.output.make_contiguous(),
-                &mut self.scratch,
-            );
-
-        let fft_size_f32 = self.fft_size as f32;
+        self.output.resize(self.window_size, 0.0);
+
+        let _ = self.inverse.process_with_scratch(
+            &mut self.spectrum,
+            self.output.make_contiguous(),
+            &mut self.scratch,
+        );
+
+        let window_size_f32 = self.window_size as f32;
         for sample in &mut self.output {
-            *sample = *sample / fft_size_f32;
+            *sample = *sample / window_size_f32;
         }
 
         self.window_function.reverse(self.output.make_contiguous());