Initial commit, basic working version

.gitignore

@@ -0,0 +1 @@
+/target

Cargo.toml

@@ -0,0 +1,17 @@
+[package]
+name = "aa"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+apodize = "1.0.0"
+cpal = { version = "0.15.2" }
+eframe = { version = "0.23.0" }
+egui = { version = "0.23.0", features = ["puffin"] }
+egui_plot = "0.23.0"
+num = "0.4.1"
+realfft = "3.3.0"
+ringbuf = "0.3.3"
+rustfft = "6.1.0"

src/analyzer.rs

@@ -0,0 +1,236 @@
+use num::complex::ComplexFloat;
+use num::integer::Roots;
+use num::traits::real::Real;
+use num::{Float, FromPrimitive, Signed};
+use realfft::{RealFftPlanner, RealToComplex};
+use ringbuf::{HeapRb, Rb};
+use rustfft::Fft;
+use rustfft::{num_complex::Complex};
+
+use std::iter::{repeat};
+
+
+
+use std::{
+    fmt::Debug,
+    sync::{Arc, Mutex},
+};
+
+use std::f64::consts::SQRT_2;
+
+type FloatType = f32;
+type SampleType = f32;
+
+pub trait Analyzer {
+    /// Process some data, and return whether the analysis was updated as a result
+    fn process_data(&mut self, data: &[SampleType]) -> bool;
+    fn set_samplerate(&mut self, rate: FloatType);
+    fn set_length(&mut self, length: usize);
+    fn set_channels(&mut self, channels: usize);
+}
+
+#[derive(Debug, Clone, Default)]
+pub struct FftAnalysis<T> {
+    pub signal_frequency: T,
+    pub signal_power: T,
+    pub total_power: T,
+    pub noise_power: T,
+    pub thd_plus_noise: T,
+    pub max: T,
+    pub min: T,
+    pub fft: Vec<(T, T, T)>, // Frequency, Magnitude, Energy
+}
+
+struct AudioBuf<T> {
+    samples: Vec<HeapRb<T>>,
+    read_since_last_update: usize,
+}
+
+pub trait WindowFunction<T: Real> {
+    fn amplitude_correction_factor(&self) -> T;
+    fn energy_correction_factor(&self) -> T;
+    fn apply(&self, buf: &mut [T]);
+}
+
+struct HanningWindow<T: Real>(Vec<T>);
+
+impl<T: Real + FromPrimitive> HanningWindow<T> {
+    fn new(fft_length: usize) -> Self {
+        Self(
+            apodize::hanning_iter(fft_length)
+                .map(|x| T::from_f64(x).unwrap())
+                .collect(),
+        )
+    }
+}
+
+impl<T: Real + FromPrimitive> WindowFunction<T> for HanningWindow<T> {
+    fn amplitude_correction_factor(&self) -> T {
+        T::from_f64(2.0).unwrap()
+    }
+    fn energy_correction_factor(&self) -> T {
+        T::from_f64(2.0 / num::Float::sqrt(3.0 / 2.0)).unwrap()
+    }
+    fn apply(&self, buf: &mut [T]) {
+        for (a, b) in buf.iter_mut().zip(self.0.iter()) {
+            *a = *a * *b
+        }
+    }
+}
+
+pub fn lin_to_db(lin: f64) -> f64 {
+    20.0 * lin.log10()
+}
+
+pub fn lin_to_dbfs(lin: f64) -> f64 {
+    20.0 * (lin * SQRT_2).log10()
+}
+
+pub struct FftAnalyzer {
+    pub last_analysis: Arc<Mutex<Vec<FftAnalysis<SampleType>>>>,
+
+    fft_length: usize,
+    sample_rate: FloatType,
+    channels: usize,
+    norm_factor: FloatType,
+    bin_freqs: Vec<FloatType>,
+
+    planner: RealFftPlanner<FloatType>,
+    processor: Arc<dyn RealToComplex<FloatType>>,
+    window: HanningWindow<FloatType>,
+
+    audio_buf: Arc<Mutex<AudioBuf<FloatType>>>,
+}
+
+impl FftAnalyzer {
+    pub fn new(fft_length: usize, sample_rate: FloatType, channels: usize) -> Self {
+        let mut planner = RealFftPlanner::new();
+        let bin_freqs = (0..fft_length)
+            .map(|x| {
+                FloatType::from_f64(x as f64).unwrap() * sample_rate
+                    / FloatType::from_f64(fft_length as f64).unwrap()
+            })
+            .collect();
+        Self {
+            last_analysis: Arc::new(Mutex::new(Vec::with_capacity(channels))),
+            fft_length,
+            sample_rate,
+            channels,
+            norm_factor: FloatType::from_f64(1.0 / (fft_length as f64).sqrt()).unwrap(),
+            bin_freqs,
+
+            processor: planner.plan_fft_forward(fft_length),
+            planner,
+            window: HanningWindow::new(fft_length),
+            audio_buf: Arc::new(Mutex::new(AudioBuf::<SampleType> {
+                samples: repeat(0)
+                    .take(channels)
+                    .map(|_| HeapRb::new(fft_length))
+                    .collect(),
+                read_since_last_update: 0,
+            })),
+        }
+    }
+    fn time_domain_process(&mut self, buf: &Vec<FloatType>) -> (FloatType, FloatType, FloatType) {
+        let mut rms_accumulator: FloatType = 0.0;
+        let mut local_max: FloatType = 0.0;
+        let mut local_min: FloatType = 0.0;
+
+        for sample in buf {
+            rms_accumulator += sample * sample;
+            if *sample > local_max {
+                local_max = *sample
+            }
+            if *sample < local_min {
+                local_min = *sample
+            }
+        }
+
+        let total_power = (rms_accumulator / buf.len() as f32).sqrt();
+
+        (local_min, local_max, total_power)
+    }
+    fn process_frame(&mut self) {
+        let _channels_float = FloatType::from(self.channels as f32);
+        let mut result = Vec::with_capacity(self.channels);
+
+        for chan in 0..self.channels {
+            let mut in_buf: Vec<FloatType> = self.audio_buf.lock().unwrap().samples[chan]
+                .iter()
+                .map(|x| *x)
+                .collect();
+
+            let td_result = self.time_domain_process(&in_buf);
+
+            // apply window
+            self.window.apply(&mut in_buf);
+
+            let mut out_buf =
+                Vec::from_iter(repeat(Complex { re: 0.0, im: 0.0 }).take(self.fft_length / 2 + 1));
+
+            // do fft
+            self.processor
+                .process(&mut in_buf, out_buf.as_mut_slice())
+                .expect("fft failed?");
+
+            let analysis_buf = out_buf
+                .iter()
+                .enumerate()
+                .map(|(bin, complex)| {
+                    // get absolute value and normalize based on length
+                    let raw_mag = (complex
+                        / FloatType::from_f64(self.fft_length as f64 / 2.0).unwrap())
+                    .abs();
+                    (
+                        self.bin_freqs[bin],
+                        raw_mag * self.window.amplitude_correction_factor(),
+                        raw_mag * self.window.energy_correction_factor(),
+                    )
+                })
+                .collect();
+            result.push(FftAnalysis {
+                signal_frequency: 0.0,
+                signal_power: 0.0,
+                total_power: td_result.2,
+                noise_power: 0.0,
+                thd_plus_noise: 0.0,
+                max: td_result.1,
+                min: td_result.0,
+                fft: analysis_buf,
+            });
+        }
+
+        let mut lock = self.last_analysis.lock().unwrap();
+        *lock = result;
+    }
+}
+
+impl Analyzer for FftAnalyzer {
+    fn process_data(&mut self, data: &[SampleType]) -> bool {
+        let _channels_float = FloatType::from(self.channels as f32);
+
+        let mut buf = self.audio_buf.lock().unwrap();
+        for chan in 0..self.channels {
+            buf.samples[chan].push_iter_overwrite(data.iter().skip(chan).step_by(self.channels).map(|x| *x));
+        }
+        buf.read_since_last_update += data.len() / self.channels;
+
+        if buf.samples[0].len() >= self.fft_length && buf.read_since_last_update >= self.fft_length / 4 {
+            buf.read_since_last_update = 0;
+            drop(buf);
+            self.process_frame();
+            true
+        } else {
+            false
+        }
+    }
+    fn set_channels(&mut self, channels: usize) {
+        self.channels = channels;
+    }
+    fn set_length(&mut self, _length: usize) {
+        unimplemented!(); // Need to rebuild the window and plan
+    }
+    fn set_samplerate(&mut self, rate: FloatType) {
+        self.sample_rate = rate;
+    }
+}

