Add sliding hann window

5 files changed, 150 insertions, 35 deletions

diff --git a/src/gui.rs b/src/gui.rs
index 3d2dfb2..9216d5a 100644
--- a/src/gui.rs
+++ b/src/gui.rs
@@ -45,7 +45,7 @@ impl RetainPluginGui {
     pub fn new(parent: Window<'_>, plugin_state: &RetainPluginShared) -> Self {
         let settings = WindowOpenOptions {
             title: "Retain".to_string(),
-            size: Size::new(1000.0, 500.0),
+            size: Size::new(600.0, 350.0),
             scale: WindowScalePolicy::SystemScaleFactor,
             gl_config: Some(GlConfig::default()),
         };
@@ -166,8 +166,8 @@ impl PluginGuiImpl for RetainPluginMainThread<'_> {
 
     fn get_size(&mut self) -> Option<GuiSize> {
         Some(GuiSize {
-            width: 1000,
-            height: 500,
+            width: 600,
+            height: 350,
         })
     }
 
diff --git a/src/params.rs b/src/params.rs
index 041f838..2916ace 100644
--- a/src/params.rs
+++ b/src/params.rs
@@ -48,7 +48,7 @@ impl RetainParamsShared {
     pub fn new() -> Self {
         Self {
             order: AtomicUsize::new(DEFAULT_ORDER),
-            window_size: AtomicUsize::new(DEFAULT_WINDOW_SIZE.into_inner()),
+            window_size: AtomicUsize::new(DEFAULT_WINDOW_SIZE.inner()),
             has_gesture: AtomicBool::new(false),
         }
     }
@@ -313,7 +313,7 @@ impl PluginMainThreadParams for RetainPluginMainThread<'_> {
             module: b"",
             min_value: 128.0,
             max_value: usize::MAX as f64,
-            default_value: DEFAULT_WINDOW_SIZE.into_inner() as f64,
+            default_value: DEFAULT_WINDOW_SIZE.inner() as f64,
         });
     }
 
diff --git a/src/window_function.rs b/src/window_function.rs
index a633af1..3fbcedf 100644
--- a/src/window_function.rs
+++ b/src/window_function.rs
@@ -1,18 +1,148 @@
+use crate::window_size::WindowSize;
+use std::{collections::VecDeque, f32::consts::PI, mem};
+
 pub trait WindowFunction: Send {
-    fn apply(&self, data: &mut [f32]);
-    fn reverse(&self, data: &mut [f32]);
+    fn apply(&mut self, data: &mut VecDeque<f32>);
+    fn reverse(&mut self, data: &mut [f32]);
+    fn needed(&self) -> usize;
 }
 
-pub struct RectangularWindow {}
+pub struct RectangularWindow {
+    window_size: usize,
+}
 
 impl RectangularWindow {
-    pub fn new(_fft_size: usize) -> Self {
-        Self {}
+    pub fn new(window_size: &WindowSize) -> Self {
+        Self {
+            window_size: window_size.inner(),
+        }
     }
 }
 
 impl WindowFunction for RectangularWindow {
-    fn apply(&self, _data: &mut [f32]) {}
+    fn apply(&mut self, _data: &mut VecDeque<f32>) {}
+
+    fn reverse(&mut self, _data: &mut [f32]) {}
+
+    fn needed(&self) -> usize {
+        self.window_size
+    }
+}
+
+pub struct HannWindow {
+    window_size: usize,
+    function: Vec<f32>,
+    normalize: Vec<f32>,
+    previous: Vec<f32>,
+    previous_overlap: Vec<f32>,
+}
+
+impl HannWindow {
+    pub fn new(window_size: &WindowSize) -> Self {
+        let window_size = window_size.inner();
+        let half_window_size = window_size / 2;
+        let window_size_f32 = window_size as f32;
+
+        // 50% overlap
+        let previous_overlap = vec![0.0; half_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;
+
+                0.5 * (1.0 - f32::cos((2.0 * PI * n) / (window_size_f32 - 1.0)))
+            })
+            .collect::<Vec<f32>>();
+
+        let normalize = (0..half_window_size)
+            .map(|i| {
+                (function[i] * function[i])
+                    + (function[i + half_window_size] * function[i + half_window_size])
+            })
+            .collect::<Vec<f32>>();
+
+        Self {
+            window_size,
+            function,
+            normalize,
+            previous,
+            previous_overlap,
+        }
+    }
+}
 
