account for latency

1 files changed, 11 insertions, 4 deletions

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;
         }
     }
 }