From 179ff6a035281c7afd4939125e3e51dd087834d3 Mon Sep 17 00:00:00 2001
From: Robbert van der Helm <mail@robbertvanderhelm.nl>
Date: Wed, 9 Nov 2022 17:44:15 +0100
Subject: [PATCH] Add automatic normalization to Buffr Glitch

---
 plugins/buffr_glitch/src/buffer.rs | 37 +++++++++++++++++++++++++++++-
 plugins/buffr_glitch/src/lib.rs    | 28 ++++++++++++++++++----
 2 files changed, 60 insertions(+), 5 deletions(-)

diff --git a/plugins/buffr_glitch/src/buffer.rs b/plugins/buffr_glitch/src/buffer.rs
index 215da35a..704e3c45 100644
--- a/plugins/buffr_glitch/src/buffer.rs
+++ b/plugins/buffr_glitch/src/buffer.rs
@@ -16,6 +16,8 @@
 
 use nih_plug::prelude::*;
 
+use crate::NormalizationMode;
+
 /// A super simple ring buffer abstraction that records audio into a recording ring buffer, and then
 /// copies audio to a playback buffer when a note is pressed so audio can be repeated while still
 /// recording audio for further key presses. This needs to be able to store at least the number of
@@ -94,7 +96,7 @@ impl RingBuffer {
 
     /// Prepare the playback buffers to play back audio at the specified frequency. This copies
     /// audio from the ring buffers to the playback buffers.
-    pub fn prepare_playback(&mut self, frequency: f32) {
+    pub fn prepare_playback(&mut self, frequency: f32, normalization_mode: NormalizationMode) {
         let note_period_samples = (frequency.recip() * self.sample_rate).ceil() as usize;
 
         // We'll copy the last `note_period_samples` samples from the recording ring buffers to the
@@ -117,6 +119,22 @@ impl RingBuffer {
                 .copy_from_slice(&recording_buffer[..copy_num_from_start]);
         }
 
+        // The playback buffer is normalized as necessary. This prevents small grains from being
+        // either way quieter or way louder than the origianl audio.
+        match normalization_mode {
+            NormalizationMode::None => (),
+            NormalizationMode::Auto => {
+                let normalization_factor =
+                    calculate_rms(&self.recording_buffers) / calculate_rms(&self.playback_buffers);
+
+                for buffer in self.playback_buffers.iter_mut() {
+                    for sample in buffer.iter_mut() {
+                        *sample *= normalization_factor;
+                    }
+                }
+            }
+        }
+
         // Reading from the buffer should always start at the beginning
         self.playback_buffer_pos = 0;
     }
@@ -139,3 +157,20 @@ impl RingBuffer {
         sample
     }
 }
+
+/// Get the RMS value of an entire buffer. This is used for (automatic) normalization.
+///
+/// # Panics
+///
+/// This will panic of `buffers` is empty.
+fn calculate_rms(buffers: &[Vec<f32>]) -> f32 {
+    nih_debug_assert_ne!(buffers.len(), 0);
+
+    let sum_of_squares: f32 = buffers
+        .iter()
+        .map(|buffer| buffer.iter().map(|sample| (sample * sample)).sum::<f32>())
+        .sum();
+    let num_samples = buffers.len() * buffers[0].len();
+
+    (sum_of_squares / num_samples as f32).sqrt()
+}
diff --git a/plugins/buffr_glitch/src/lib.rs b/plugins/buffr_glitch/src/lib.rs
index 49bf885b..637b265a 100644
--- a/plugins/buffr_glitch/src/lib.rs
+++ b/plugins/buffr_glitch/src/lib.rs
@@ -33,9 +33,24 @@ struct BuffrGlitch {
     midi_note_id: Option<u8>,
 }
 
-// TODO: Normalize option
 #[derive(Params)]
-struct BuffrGlitchParams {}
+struct BuffrGlitchParams {
+    /// Controls if and how grains are normalization.
+    #[id = "normalization_mode"]
+    normalization_mode: EnumParam<NormalizationMode>,
+}
+
+/// Controls how grains are normalized.
+#[derive(Enum, Debug, PartialEq, Eq)]
+pub enum NormalizationMode {
+    /// Don't normalize at all
+    #[id = "none"]
+    None,
+    /// Automatically normalize based on the recording buffer's RMS value.
+    #[id = "auto"]
+    Auto,
+    // TODO: Explicit RMS target
+}
 
 impl Default for BuffrGlitch {
     fn default() -> Self {
@@ -52,7 +67,9 @@ impl Default for BuffrGlitch {
 
 impl Default for BuffrGlitchParams {
     fn default() -> Self {
-        Self {}
+        Self {
+            normalization_mode: EnumParam::new("Normalization", NormalizationMode::Auto),
+        }
     }
 }
 
@@ -124,7 +141,10 @@ impl Plugin for BuffrGlitch {
 
                         // We'll copy audio to the playback buffer to match the pitch of the note
                         // that was just played
-                        self.buffer.prepare_playback(util::midi_note_to_freq(note));
+                        self.buffer.prepare_playback(
+                            util::midi_note_to_freq(note),
+                            self.params.normalization_mode.value(),
+                        );
                     }
                     NoteEvent::NoteOff { note, .. } if self.midi_note_id == Some(note) => {
                         // A NoteOff for the currently playing note immediately ends playback