-    fn reverse(&self, _data: &mut [f32]) {}
+impl WindowFunction for HannWindow {
+    fn apply(&mut self, data: &mut VecDeque<f32>) {
+        let half_window_size = self.window_size / 2;
+
+        // 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);
+            }
+
+            // apply the function
+            for (i, sample) in data.iter_mut().enumerate() {
+                *sample *= self.function[i];
+            }
+
+            return;
+        }
+
+        // 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);
+        }
+
+        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);
+        }
+
+        // apply the function
+        for (i, sample) in data.iter_mut().enumerate() {
+            *sample *= self.function[i];
+        }
+    }
+
+    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;
+        }
+
+        // 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];
+        }
+
+        self.previous_overlap.clear();
+
+        // copy this window's tail for future overlap
+        for sample in data.iter().skip(half_window_size) {
+            self.previous_overlap.push(*sample);
+        }
+    }
+
+    fn needed(&self) -> usize {
+        if self.previous.is_empty() {
+            self.window_size
+        } else {
+            self.window_size / 2
+        }
+    }
 }
diff --git a/src/window_size.rs b/src/window_size.rs
index 9d3f66b..b3547d6 100644
--- a/src/window_size.rs
+++ b/src/window_size.rs
@@ -49,21 +49,6 @@ impl WindowSize {
         }
     }
 
-    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,
diff --git a/src/windowed_fft.rs b/src/windowed_fft.rs
index 709a68b..114ea0b 100644
--- a/src/windowed_fft.rs
+++ b/src/windowed_fft.rs
@@ -5,7 +5,7 @@
     BlackmanHarrisWindow, HammingWindow, HannWindow, RectangularWindow, WindowFunction,
 };*/
 use crate::{
-    window_function::{RectangularWindow, WindowFunction},
+    window_function::{HannWindow, RectangularWindow, WindowFunction},
     window_size::WindowSize,
 };
 use num_complex::Complex;
@@ -26,14 +26,13 @@ pub struct WindowedRealFft {
 
 impl WindowedRealFft {
     pub fn new(window_size: WindowSize) -> Self {
-        let window_size = window_size.into_inner();
+        let window_function = Box::new(HannWindow::new(&window_size));
+        let window_size = window_size.inner();
 
         let mut planner = RealFftPlanner::new();
         let forward = planner.plan_fft_forward(window_size);
         let inverse = planner.plan_fft_inverse(window_size);
 
-        let window_function = Box::new(RectangularWindow::new(window_size));
-
         let input = VecDeque::with_capacity(window_size);
         let output = VecDeque::with_capacity(window_size);
 
@@ -60,7 +59,8 @@ impl WindowedRealFft {
             return;
         }
 
-        self.window_size = window_size.into_inner();
+        self.window_function = Box::new(HannWindow::new(&window_size));
+        self.window_size = window_size.inner();
 
         self.forward = self.planner.plan_fft_forward(self.window_size);
         self.inverse = self.planner.plan_fft_inverse(self.window_size);
@@ -85,7 +85,7 @@ impl WindowedRealFft {
     pub fn push_back_input(&mut self, value: f32) -> bool {
         self.input.push_back(value);
 
-        self.input.len() >= self.window_size
+        self.input.len() >= self.window_function.needed()
     }
 
     pub fn pop_front_output(&mut self) -> f32 {
@@ -97,7 +97,7 @@ impl WindowedRealFft {
     }
 
     pub fn forward(&mut self) {
-        self.window_function.apply(self.input.make_contiguous());
+        self.window_function.apply(&mut self.input);
 
         let _ = self.forward.process_with_scratch(
             self.input.make_contiguous(),
@@ -117,7 +117,7 @@ impl WindowedRealFft {
 
         let window_size_f32 = self.window_size as f32;
         for sample in &mut self.output {
-            *sample = *sample / window_size_f32;
+            *sample /= window_size_f32;
         }
 
         self.window_function.reverse(self.output.make_contiguous());