src/main.rs

@@ -0,0 +1,432 @@
+use cpal::traits::{DeviceTrait, HostTrait, StreamTrait};
+use cpal::SampleFormat;
+
+use egui::epaint::Hsva;
+use egui::Color32;
+use egui_plot::{Legend, Line, Plot};
+
+use std::mem::swap;
+use std::sync::mpsc::{channel, Receiver, Sender};
+use std::sync::{Arc, Mutex};
+use std::thread::{self, JoinHandle};
+
+mod analyzer;
+use analyzer::{lin_to_db, Analyzer, FftAnalysis, FftAnalyzer};
+
+const SAMPLE_RATES: &[u32] = &[
+    11025, 22050, 24000, 32000, 44100, 48000, 88200, 96000, 176400, 192000,
+];
+const CHANNEL_NAMES: &[&str] = &["Left", "Right"];
+const CHANNEL_COLOURS: &[Hsva] = &[
+    Hsva {
+        h: 0.618,
+        s: 0.85,
+        v: 0.5,
+        a: 0.75,
+    },
+    Hsva {
+        h: 0.0,
+        s: 0.85,
+        v: 0.5,
+        a: 0.75,
+    },
+];
+const DEFAULT_SAMPLE_RATE: u32 = 48000;
+
+const AUDIO_FRAME_SIZE: cpal::BufferSize = cpal::BufferSize::Fixed(512);
+
+type SampleType = f32;
+const FFT_LENGTHS: &[usize] = &[512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072];
+const DEFAULT_FFT_LENGTH: usize = 1;
+
+#[derive(Default)]
+struct FftPlotData {
+    plot_points: Vec<Vec<[f64; 2]>>,
+    plot_min: f64,
+    plot_max: f64,
+}
+
+#[derive(Clone, Default)]
+struct FftResult<T> {
+    fft_analyses: Arc<Mutex<Vec<FftAnalysis<T>>>>,
+    plot_data: Arc<Mutex<FftPlotData>>,
+}
+
+fn ui_plot(ui: &mut egui::Ui, last_fft_result: FftResult<f32>) {
+    let plot_data = last_fft_result.plot_data.lock().unwrap();
+    let plot_min = plot_data.plot_min.min(-120.0);
+    let plot_max = plot_data.plot_max.max(0.0);
+    let _plot = Plot::new("FFT")
+        .allow_drag(false)
+        .include_y(plot_min)
+        .include_y(plot_max)
+        .set_margin_fraction(egui::vec2(0.0, 0.0))
+        .x_axis_label("Frequency (KHz)")
+        .y_axis_label("Magnitude (dBFS)")
+        .x_axis_formatter(|x, _n, _r| format!("{:.0}", x / 1000.0))
+        .legend(Legend::default())
+        .show(ui, |plot_ui| {
+            for (chan, chart) in plot_data.plot_points.iter().enumerate() {
+                plot_ui.line(
+                    Line::new(chart.to_vec())
+                        .fill(plot_min as f32)
+                        .name(CHANNEL_NAMES[chan])
+                        .color(CHANNEL_COLOURS[chan]),
+                );
+            }
+        });
+}
+
+#[derive(Default)]
+struct MyAppUiSelections {
+    audio_device: usize,
+    sample_rate: usize,
+    fft_length: usize,
+}
+
+struct MyApp {
+    egui_ctx: egui::Context,
+
+    fft_thread: JoinHandle<()>,
+    fft_close_channel: Sender<()>,
+    fft_length: usize,
+
+    audio_device: cpal::Device,
+    audio_config: cpal::StreamConfig,
+
+    audio_devices: Vec<String>,
+    sample_rates: Vec<u32>,
+    audio_stream: Option<cpal::Stream>,
+    sample_channel: Sender<Vec<SampleType>>, // Store this here so we can restart the audio stream
+    last_result: FftResult<f32>,
+
+    ui_selections: MyAppUiSelections,
+}
+
+impl eframe::App for MyApp {
+    fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
+        egui::SidePanel::left("left_panel").show(ctx, |ui| {
+            egui::ComboBox::from_label("Source")
+                .selected_text(self.audio_devices[self.ui_selections.audio_device].clone())
+                .show_ui(ui, |ui| {
+                    self.audio_devices
+                        .iter()
+                        .enumerate()
+                        .for_each(|(index, dev)| {
+                            ui.selectable_value(&mut self.ui_selections.audio_device, index, dev);
+                        });
+                });
+            self.sample_rate_box(ui);
+            self.fft_length_box(ui);
+        });
+        egui::CentralPanel::default().show(ctx, |ui| {
+            ui_plot(ui, self.last_result.clone());
+        });
+    }
+}
+
+impl MyApp {
+    fn sample_rate_box(&mut self, ui: &mut egui::Ui) {
+        let sample_rate_box = egui::ComboBox::from_label("Sample Rate")
+            .selected_text(self.sample_rates[self.ui_selections.sample_rate].to_string());
+
+        sample_rate_box.show_ui(ui, |ui| {
+            self.sample_rates
+                .iter()
+                .enumerate()
+                .for_each(|(index, rate)| {
+                    ui.selectable_value(
+                        &mut self.ui_selections.sample_rate,
+                        index,
+                        rate.to_string(),
+                    );
+                });
+        });
+        if self.sample_rates[self.ui_selections.sample_rate] != self.audio_config.sample_rate.0 {
+            self.set_sample_rate(self.sample_rates[self.ui_selections.sample_rate]);
+        }
+    }
+    fn fft_length_box(&mut self, ui: &mut egui::Ui) {
+        let fft_length_box = egui::ComboBox::from_label("FFT Length")
+            .selected_text(FFT_LENGTHS[self.ui_selections.fft_length].to_string());
+
+        fft_length_box.show_ui(ui, |ui| {
+            for (index, length) in FFT_LENGTHS.iter().enumerate() {
+                ui.selectable_value(
+                    &mut self.ui_selections.fft_length,
+                    index,
+                    length.to_string(),
+                );
+            }
+        });
+        if FFT_LENGTHS[self.ui_selections.fft_length] != self.fft_length {
+            self.set_fft_length(FFT_LENGTHS[self.ui_selections.fft_length]);
+        }
+    }
+    fn new(cc: &eframe::CreationContext) -> Self {
+        let host = cpal::default_host();
+        println!("Audio host: {}", host.id().name());
+
+        let (sample_tx, sample_rx) = channel::<Vec<SampleType>>();
+
+        let device_names: Vec<_> = host
+            .input_devices()
+            .expect("No available input devices")
+            .map(|x| x.name().unwrap())
+            .collect();
+        let default_device = host.default_input_device().unwrap();
+        let mut supported_configs = default_device.supported_input_configs().unwrap();
+        let supported_sample_rates: Vec<_> = SAMPLE_RATES
+            .iter()
+            .filter(|x| {
+                supported_configs
+                    .any(|y| **x <= y.max_sample_rate().0 && **x >= y.min_sample_rate().0)
+            })
+            .map(|x| *x)
+            .collect();
+
+        let sample_rate_idx = supported_sample_rates
+            .iter()
+            .position(|x| *x == DEFAULT_SAMPLE_RATE)
+            .unwrap_or(0);
+
+        let audio_config = cpal::StreamConfig {
+            channels: 2,
+            sample_rate: cpal::SampleRate(supported_sample_rates[sample_rate_idx]),
+            buffer_size: AUDIO_FRAME_SIZE,
+        };
+
+        let fft_thread_ctx = cc.egui_ctx.clone();
+        let last_result = FftResult::<f32>::default();
+        let fft_thread_result = last_result.clone();
+
+        let fft_close_channel = channel::<()>();
+
+        let fft_analyzer = FftAnalyzer::new(
+            FFT_LENGTHS[DEFAULT_FFT_LENGTH],
+            supported_sample_rates[sample_rate_idx] as f32,
+            2,
+        );
+
+        Self {
+            egui_ctx: cc.egui_ctx.clone(),
+            fft_thread: thread::Builder::new()
+                .name("fft".into())
+                .spawn(move || {
+                    fft_thread_impl(
+                        fft_analyzer,
+                        fft_thread_ctx,
+                        sample_rx,
+                        fft_close_channel.1,
+                        fft_thread_result,
+                    )
+                })
+                .unwrap(),
+            fft_length: FFT_LENGTHS[DEFAULT_FFT_LENGTH],
+            fft_close_channel: fft_close_channel.0,
+            ui_selections: MyAppUiSelections {
+                audio_device: device_names
+                    .iter()
+                    .position(|x| *x == default_device.name().unwrap())
+                    .unwrap_or(0),
+                sample_rate: sample_rate_idx,
+                fft_length: DEFAULT_FFT_LENGTH,
+            },
+            sample_rates: supported_sample_rates,
+            audio_device: default_device,
+            audio_stream: None,
+            audio_devices: device_names,
+            last_result,
+            sample_channel: sample_tx,
+            audio_config,
+        }
+    }
+    fn run(&mut self) {
+        self.audio_thread();
+    }
+    fn restart_threads(&mut self) {
+        let sample_channel = channel::<Vec<SampleType>>();
+        let close_channel = channel::<()>();
+
+        // Stop the audio first so it doesn't try to send to a closed channel
+        self.audio_stream = None;
+
+        // Stop fft thread, rebuild it, swap with the old and join on the old thread
+        self.fft_close_channel.send(()).unwrap();
+        let mut new_thread = self.fft_thread(sample_channel.1, close_channel.1);
+        swap(&mut self.fft_thread, &mut new_thread);
+        new_thread.join().unwrap();
+
+        // Replace the channels in self
+        self.fft_close_channel = close_channel.0;
+        self.sample_channel = sample_channel.0;
+
+        // Restart the audio stream
+        self.audio_thread();
+    }
+    fn set_sample_rate(&mut self, rate: u32) {
+        match self.sample_rates.iter().position(|x| *x == rate) {
+            Some(rate_idx) => {
+                self.ui_selections.sample_rate = rate_idx;
+                self.audio_config.sample_rate = cpal::SampleRate(SAMPLE_RATES[rate_idx]);
+                self.restart_threads();
+            }
+            None => unimplemented!("Make some graceful error handler"),
+        }
+    }
+    fn set_fft_length(&mut self, length: usize) {
+        match FFT_LENGTHS.iter().position(|x| *x == length) {
+            Some(length_idx) => {
+                self.fft_length = length;
+                self.ui_selections.fft_length = length_idx;
+                self.restart_threads();
+            }
+            None => unimplemented!("Make some graceful error handler"),
+        }
+    }
+    fn audio_thread(&mut self) {
+        if self.audio_stream.is_some() {
+            self.audio_stream = None;
+        }
+        println!(
+            "Starting audio thread on device `{}` with config `{:?}`",
+            self.audio_device.name().unwrap(),
+            self.audio_config
+        );
+
+        let new_channel = self.sample_channel.clone();
+
+        self.audio_stream = Some(
+            self.audio_device
+                .build_input_stream(
+                    &self.audio_config,
+                    move |data, _| handle_audio_frame(data, new_channel.clone()),
+                    move |err| panic!("{:?}", err),
+                    Some(std::time::Duration::from_secs(1)),
+                )
+                .unwrap(),
+        );
+        println!("Stream built");
+        self.audio_stream.as_ref().unwrap().play().unwrap(); // this doesn't block
+    }
+
+    fn fft_thread(
+        &mut self,
+        data_pipe: Receiver<Vec<SampleType>>,
+        close_pipe: Receiver<()>,
+    ) -> JoinHandle<()> {
+        let fft_analyzer = FftAnalyzer::new(
+            FFT_LENGTHS[self.ui_selections.fft_length],
+            self.sample_rates[self.ui_selections.sample_rate] as f32,
+            2,
+        );
+        let result_ref = self.last_result.clone();
+        let ctx_ref = self.egui_ctx.clone();
+        thread::Builder::new()
+            .name("fft".into())
+            .spawn(move || {
+                fft_thread_impl(fft_analyzer, ctx_ref, data_pipe, close_pipe, result_ref)
+            })
+            .unwrap()
+    }
+}
+
+fn plot_min(data: &Vec<Vec<[f64; 2]>>) -> f64 {
+    let mut min = f64::INFINITY;
+    for channel in data {
+        for point in channel {
+            if point[1] < min && point[1] > f64::NEG_INFINITY {
+                min = point[1];
+            }
+        }
+    }
+    min.clamp(-160.0, 160.0)
+}
+
+fn plot_max(data: &Vec<Vec<[f64; 2]>>) -> f64 {
+    let mut max = f64::NEG_INFINITY;
+    for channel in data {
+        for point in channel {
+            if point[1] > max && point[1] < f64::INFINITY {
+                max = point[1];
+            }
+        }
+    }
+    max.clamp(-160.0, 160.0)
+}
+
+fn fft_thread_impl(
+    mut analyzer: FftAnalyzer,
+    ctx: egui::Context,
+    data_pipe: Receiver<Vec<SampleType>>,
+    close_pipe: Receiver<()>,
+    last_result: FftResult<f32>,
+) {
+    println!("FFT thread initialized");
+    loop {
+        if close_pipe.try_recv().is_ok() {
+            println!("FFT thread ended");
+            return;
+        }
+        if let Ok(buf) = data_pipe.recv_timeout(std::time::Duration::from_secs_f64(0.1)) {
+            if analyzer.process_data(&buf) {
+                // Prepare the data for plotting
+                let charts: Vec<Vec<[f64; 2]>> =
+                    analyzer
+                        .last_analysis
+                        .lock()
+                        .unwrap()
+                        .iter()
+                        .map(|analysis| {
+                            Vec::from_iter(analysis.fft.iter().map(|(freq, amp, _energy)| {
+                                [(*freq as f64), lin_to_db(*amp as f64)]
+                            }))
+                        })
+                        .collect();
+
+                let plot_min = plot_min(&charts);
+                let plot_max = plot_max(&charts);
+
+                let mut plot_data = last_result.plot_data.lock().unwrap();
+                *plot_data = FftPlotData {
+                    plot_points: charts,
+                    plot_min,
+                    plot_max,
+                };
+                // while holding the plot_lock which will block the gui thread, update the rest
+                *last_result.fft_analyses.lock().unwrap() =
+                    analyzer.last_analysis.lock().unwrap().clone();
+                drop(plot_data);
+                ctx.request_repaint();
+            }
+        } else {
+            println!("Didn't receive any data, looping around");
+        }
+    }
+}
+
+fn handle_audio_frame(data: &[SampleType], chan: Sender<Vec<SampleType>>) {
+    // Copy and send to processor thread
+    // FIXME: Fails at fft lengths > 16384. Are we too slow collecting data?
+    chan.send(data.to_vec()).unwrap_or_default()
+}
+
+fn main() {
+    let options = eframe::NativeOptions {
+        initial_window_size: Some(egui::vec2(640.0, 480.0)),
+        ..Default::default()
+    };
+
+    println!("App initialized");
+
+    eframe::run_native(
+        "rust-aa",
+        options,
+        Box::new(|cc| {
+            let mut app = Box::new(MyApp::new(cc));
+            app.run();
+            app
+        }),
+    )
+    .unwrap();
+}