Add Blackman-Harris window

6 files changed, 257 insertions, 41 deletions

diff --git a/src/lib.rs b/src/lib.rs
index 3288088..7157c40 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -22,6 +22,7 @@ mod params;
 mod retain;
 mod window_function;
 mod window_size;
+mod window_type;
 mod windowed_fft;
 
 /// The type that represents our plugin in Clack.
diff --git a/src/retain.rs b/src/retain.rs
index 49730e1..8c4c8eb 100644
--- a/src/retain.rs
+++ b/src/retain.rs
@@ -2,26 +2,26 @@ use num_complex::Complex;
 use rustc_hash::FxHashSet;
 
 #[inline(always)]
-pub fn retain_top_n_magnitudes(fft_real_signal: &mut [Complex<f32>], n: usize) {
+pub fn retain_top_n_magnitudes(spectrum: &mut [Complex<f32>], n: usize) {
     // you'd be keeping the entire signal
     // or keeping nothing at all
-    if n >= fft_real_signal.len() {
+    if n >= spectrum.len() {
         return;
     } else if n == 0 {
-        for phasor in fft_real_signal {
+        for phasor in spectrum {
             *phasor = Complex::ZERO;
         }
 
         return;
     }
 
-    let mut indexed = fft_real_signal
+    let mut indexed = spectrum
         .iter()
         .copied()
         .enumerate()
         .collect::<Vec<(usize, Complex<f32>)>>();
 
-    let target = fft_real_signal.len() - n;
+    let target = spectrum.len() - n;
     indexed.select_nth_unstable_by(target, |(_, c0), (_, c1)| c0.norm().total_cmp(&c1.norm()));
 
     let top_indices = indexed[target..]
@@ -29,7 +29,7 @@ pub fn retain_top_n_magnitudes(fft_real_signal: &mut [Complex<f32>], n: usize) {
         .map(|&(i, _)| i)
         .collect::<FxHashSet<usize>>();
 
-    for (i, phasor) in fft_real_signal.iter_mut().enumerate() {
+    for (i, phasor) in spectrum.iter_mut().enumerate() {
         if !top_indices.contains(&i) {
             *phasor = Complex::ZERO;
         }
diff --git a/src/window_function.rs b/src/window_function.rs
index 3fbcedf..8e21694 100644
--- a/src/window_function.rs
+++ b/src/window_function.rs
@@ -1,10 +1,15 @@
 use crate::window_size::WindowSize;
-use std::{collections::VecDeque, f32::consts::PI, mem};
+use std::{collections::VecDeque, f32::consts::PI};
 
 pub trait WindowFunction: Send {
+    /// Apply the window in-place
     fn apply(&mut self, data: &mut VecDeque<f32>);
+    /// Reverse the window in-place
     fn reverse(&mut self, data: &mut [f32]);
+    /// Yields the amount of samples still required to apply the window
     fn needed(&self) -> usize;
+    /// Updates the window size
+    fn window_size(&mut self, window_size: WindowSize);
 }
 
 pub struct RectangularWindow {
@@ -27,6 +32,8 @@ impl WindowFunction for RectangularWindow {
     fn needed(&self) -> usize {
         self.window_size
     }
+
+    fn window_size(&mut self, _window_size: WindowSize) {}
 }
 
 pub struct HannWindow {
@@ -34,7 +41,7 @@ pub struct HannWindow {
     function: Vec<f32>,
     normalize: Vec<f32>,
     previous: Vec<f32>,
-    previous_overlap: Vec<f32>,
+    overlap_add: Vec<f32>,
 }
 
 impl HannWindow {
@@ -44,16 +51,16 @@ impl HannWindow {
         let window_size_f32 = window_size as f32;
 
         // 50% overlap
-        let previous_overlap = vec![0.0; half_window_size];
+        let overlap_add = vec![0.0; window_size];
 
         // with capacity is important for the first needed samples
         let previous = Vec::with_capacity(window_size);
 
         let function = (0..window_size)
-            .map(|n| {
-                let n = n as f32;
+            .map(|i| {
+                let i = i as f32;
 
-                0.5 * (1.0 - f32::cos((2.0 * PI * n) / (window_size_f32 - 1.0)))
+                0.5 * (1.0 - f32::cos((2.0 * PI * i) / (window_size_f32 - 1.0)))
             })
             .collect::<Vec<f32>>();
 
@@ -69,7 +76,7 @@ impl HannWindow {
             function,
             normalize,
             previous,
-            previous_overlap,
+            overlap_add,
         }
     }
 }
@@ -81,8 +88,8 @@ impl WindowFunction for HannWindow {
         // the full window size is needed to window over the samples the first time
         if self.previous.is_empty() {
             // store the latter half for when only half the size is needed
-            for sample in data.iter().skip(half_window_size) {
-                self.previous.push(*sample);
+            for sample in data.iter().skip(half_window_size).copied() {
+                self.previous.push(sample);
             }
 
             // apply the function
@@ -95,15 +102,15 @@ impl WindowFunction for HannWindow {
 
         // In a real time "sliding window" it's necessary to add in the previous samples
         // to get the full resolution
-        for sample in self.previous.iter().rev() {
-            data.push_front(*sample);
+        for sample in self.previous.iter().rev().copied() {
+            data.push_front(sample);
         }
 
         self.previous.clear();
 
         // keep this iteration's samples for the next window
-        for sample in data.iter().skip(half_window_size) {
-            self.previous.push(*sample);
+        for sample in data.iter().skip(half_window_size).copied() {
+            self.previous.push(sample);
         }
 
         // apply the function
@@ -115,26 +122,22 @@ impl WindowFunction for HannWindow {
     fn reverse(&mut self, data: &mut [f32]) {
         let half_window_size = self.window_size / 2;
 
-        for (sample, function) in data.iter_mut().zip(self.function.iter()) {
-            *sample *= *function;
+        for i in 0..self.window_size {
+            self.overlap_add[i] += data[i] * self.function[i];
         }
 
-        // add back in the previous applied window function
-        // this will reverse the 50% overlap
-        for (i, (previous, sample)) in self
-            .previous_overlap
-            .iter()
-            .zip(data.iter_mut())
-            .enumerate()
-        {
-            *sample = (*sample + previous) / self.normalize[i];
+        for i in 0..half_window_size {
+            if self.normalize[i] > 1e-6 {
+                data[i] = self.overlap_add[i] / self.normalize[i];
+            } else {
+                data[i] = 0.0;
+            }
         }
 
-        self.previous_overlap.clear();
+        self.overlap_add.rotate_left(half_window_size);
 
-        // copy this window's tail for future overlap
-        for sample in data.iter().skip(half_window_size) {
-            self.previous_overlap.push(*sample);
+        for sample in &mut self.overlap_add[half_window_size..] {
+            *sample = 0.0;
         }
     }
 
@@ -145,4 +148,185 @@ impl WindowFunction for HannWindow {
             self.window_size / 2
         }
     }
+
+    fn window_size(&mut self, window_size: WindowSize) {
+        let window_size = window_size.inner();
+        let half_window_size = window_size / 2;
+        let window_size_f32 = window_size as f32;
+
+        self.overlap_add.resize(window_size, 0.0);
+
+        self.previous.clear();
+        self.previous.reserve_exact(window_size);
+
+        self.function = (0..window_size)
+            .map(|i| {
+                let i = i as f32;
+
+                0.5 * (1.0 - f32::cos((2.0 * PI * i) / (window_size_f32 - 1.0)))
+            })
+            .collect::<Vec<f32>>();
+
+        self.normalize = (0..half_window_size)
+            .map(|i| {
+                (self.function[i] * self.function[i])
+                    + (self.function[i + half_window_size] * self.function[i + half_window_size])
+            })
+            .collect::<Vec<f32>>();
+    }
+}
+
+pub struct BlackmanHarrisWindow {
+    window_size: usize,
+    function: Vec<f32>,
+    normalize: Vec<f32>,
+    previous: Vec<f32>,
+    overlap_add: Vec<f32>,
+}
+
+impl BlackmanHarrisWindow {
+    const A_0: f32 = 0.35875;
+    const A_1: f32 = 0.48829;
+    const A_2: f32 = 0.14128;
+    const A_3: f32 = 0.01168;
+
+    pub fn new(window_size: &WindowSize) -> Self {
+        let window_size = window_size.inner();
+        let quarter_window_size = window_size / 4;
+        let three_quarters_window_size = 3 * quarter_window_size;
+        let window_size_f32 = window_size as f32;
+
+        let overlap_add = vec![0.0; window_size];
+
+        let previous = Vec::with_capacity(three_quarters_window_size);
+
+        let function = (0..window_size)
+            .map(|i| {
+                let i = i as f32;
+
+                let two = f32::cos((2.0 * PI * i) / (window_size_f32 - 1.0));
+                let four = f32::cos((4.0 * PI * i) / (window_size_f32 - 1.0));
+                let six = f32::cos((6.0 * PI * i) / (window_size_f32 - 1.0));
+
+                Self::A_0 - (Self::A_1 * two) + (Self::A_2 * four) - (Self::A_3 * six)
+            })
+            .collect::<Vec<f32>>();
+
+        #[rustfmt::skip]
+        let normalize = (0..quarter_window_size)
+            .map(|i| {
+                (function[i] * function[i])
+                    + (function[i + quarter_window_size] * function[i + quarter_window_size])
+                    + (function[i + 2 * quarter_window_size] * function[i + 2 * quarter_window_size])
+                    + (function[i + 3 * quarter_window_size] * function[i + 3 * quarter_window_size])
+            })
+            .collect::<Vec<f32>>();
+
+        Self {
+            window_size,
+            function,
+            normalize,
+            previous,
+            overlap_add,
+        }
+    }
+}
+
+impl WindowFunction for BlackmanHarrisWindow {
+    fn apply(&mut self, data: &mut VecDeque<f32>) {
+        let quarter_window_size = self.window_size / 4;
+
+        if self.previous.is_empty() {
+            for sample in data.iter().skip(quarter_window_size).copied() {
+                self.previous.push(sample);
+            }
+
+            for (i, sample) in data.iter_mut().enumerate() {
+                *sample *= self.function[i];
+            }
+
+            return;
+        }
+
+        for sample in self.previous.iter().rev().copied() {
+            data.push_front(sample);
+        }
+
+        self.previous.clear();
+
+        for sample in data.iter().skip(quarter_window_size).copied() {
+            self.previous.push(sample);
+        }
+
+        for (i, sample) in data.iter_mut().enumerate() {
+            *sample *= self.function[i];
+        }
+    }
+
+    fn reverse(&mut self, data: &mut [f32]) {
+        let quarter_window_size = self.window_size / 4;
+        let three_quarters_window_size = 3 * quarter_window_size;
+
+        for i in 0..self.window_size {
+            self.overlap_add[i] += data[i] * self.function[i];
+        }
+
+        for i in 0..quarter_window_size {
+            if self.normalize[i] > 1e-6 {
+                data[i] = self.overlap_add[i] / self.normalize[i];
+            } else {
+                data[i] = 0.0;
+            }
+        }
+
+        self.overlap_add.rotate_left(quarter_window_size);
+
+        for sample in &mut self.overlap_add[three_quarters_window_size..] {
+            *sample = 0.0;
+        }
+    }
+
+    fn needed(&self) -> usize {
+        if self.previous.is_empty() {
+            self.window_size
+        } else {
+            self.window_size / 4
+        }
+    }
+
+    fn window_size(&mut self, window_size: WindowSize) {
+        let window_size = window_size.inner();
+        let quarter_window_size = window_size / 4;
+        let three_quarters_window_size = 3 * quarter_window_size;
+        let window_size_f32 = window_size as f32;
+
+        self.overlap_add.resize(window_size, 0.0);
+
+        self.previous.clear();
+        self.previous.reserve_exact(window_size);
+
+        self.function = (0..window_size)
+            .map(|i| {
+                let i = i as f32;
+
+                let two = f32::cos((2.0 * PI * i) / (window_size_f32 - 1.0));
+                let four = f32::cos((4.0 * PI * i) / (window_size_f32 - 1.0));
+                let six = f32::cos((6.0 * PI * i) / (window_size_f32 - 1.0));
+
+                Self::A_0 - (Self::A_1 * two) + (Self::A_2 * four) - (Self::A_3 * six)
+            })
+            .collect::<Vec<f32>>();
+
+        self.normalize = (0..quarter_window_size)
+            .map(|i| {
+                (self.function[i] * self.function[i])
+                    + (self.function[i + quarter_window_size]
+                        * self.function[i + quarter_window_size])
+                    + (self.function[i + 2 * quarter_window_size]
+                        * self.function[i + 2 * quarter_window_size])
+                    + (self.function[i + 3 * quarter_window_size]
+                        * self.function[i + 3 * quarter_window_size])
+            })
+            .collect::<Vec<f32>>();
+    }
 }
diff --git a/src/window_size.rs b/src/window_size.rs
index b3547d6..bfc3957 100644
--- a/src/window_size.rs
+++ b/src/window_size.rs
@@ -1,5 +1,5 @@
-/// The max value of the custom window size is u32::MAX.
-/// That's because it's the maximum latency one can report.
+/// The max value of the custom window size is `u32::MAX`.
+/// This is because it's the maximum latency one can report to a CLAP host.
 #[derive(Debug, PartialEq)]
 pub enum WindowSize {
     Size128,
diff --git a/src/window_type.rs b/src/window_type.rs
new file mode 100644
index 0000000..371a410
--- /dev/null
+++ b/src/window_type.rs
@@ -0,0 +1,32 @@
+use crate::{
+    window_function::{BlackmanHarrisWindow, HannWindow, RectangularWindow, WindowFunction},
+    window_size::WindowSize,
+};
+
+pub enum WindowType {
+    Hann,
+    Rectangular,
+    BlackmanHarris,
+}
+
+impl WindowType {
+    fn into_function(self, window_size: &WindowSize) -> Box<dyn WindowFunction> {
+        match self {
+            Self::Hann => Box::new(HannWindow::new(window_size)),
+            Self::Rectangular => Box::new(RectangularWindow::new(window_size)),
+            Self::BlackmanHarris => Box::new(BlackmanHarrisWindow::new(window_size)),
+        }
+    }
+
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            Self::Hann => "Hann",
+            Self::Rectangular => "Rectangular",
+            Self::BlackmanHarris => "Blackman-Harris",
+        }
+    }
+
+    pub fn iter() -> impl Iterator<Item = Self> {
+        [Self::Hann, Self::Rectangular].into_iter()
+    }
+}
diff --git a/src/windowed_fft.rs b/src/windowed_fft.rs
index 114ea0b..c08c2b4 100644
--- a/src/windowed_fft.rs
+++ b/src/windowed_fft.rs
@@ -1,11 +1,8 @@
 #![allow(clippy::must_use_candidate)]
 #![allow(clippy::return_self_not_must_use)]
 
-/*use crate::window::{
-    BlackmanHarrisWindow, HammingWindow, HannWindow, RectangularWindow, WindowFunction,
-};*/
 use crate::{
-    window_function::{HannWindow, RectangularWindow, WindowFunction},
+    window_function::{BlackmanHarrisWindow, HannWindow, RectangularWindow, WindowFunction},
     window_size::WindowSize,
 };
 use num_complex::Complex;
@@ -26,7 +23,7 @@ pub struct WindowedRealFft {
 
 impl WindowedRealFft {
     pub fn new(window_size: WindowSize) -> Self {
-        let window_function = Box::new(HannWindow::new(&window_size));
+        let window_function = Box::new(BlackmanHarrisWindow::new(&window_size));
         let window_size = window_size.inner();
 
         let mut planner = RealFftPlanner::new();
@@ -54,6 +51,8 @@ impl WindowedRealFft {
         }
     }
 
+    pub fn window_function(&mut self, window_function: Box<dyn WindowFunction>) {}
+
     pub fn window_size(&mut self, window_size: WindowSize) {
         if window_size == self.window_size.into() {
             return;