real time processing

8 files changed, 192 insertions, 77 deletions

diff --git a/Cargo.toml b/Cargo.toml
index ed301af..8fea9d7 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -26,6 +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"
 
 [build-dependencies]
 prost-build = "0.14.3"
diff --git a/src/audio.rs b/src/audio.rs
index 495cc29..4858346 100644
--- a/src/audio.rs
+++ b/src/audio.rs
@@ -3,8 +3,16 @@
 use crate::{
     RetainPluginMainThread, RetainPluginShared,
     params::{GestureChange, RetainParamsLocal, RetainParamsShared},
+    retain::retain_top_n_magnitudes,
+    windowed_fft::WindowedRealFft,
+};
+use clack_extensions::{
+    audio_ports::{
+        AudioPortFlags, AudioPortInfo, AudioPortInfoWriter, AudioPortType, PluginAudioPortsImpl,
+    },
+    latency::HostLatency,
+    params::PluginAudioProcessorParams,
 };
-use clack_extensions::{audio_ports::*, latency::HostLatency, params::PluginAudioProcessorParams};
 use clack_plugin::{
     events::event_types::{ParamGestureBeginEvent, ParamGestureEndEvent},
     prelude::*,
@@ -21,6 +29,9 @@ pub struct RetainPluginAudioProcessor<'a> {
     shared: &'a RetainPluginShared<'a>,
     /// Our handle to the host
     host: HostAudioProcessorHandle<'a>,
+    /// Fft for the left channel
+    fft_left: WindowedRealFft,
+	fft_right: WindowedRealFft,
 }
 
 impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread<'a>>
@@ -36,11 +47,16 @@ impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread
             latency.changed(&mut main_thread.host_main_thread);
         }
 
+        let fft_left = WindowedRealFft::new(32768);
+		let fft_right = WindowedRealFft::new(32768);
+
         // This is where we would allocate intermediate buffers and such if we needed them.
         Ok(Self {
-            host,
-            shared,
             params: RetainParamsLocal::new(&shared.params),
+            shared,
+            host,
+            fft_left,
+			fft_right,
         })
     }
 
@@ -89,14 +105,35 @@ impl<'a> PluginAudioProcessor<'a, RetainPluginShared<'a>, RetainPluginMainThread
             }
 
             // Get the parameters after all changes have been handled.
-            let _order = self.params.get_order();
-            let _window_size = self.params.get_window_size();
+            let order = self.params.get_order();
 
             // process in place here
             if let [Some(left), Some(right)] = &mut channel_buffers {
-                for _ in left.iter_mut() {}
+                for sample in left.iter_mut() {
+                    if self.fft_left.push_front_input(*sample) {
+                        self.fft_left.forward();
+
+                        retain_top_n_magnitudes(self.fft_left.get_spectrum(), order);
+
+                        self.fft_left.inverse();
+                        self.fft_left.clear_input();
+                    }
+
+                    *sample = self.fft_left.pop_back_output();
+                }
+
+                for sample in right.iter_mut() {
+					if self.fft_right.push_front_input(*sample) {
+                        self.fft_right.forward();
+
+                        retain_top_n_magnitudes(self.fft_right.get_spectrum(), order);
+
+                        self.fft_right.inverse();
+                        self.fft_right.clear_input();
+                    }
 
-                for _ in right.iter_mut() {}
+                    *sample = self.fft_right.pop_back_output();
+				}
             }
         }
 
diff --git a/src/gui.rs b/src/gui.rs
index 310fbef..b10705b 100644
--- a/src/gui.rs
+++ b/src/gui.rs
@@ -6,7 +6,7 @@ use crate::{
     window_size::WindowSize,
 };
 use baseview::{Size, WindowHandle, WindowOpenOptions, WindowScalePolicy, gl::GlConfig};
