use rustfft::{num_complex::Complex, FftPlanner};
use std::sync::Arc;
pub struct SpectrumAnalyzer {
fft_size: usize,
fft: Arc<dyn rustfft::Fft<f64>>,
window: Vec<f64>,
num_bins: usize,
fft_buffer: Vec<Complex<f64>>,
magnitudes: Vec<f64>,
result: Vec<f32>,
bin_ranges: Vec<(usize, usize)>,
bin_sample_rate: Option<u32>,
}
impl SpectrumAnalyzer {
pub fn new(fft_size: usize, num_bins: usize) -> Self {
let mut planner = FftPlanner::new();
let fft = planner.plan_fft_forward(fft_size);
let window: Vec<f64> = (0..fft_size)
.map(|i| 0.5 * (1.0 - (2.0 * std::f64::consts::PI * i as f64 / fft_size as f64).cos()))
.collect();
Self {
fft_size,
fft,
window,
num_bins,
fft_buffer: vec![Complex::new(0.0, 0.0); fft_size],
magnitudes: vec![0.0; fft_size.saturating_div(2).saturating_sub(1)],
result: vec![0.0; num_bins],
bin_ranges: Vec::with_capacity(num_bins),
bin_sample_rate: None,
}
}
pub fn analyze(&mut self, samples: &[f64], sample_rate: u32) -> &[f32] {
if samples.len() < self.fft_size {
self.result.fill(0.0);
return &self.result;
}
for ((slot, &sample), &window) in self
.fft_buffer
.iter_mut()
.zip(samples.iter().take(self.fft_size))
.zip(&self.window)
{
*slot = Complex::new(sample * window, 0.0);
}
self.fft.process(&mut self.fft_buffer);
for (dst, c) in self
.magnitudes
.iter_mut()
.zip(self.fft_buffer[1..self.fft_size / 2].iter())
{
*dst = c.norm() / self.fft_size as f64;
}
self.ensure_bin_ranges(sample_rate);
self.log_bin();
&self.result
}
fn ensure_bin_ranges(&mut self, sample_rate: u32) {
if self.bin_sample_rate == Some(sample_rate) && self.bin_ranges.len() == self.num_bins {
return;
}
let nyquist = sample_rate as f64 / 2.0;
let min_freq = 20.0f64;
let max_freq = nyquist;
let log_min = min_freq.log10();
let log_max = max_freq.log10();
let freq_per_bin = nyquist / self.magnitudes.len().max(1) as f64;
self.bin_ranges.clear();
for bin_idx in 0..self.num_bins {
let freq_low = 10.0_f64
.powf(log_min + (log_max - log_min) * bin_idx as f64 / self.num_bins as f64);
let freq_high = 10.0_f64
.powf(log_min + (log_max - log_min) * (bin_idx + 1) as f64 / self.num_bins as f64);
let idx_low = ((freq_low / freq_per_bin) as usize)
.clamp(0, self.magnitudes.len().saturating_sub(1));
let idx_high =
((freq_high / freq_per_bin) as usize).clamp(idx_low + 1, self.magnitudes.len());
self.bin_ranges.push((idx_low, idx_high));
}
self.bin_sample_rate = Some(sample_rate);
}
fn log_bin(&mut self) {
self.result.fill(0.0);
for (result_val, &(idx_low, idx_high)) in self.result.iter_mut().zip(&self.bin_ranges) {
if idx_high > idx_low {
let sum: f64 = self.magnitudes[idx_low..idx_high]
.iter()
.map(|m| m * m)
.sum();
let rms = (sum / (idx_high - idx_low) as f64).sqrt();
let db = 20.0 * (rms + 1e-9).log10();
*result_val = ((db + 90.0) / 90.0).clamp(0.0, 1.0) as f32;
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use rustfft::FftPlanner;
#[test]
fn short_input_returns_reused_zero_bins() {
let mut analyzer = SpectrumAnalyzer::new(16, 4);
let first_ptr = analyzer.analyze(&[0.0; 8], 48_000).as_ptr();
assert_eq!(analyzer.analyze(&[0.0; 8], 48_000), &[0.0; 4]);
assert_eq!(analyzer.analyze(&[0.0; 8], 48_000).as_ptr(), first_ptr);
}
#[test]
fn analyze_reuses_result_and_recomputes_ranges_on_sample_rate_change() {
let mut analyzer = SpectrumAnalyzer::new(64, 8);
let samples: Vec<f64> = (0..64).map(|i| (i as f64 * 0.1).sin()).collect();
let first_ptr = analyzer.analyze(&samples, 48_000).as_ptr();
let first_ranges = analyzer.bin_ranges.clone();
assert!(analyzer.analyze(&samples, 48_000).iter().any(|&v| v > 0.0));
assert_eq!(analyzer.analyze(&samples, 48_000).as_ptr(), first_ptr);
assert_eq!(analyzer.bin_ranges, first_ranges);
analyzer.analyze(&samples, 96_000);
assert_ne!(analyzer.bin_ranges, first_ranges);
}
#[test]
fn analyzer_output_matches_legacy_allocation_path() {
let mut analyzer = SpectrumAnalyzer::new(128, 16);
let samples: Vec<f64> = (0..128)
.map(|i| {
let t = i as f64 / 48_000.0;
(2.0 * std::f64::consts::PI * 997.0 * t).sin() * 0.4
})
.collect();
let actual = analyzer.analyze(&samples, 48_000).to_vec();
let expected = legacy_analyze(&samples, 128, 16, 48_000);
for (idx, (actual, expected)) in actual.iter().zip(expected.iter()).enumerate() {
assert!(
(actual - expected).abs() <= 1e-6,
"bin {idx}: actual={actual}, expected={expected}"
);
}
}
fn legacy_analyze(
samples: &[f64],
fft_size: usize,
num_bins: usize,
sample_rate: u32,
) -> Vec<f32> {
if samples.len() < fft_size {
return vec![0.0; num_bins];
}
let mut planner = FftPlanner::new();
let fft = planner.plan_fft_forward(fft_size);
let window: Vec<f64> = (0..fft_size)
.map(|i| 0.5 * (1.0 - (2.0 * std::f64::consts::PI * i as f64 / fft_size as f64).cos()))
.collect();
let mut buffer: Vec<Complex<f64>> = samples[..fft_size]
.iter()
.zip(&window)
.map(|(&s, &w)| Complex::new(s * w, 0.0))
.collect();
fft.process(&mut buffer);
let magnitudes: Vec<f64> = buffer[1..fft_size / 2]
.iter()
.map(|c| c.norm() / fft_size as f64)
.collect();
legacy_log_bin(&magnitudes, sample_rate, num_bins)
}
fn legacy_log_bin(magnitudes: &[f64], sample_rate: u32, num_bins: usize) -> Vec<f32> {
let mut result = vec![0.0f32; num_bins];
let nyquist = sample_rate as f64 / 2.0;
let min_freq = 20.0f64;
let max_freq = nyquist;
let log_min = min_freq.log10();
let log_max = max_freq.log10();
for (bin_idx, result_val) in result.iter_mut().enumerate() {
let freq_low =
10.0_f64.powf(log_min + (log_max - log_min) * bin_idx as f64 / num_bins as f64);
let freq_high = 10.0_f64
.powf(log_min + (log_max - log_min) * (bin_idx + 1) as f64 / num_bins as f64);
let freq_per_bin = nyquist / magnitudes.len() as f64;
let idx_low =
((freq_low / freq_per_bin) as usize).clamp(0, magnitudes.len().saturating_sub(1));
let idx_high =
((freq_high / freq_per_bin) as usize).clamp(idx_low + 1, magnitudes.len());
if idx_high > idx_low {
let sum: f64 = magnitudes[idx_low..idx_high].iter().map(|m| m * m).sum();
let rms = (sum / (idx_high - idx_low) as f64).sqrt();
let db = 20.0 * (rms + 1e-9).log10();
*result_val = ((db + 90.0) / 90.0).clamp(0.0, 1.0) as f32;
}
}
result
}
}