Avoid allocations in the real-time audio thread

2 files changed, 76 insertions, 59 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 0908ee7..88e1df0 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -7,7 +7,7 @@ mod window_type;
 mod windowed_fft;
 
 use crate::{
-    retain::retain_top_n_magnitudes, window_size::WindowSize, window_type::WindowType,
+    retain::RetainAlgorithm, window_size::WindowSize, window_type::WindowType,
     windowed_fft::WindowedRealFft,
 };
 use egui::{Align, CentralPanel, ComboBox, Layout, Slider, Vec2};
@@ -31,6 +31,7 @@ pub struct Retain {
     initial_gui_state: Option<GuiState>,
     fft_left: WindowedRealFft,
     fft_right: WindowedRealFft,
+	retain_algorithm: RetainAlgorithm,
 }
 
 impl Default for Retain {
@@ -50,6 +51,7 @@ impl Default for Retain {
             }),
             fft_left,
             fft_right,
+			retain_algorithm: RetainAlgorithm::new(),
         }
     }
 }
@@ -57,7 +59,7 @@ impl Default for Retain {
 #[derive(Params)]
 pub struct RetainParams {
     /// This allows the resizing of the plugin to be restored
-    #[persist = "editor-state"]
+    #[persist = "editor_state"]
     editor_state: Arc<EguiState>,
 
     /// the number of top frequencies to keep
@@ -70,11 +72,11 @@ pub struct RetainParams {
     pub window_size: IntParam,
 
     /// the type of window function used to reduce spectral leakage or clicking sounds
-    /// changing this makes a big difference
+    /// playing with this makes a big difference in output
     #[id = "window_type"]
     pub window_type: EnumParam<WindowType>,
 
-    /// whether the plugin retains or removes the highest magnitudes
+    /// whether the plugin retains or removes the highest magnitude frequencies
     #[id = "complement"]
     pub complement: BoolParam,
 }
@@ -304,7 +306,7 @@ impl Plugin for Retain {
 					let order = self.params.order.value() as usize;
 					let complement = self.params.complement.value();
 
-                    retain_top_n_magnitudes(self.fft_left.get_spectrum(), order, complement);
+                    self.retain_algorithm.calculate(self.fft_left.get_spectrum(), order, complement);
 
                     self.fft_left.inverse();
                     self.fft_left.clear_input();
@@ -320,7 +322,7 @@ impl Plugin for Retain {
 					let order = self.params.order.value() as usize;
 					let complement = self.params.complement.value();
 
-                    retain_top_n_magnitudes(self.fft_right.get_spectrum(), order, complement);
+                    self.retain_algorithm.calculate(self.fft_right.get_spectrum(), order, complement);
 
                     self.fft_right.inverse();
                     self.fft_right.clear_input();
diff --git a/src/retain.rs b/src/retain.rs
index 5c14246..8e11164 100644
--- a/src/retain.rs
+++ b/src/retain.rs
@@ -1,55 +1,70 @@
 use num_complex::Complex;
-use rustc_hash::FxHashSet;
-
-#[inline(always)]
-pub fn retain_top_n_magnitudes(spectrum: &mut [Complex<f32>], n: usize, complement: bool) {
-    // you'd be keeping the entire signal
-    if n >= spectrum.len() && !complement {
-        return;
-    }
-
-    if n == 0 && complement {
-        return;
-    }
-
-    // keeping nothing at all
-    if n == 0 && !complement {
-        for phasor in spectrum {
-            *phasor = Complex::ZERO;
-        }
-
-        return;
-    }
-
-    if n >= spectrum.len() && complement {
-        for phasor in spectrum {
-            *phasor = Complex::ZERO;
-        }
-
-        return;
-    }
-
-    let mut indexed = spectrum
-        .iter()
-        .copied()
-        .enumerate()
-        .collect::<Vec<(usize, Complex<f32>)>>();
-
-    let target = spectrum.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::<FxHashSet<usize>>();
-
-    for (i, phasor) in spectrum.iter_mut().enumerate() {
-        if top_indices.contains(&i) && complement {
-            *phasor = Complex::ZERO;
-        }
-
-        if !top_indices.contains(&i) && !complement {
-            *phasor = Complex::ZERO;
-        }
-    }
+use rustc_hash::{FxBuildHasher, FxHashSet};
+
+/// I thought about doing this to avoid allocating new memory in the audio thread
+pub struct RetainAlgorithm {
+	indexed_spectrum: Vec<(usize, Complex<f32>)>,
+	top_indices: FxHashSet<usize>,
 }
+
+impl RetainAlgorithm {
+	pub fn new() -> Self {
+		let max_size = (32768 / 2) + 1;
+		
+		Self {
+			indexed_spectrum: Vec::with_capacity(max_size),
+			top_indices: FxHashSet::with_capacity_and_hasher(max_size, FxBuildHasher::default()),
+		}
+	}
+	
+	pub fn calculate(&mut self, spectrum: &mut [Complex<f32>], n: usize, complement: bool) {
+		// you'd be keeping the entire signal
+		if n >= spectrum.len() && !complement {
+			return;
+		}
+
+		if n == 0 && complement {
+			return;
+		}
+
+		// keeping nothing at all
+		if n == 0 && !complement {
+			for phasor in spectrum {
+				*phasor = Complex::ZERO;
+			}
+
+			return;
+		}
+
+		if n >= spectrum.len() && complement {
+			for phasor in spectrum {
+				*phasor = Complex::ZERO;
+			}
+
+			return;
+		}
+		
+		self.indexed_spectrum.clear();
+		for (i, phasor) in spectrum.iter().copied().enumerate() {
+			self.indexed_spectrum.push((i, phasor));
+		}
+
+		let target = spectrum.len() - n;
+		self.indexed_spectrum.select_nth_unstable_by(target, |(_, c0), (_, c1)| c0.norm().total_cmp(&c1.norm()));
+		
+		self.top_indices.clear();
+		for i in self.indexed_spectrum[target..].iter().map(|&(i, _)| i) {
+			self.top_indices.insert(i);
+		}
+
+		for (i, phasor) in spectrum.iter_mut().enumerate() {
+			if self.top_indices.contains(&i) && complement {
+				*phasor = Complex::ZERO;
+			}
+
+			if !self.top_indices.contains(&i) && !complement {
+				*phasor = Complex::ZERO;
+			}
+		}
+	}
+}
\ No newline at end of file