-use clack_extensions::gui::*;
+use clack_extensions::gui::{GuiApiType, GuiConfiguration, GuiSize, PluginGuiImpl, Window};
 use clack_plugin::prelude::*;
 use egui_baseview::{
     EguiWindow, GraphicsConfig, Queue,
diff --git a/src/lib.rs b/src/lib.rs
index 9f73aa9..e14a81e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -4,19 +4,25 @@ use crate::{
     params::{RetainParamsLocal, RetainParamsShared},
 };
 use clack_extensions::{
-    audio_ports::*,
+    audio_ports::PluginAudioPorts,
     gui::PluginGui,
     latency::{PluginLatency, PluginLatencyImpl},
-    params::*,
+    params::PluginParams,
     state::PluginState,
 };
-use clack_plugin::prelude::*;
+use clack_plugin::{
+    plugin::features::{AUDIO_EFFECT, STEREO},
+    prelude::*,
+};
 use std::sync::Arc;
 
 mod audio;
 mod gui;
 mod params;
+mod retain;
+mod window_function;
 mod window_size;
+mod windowed_fft;
 
 /// The type that represents our plugin in Clack.
 ///
@@ -43,8 +49,6 @@ impl Plugin for RetainPlugin {
 
 impl DefaultPluginFactory for RetainPlugin {
     fn get_descriptor() -> PluginDescriptor {
-        use clack_plugin::plugin::features::*;
-
         PluginDescriptor::new("com.haemolacriaa.retain", "Retain")
             .with_description("Retains only the nth largest magnitude frequencies in the signal")
             .with_features([AUDIO_EFFECT, STEREO])
@@ -79,7 +83,7 @@ impl<'a> RetainPluginShared<'a> {
     fn new(host: HostSharedHandle<'a>) -> Self {
         RetainPluginShared {
             params: Arc::new(RetainParamsShared::new()),
-            host: host,
+            host,
         }
     }
 }
diff --git a/src/params.rs b/src/params.rs
index 00782ba..c895bca 100644
--- a/src/params.rs
+++ b/src/params.rs
@@ -1,7 +1,12 @@
 //! Contains all types and implementations related to parameter management.
 
 use crate::{RetainPluginMainThread, window_size::WindowSize};
-use clack_extensions::{params::*, state::PluginStateImpl};
+use clack_extensions::{
+    params::{
+        ParamDisplayWriter, ParamInfo, ParamInfoFlags, ParamInfoWriter, PluginMainThreadParams,
+    },
+    state::PluginStateImpl,
+};
 use clack_plugin::{
     events::{event_types::ParamValueEvent, spaces::CoreEventSpace},
     prelude::*,
@@ -18,7 +23,7 @@ use std::{
 
 /// The default value of the order parameter.
 const DEFAULT_ORDER: usize = 1000;
-pub const DEFAULT_WINDOW_SIZE: WindowSize = WindowSize::Size128;
+pub const DEFAULT_WINDOW_SIZE: WindowSize = WindowSize::Size32768;
 
 /// A struct that manages the parameters for our plugin.
 ///
diff --git a/src/retain.rs b/src/retain.rs
new file mode 100644
index 0000000..bc3934a
--- /dev/null
+++ b/src/retain.rs
@@ -0,0 +1,30 @@
+use hashbrown::HashSet;
+use num_complex::Complex;
+
+#[inline(always)]
+pub fn retain_top_n_magnitudes(fft_real_signal: &mut [Complex<f32>], n: usize) {
+    // you'd be keeping the entire signal
+    if n >= fft_real_signal.len() {
+        return;
+    }
+
+    let mut indexed = fft_real_signal
+        .iter()
+        .copied()
+        .enumerate()
+        .collect::<Vec<(usize, Complex<f32>)>>();
+
+    let target = fft_real_signal.len() - n;
+    indexed.select_nth_unstable_by(target, |(_, c0), (_, c1)| c0.norm().total_cmp(&c1.norm()));
+
+    let top_indices = indexed[target..]
+        .iter()
+        .map(|&(i, _)| i)
+        .collect::<HashSet<usize>>();
+
+    for (i, val) in fft_real_signal.iter_mut().enumerate() {
+        if !top_indices.contains(&i) {
+            *val = Complex::ZERO;
+        }
+    }
+}
diff --git a/src/window_function.rs b/src/window_function.rs
new file mode 100644
index 0000000..a633af1
--- /dev/null
+++ b/src/window_function.rs
@@ -0,0 +1,18 @@
+pub trait WindowFunction: Send {
+    fn apply(&self, data: &mut [f32]);
+    fn reverse(&self, data: &mut [f32]);
+}
+
+pub struct RectangularWindow {}
+
+impl RectangularWindow {
+    pub fn new(_fft_size: usize) -> Self {
+        Self {}
+    }
+}
+
+impl WindowFunction for RectangularWindow {
+    fn apply(&self, _data: &mut [f32]) {}
+
+    fn reverse(&self, _data: &mut [f32]) {}
+}
diff --git a/src/windowed_fft.rs b/src/windowed_fft.rs
index 161d3f0..6d11564 100644
--- a/src/windowed_fft.rs
+++ b/src/windowed_fft.rs
@@ -1,12 +1,13 @@
 #![allow(clippy::must_use_candidate)]
 #![allow(clippy::return_self_not_must_use)]
 
-use crate::window::{
+/*use crate::window::{
     BlackmanHarrisWindow, HammingWindow, HannWindow, RectangularWindow, WindowFunction,
-};
+};*/
+use crate::window_function::{RectangularWindow, WindowFunction};
 use num_complex::Complex;
 use realfft::{ComplexToReal, RealFftPlanner, RealToComplex};
-use std::sync::Arc;
+use std::{collections::VecDeque, sync::Arc};
 
 pub struct WindowedRealFft {
     fft_size: usize,
@@ -14,16 +15,31 @@ pub struct WindowedRealFft {
     forward: Arc<dyn RealToComplex<f32>>,
     inverse: Arc<dyn ComplexToReal<f32>>,
     window_function: Box<dyn WindowFunction>,
-    original_length: usize,
+    input: VecDeque<f32>,
+    output: VecDeque<f32>,
+    spectrum: Vec<Complex<f32>>,
+    scratch: Vec<Complex<f32>>,
 }
 
 impl WindowedRealFft {
-    pub fn new(fft_size: usize) -> Self {
+    pub fn new(mut fft_size: usize) -> Self {
+		if fft_size == 0 {
+			fft_size = 1;
+		}
+		
         let mut planner = RealFftPlanner::new();
         let forward = planner.plan_fft_forward(fft_size);
         let inverse = planner.plan_fft_inverse(fft_size);
 
-        let window_function = Box::new(HannWindow::new(fft_size));
+        let window_function = Box::new(RectangularWindow::new(fft_size));
+
+        let input = VecDeque::with_capacity(fft_size);
+        let output = VecDeque::with_capacity(fft_size);
+
+        let spectrum = vec![Complex::ZERO; (fft_size / 2) + 1];
+
+        let max_scratch_len = forward.get_scratch_len().max(inverse.get_scratch_len());
+        let scratch = vec![Complex::ZERO; max_scratch_len];
 
         Self {
             fft_size,
@@ -31,81 +47,84 @@ impl WindowedRealFft {
             forward,
             inverse,
             window_function,
-            original_length: 0,
+            input,
+            output,
+            spectrum,
+            scratch,
         }
     }
 
-    pub fn fft_size(mut self, value: usize) -> Self {
+    pub fn fft_size(mut self, mut value: usize) -> Self {
+		if value == 0 {
+			value = 1;
+		}
+		
         self.fft_size = value;
 
-        let forward = self.planner.plan_fft_forward(value);
-        let inverse = self.planner.plan_fft_inverse(value);
-
-        self.forward = forward;
-        self.inverse = inverse;
-
-        self
-    }
-
-    pub fn original_length(&mut self, value: usize) {
-        self.original_length = value;
-    }
-
-    pub fn forward(&mut self, data: impl Into<Vec<f32>>) -> Vec<Vec<Complex<f32>>> {
-        let mut data = data.into();
-
-        let Some(new_length) = data.len().checked_next_multiple_of(self.fft_size) else {
-            return vec![];
-        };
+        self.forward = self.planner.plan_fft_forward(self.fft_size);
+        self.inverse = self.planner.plan_fft_inverse(self.fft_size);
 
-        self.original_length = data.len();
+        self.input.reserve_exact(self.fft_size);
+        self.output.reserve_exact(self.fft_size);
 
-        data.resize(new_length, 0.0);
+        self.input.clear();
+        self.output.clear();
 
-        self.window_function
-            .apply(data)
-            .into_par_iter()
-            .map(|mut chunk| {
-                let mut output = self.forward.make_output_vec();
+        self.spectrum.resize((self.fft_size / 2) + 1, Complex::ZERO);
 
-                self.forward.process(&mut chunk, &mut output).unwrap();
+        let max_scratch_len = self
+            .forward
+            .get_scratch_len()
+            .max(self.inverse.get_scratch_len());
+        self.scratch.resize(max_scratch_len, Complex::ZERO);
 
-                output
-            })
-            .collect::<Vec<Vec<Complex<f32>>>>()
+        self
     }
 
-    pub fn inverse(&mut self, data: Vec<Vec<Complex<f32>>>) -> Vec<f32> {
-        let data = data
-            .into_par_iter()
-            .map(|mut chunk| {
-                let mut real_output = self.inverse.make_output_vec();
+    pub fn clear_input(&mut self) {
+        self.input.clear();
+    }
 
-                self.inverse.process(&mut chunk, &mut real_output).unwrap();
+    pub fn push_front_input(&mut self, value: f32) -> bool {
+        self.input.push_front(value);
 
-                let chunk_size_f32 = self.fft_size as f32;
+        self.input.len() >= self.fft_size
+    }
 
-                real_output
-                    .into_iter()
-                    .map(|sample| sample / chunk_size_f32)
-                    .collect::<Vec<f32>>()
-            })
-            .collect::<Vec<Vec<f32>>>();
+    pub fn pop_back_output(&mut self) -> f32 {
+        self.output.pop_back().unwrap_or(0.0)
+    }
 
-        let mut output = self.window_function.reverse(data);
+    pub fn get_spectrum(&mut self) -> &mut [Complex<f32>] {
+        &mut self.spectrum
+    }
 
-        output.truncate(self.original_length);
+    pub fn forward(&mut self) {
+        self.window_function.apply(self.input.make_contiguous());
 
-        output
+        self.forward
+            .process_with_scratch(
+                self.input.make_contiguous(),
+                &mut self.spectrum,
+                &mut self.scratch,
+            );
     }
 
-    pub fn i32_to_f32(item: Vec<i32>) -> Vec<f32> {
-        item.into_iter().map(f32::from).collect::<Vec<f32>>()
-    }
+    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;
+        for sample in &mut self.output {
+            *sample = *sample / fft_size_f32;
+        }
 
-    pub fn f32_to_i32(item: Vec<f32>) -> Vec<i32> {
-        item.into_iter()
-            .map(|value| value as i32)
-            .collect::<Vec<i32>>()
+        self.window_function.reverse(self.output.make_contiguous());
     }
 }