#[cfg(bela_device)]
use core::ffi::c_int;
use core::fmt;
#[cfg(bela_device)]
use core::ptr::NonNull;
use crate::error::Error;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct FftLength(usize);
impl FftLength {
pub const MIN: Self = Self(8);
pub const MAX: Self = Self(65536);
#[must_use]
pub const fn new(length: usize) -> Option<Self> {
if !length.is_power_of_two() || length < Self::MIN.0 || length > Self::MAX.0 {
return None;
}
Some(Self(length))
}
#[must_use]
pub const fn rounded_up(length: usize) -> Option<Self> {
if length <= Self::MIN.0 {
return Some(Self::MIN);
}
if length > Self::MAX.0 {
return None;
}
Some(Self(length.next_power_of_two()))
}
#[must_use]
pub const fn get(self) -> usize {
self.0
}
#[cfg(bela_device)]
#[allow(
clippy::cast_possible_truncation,
clippy::cast_possible_wrap,
reason = "FftLength cannot hold more than MAX, which is 65536"
)]
const fn as_nfft(self) -> c_int {
self.0 as c_int
}
#[must_use]
pub const fn spectrum_len(self) -> usize {
self.0 / 2 + 1
}
}
const _: () = assert!(
FftLength::MAX.0 <= 2_147_483_647,
"a transform length has to fit the C int NE10 takes"
);
impl From<FftLength> for usize {
fn from(length: FftLength) -> Self {
length.0
}
}
impl TryFrom<usize> for FftLength {
type Error = Error;
fn try_from(length: usize) -> Result<Self, Error> {
Self::new(length).ok_or(Error::FftLength { value: length })
}
}
impl fmt::Display for FftLength {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}
#[repr(C)]
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct FftBin {
pub re: f32,
pub im: f32,
}
const _: () = {
assert!(
size_of::<FftBin>() == size_of::<bela_sys::ne10_fft_cpx_float32_t>(),
"a bin has to be the size NE10 reads and writes"
);
assert!(
align_of::<FftBin>() == align_of::<bela_sys::ne10_fft_cpx_float32_t>(),
"a bin has to be aligned as NE10 expects"
);
};
impl FftBin {
pub const ZERO: Self = Self { re: 0.0, im: 0.0 };
#[must_use]
pub const fn new(re: f32, im: f32) -> Self {
Self { re, im }
}
#[must_use]
pub fn magnitude(self) -> f32 {
self.magnitude_squared().sqrt()
}
#[must_use]
pub const fn magnitude_squared(self) -> f32 {
self.re * self.re + self.im * self.im
}
#[must_use]
pub fn phase(self) -> f32 {
self.im.atan2(self.re)
}
}
impl From<(f32, f32)> for FftBin {
fn from((re, im): (f32, f32)) -> Self {
Self::new(re, im)
}
}
impl From<[f32; 2]> for FftBin {
fn from([re, im]: [f32; 2]) -> Self {
Self::new(re, im)
}
}
pub struct RealFft {
length: FftLength,
#[cfg(bela_device)]
plan: NonNull<bela_sys::ne10_fft_r2c_state_float32_t>,
}
unsafe impl Send for RealFft {}
unsafe impl Sync for RealFft {}
impl RealFft {
#[cfg(bela_device)]
pub fn new(length: FftLength) -> Result<Self, Error> {
let plan = NonNull::new(unsafe { bela_sys::ne10_fft_alloc_r2c_float32(length.as_nfft()) })
.ok_or_else(|| Error::FftCreate {
length: length.get(),
})?;
Ok(Self { length, plan })
}
#[cfg(not(bela_device))]
#[allow(
clippy::missing_const_for_fn,
reason = "mirrors the device signature, which allocates"
)]
pub fn new(_length: FftLength) -> Result<Self, Error> {
Err(Error::FftUnavailable)
}
#[must_use]
pub const fn length(&self) -> FftLength {
self.length
}
#[must_use]
pub const fn spectrum_len(&self) -> usize {
self.length.spectrum_len()
}
#[must_use]
pub fn new_signal(&self) -> Vec<f32> {
vec![0.0; self.length.get()]
}
#[must_use]
pub fn new_spectrum(&self) -> Vec<FftBin> {
vec![FftBin::ZERO; self.spectrum_len()]
}
pub fn forward(&mut self, signal: &mut [f32], spectrum: &mut [FftBin]) -> Result<(), Error> {
self.check_signal(signal.len())?;
self.check_spectrum(spectrum.len())?;
self.forward_raw(signal, spectrum);
Ok(())
}
pub fn inverse(&mut self, spectrum: &mut [FftBin], signal: &mut [f32]) -> Result<(), Error> {
self.check_spectrum(spectrum.len())?;
self.check_signal(signal.len())?;
self.inverse_raw(spectrum, signal);
Ok(())
}
const fn check_signal(&self, actual: usize) -> Result<(), Error> {
let expected = self.length.get();
if actual == expected {
Ok(())
} else {
Err(Error::FftSignalLen { expected, actual })
}
}
const fn check_spectrum(&self, actual: usize) -> Result<(), Error> {
let expected = self.spectrum_len();
if actual == expected {
Ok(())
} else {
Err(Error::FftSpectrumLen { expected, actual })
}
}
#[cfg(bela_device)]
fn forward_raw(&mut self, signal: &mut [f32], spectrum: &mut [FftBin]) {
unsafe {
bela_sys::ne10_fft_r2c_1d_float32_neon(
spectrum.as_mut_ptr().cast(),
signal.as_mut_ptr(),
self.plan.as_ptr(),
);
}
}
#[cfg(not(bela_device))]
#[allow(
clippy::unused_self,
clippy::needless_pass_by_ref_mut,
clippy::missing_const_for_fn,
reason = "mirrors the device signature, which transforms through the plan"
)]
fn forward_raw(&mut self, _signal: &mut [f32], _spectrum: &mut [FftBin]) {}
#[cfg(bela_device)]
fn inverse_raw(&mut self, spectrum: &mut [FftBin], signal: &mut [f32]) {
unsafe {
bela_sys::ne10_fft_c2r_1d_float32_neon(
signal.as_mut_ptr(),
spectrum.as_mut_ptr().cast(),
self.plan.as_ptr(),
);
}
}
#[cfg(not(bela_device))]
#[allow(
clippy::unused_self,
clippy::needless_pass_by_ref_mut,
clippy::missing_const_for_fn,
reason = "mirrors the device signature, which transforms through the plan"
)]
fn inverse_raw(&mut self, _spectrum: &mut [FftBin], _signal: &mut [f32]) {}
}
#[cfg(all(test, not(bela_device)))]
impl RealFft {
const fn for_test(length: FftLength) -> Self {
Self { length }
}
}
impl fmt::Debug for RealFft {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("RealFft")
.field("length", &self.length)
.finish_non_exhaustive()
}
}
impl Drop for RealFft {
fn drop(&mut self) {
#[cfg(bela_device)]
unsafe {
bela_sys::ne10_fft_destroy_r2c_float32(self.plan.as_ptr());
}
}
}
#[cfg(test)]
mod tests {
use core::f32::consts::FRAC_PI_2;
use super::*;
#[test]
fn a_length_is_a_power_of_two_in_range() {
for length in [8, 16, 1024, 65536] {
assert_eq!(
FftLength::new(length).map(FftLength::get),
Some(length),
"{length} is supported"
);
}
for length in [0, 1, 3, 7, 9, 1000, 131_072] {
assert_eq!(FftLength::new(length), None, "{length} is not");
}
}
#[test]
fn the_lengths_that_corrupt_memory_are_not_lengths() {
assert_eq!(FftLength::new(2), None);
assert_eq!(FftLength::new(4), None);
assert_eq!(FftLength::MIN.get(), 8);
}
#[test]
fn rounding_up_lands_on_a_supported_length() {
for (asked, expected) in [
(0, 8),
(1, 8),
(2, 8),
(5, 8),
(8, 8),
(9, 16),
(1000, 1024),
(65536, 65536),
] {
assert_eq!(
FftLength::rounded_up(asked).map(FftLength::get),
Some(expected),
"rounding {asked} up"
);
}
assert_eq!(FftLength::rounded_up(65537), None);
assert_eq!(FftLength::rounded_up(usize::MAX), None);
}
#[test]
fn a_spectrum_holds_dc_nyquist_and_what_is_between() {
assert_eq!(FftLength::MIN.spectrum_len(), 5);
assert_eq!(
FftLength::new(1024)
.expect("1024 is supported")
.spectrum_len(),
513
);
assert_eq!(FftLength::MAX.spectrum_len(), 32769);
}
#[test]
fn a_length_converts_both_ways() {
let length = FftLength::new(256).expect("256 is supported");
assert_eq!(usize::from(length), 256);
assert_eq!(FftLength::try_from(256_usize), Ok(length));
assert_eq!(
FftLength::try_from(3_usize),
Err(Error::FftLength { value: 3 })
);
assert_eq!(length.to_string(), "256");
}
#[test]
fn a_bin_is_two_floats_laid_out_as_ne10_lays_them() {
assert_eq!(size_of::<FftBin>(), size_of::<f32>() * 2);
assert_eq!(align_of::<FftBin>(), align_of::<f32>());
assert_eq!(
size_of::<FftBin>(),
size_of::<bela_sys::ne10_fft_cpx_float32_t>()
);
assert_eq!(
align_of::<FftBin>(),
align_of::<bela_sys::ne10_fft_cpx_float32_t>()
);
assert_eq!(FftBin::ZERO, FftBin::default());
assert_eq!(FftBin::from((1.0, 2.0)), FftBin::new(1.0, 2.0));
assert_eq!(FftBin::from([1.0, 2.0]), FftBin::new(1.0, 2.0));
}
#[test]
fn a_bin_reports_magnitude_and_phase() {
let bin = FftBin::new(3.0, 4.0);
assert!((bin.magnitude() - 5.0).abs() < 1e-6);
assert!((bin.magnitude_squared() - 25.0).abs() < 1e-6);
assert!((FftBin::new(0.0, 1.0).phase() - FRAC_PI_2).abs() < 1e-6);
assert!(FftBin::ZERO.magnitude().abs() < f32::EPSILON);
}
#[test]
#[cfg(not(bela_device))]
fn a_plan_needs_a_board() {
let length = FftLength::new(64).expect("64 is supported");
assert_eq!(RealFft::new(length).unwrap_err(), Error::FftUnavailable);
}
#[cfg(not(bela_device))]
mod host {
use super::*;
fn plan() -> RealFft {
RealFft::for_test(FftLength::new(64).expect("64 is supported"))
}
#[test]
fn a_plan_reports_what_it_transforms() {
let fft = plan();
assert_eq!(fft.length().get(), 64);
assert_eq!(fft.spectrum_len(), 33);
assert!(
format!("{fft:?}").contains("64"),
"the Debug names the length"
);
}
#[test]
fn the_buffers_a_plan_makes_are_the_ones_it_takes() {
let fft = plan();
let mut signal = fft.new_signal();
let mut spectrum = fft.new_spectrum();
assert_eq!(signal.len(), fft.length().get());
assert_eq!(spectrum.len(), fft.spectrum_len());
assert!(signal.iter().all(|sample| *sample == 0.0));
assert!(spectrum.iter().all(|bin| *bin == FftBin::ZERO));
let mut fft = fft;
assert_eq!(fft.forward(&mut signal, &mut spectrum), Ok(()));
assert_eq!(fft.inverse(&mut spectrum, &mut signal), Ok(()));
}
#[test]
fn a_buffer_of_the_wrong_length_is_refused_and_says_which() {
let mut fft = plan();
let mut signal = fft.new_signal();
let mut spectrum = fft.new_spectrum();
assert_eq!(
fft.forward(&mut signal[..63], &mut spectrum),
Err(Error::FftSignalLen {
expected: 64,
actual: 63
})
);
assert_eq!(
fft.forward(&mut signal, &mut spectrum[..32]),
Err(Error::FftSpectrumLen {
expected: 33,
actual: 32
})
);
assert_eq!(
fft.inverse(&mut spectrum[..32], &mut signal),
Err(Error::FftSpectrumLen {
expected: 33,
actual: 32
})
);
assert_eq!(
fft.inverse(&mut spectrum, &mut signal[..63]),
Err(Error::FftSignalLen {
expected: 64,
actual: 63
})
);
}
#[test]
fn a_longer_buffer_is_refused_as_well() {
let mut fft = plan();
let mut signal = vec![0.0; 65];
let mut spectrum = vec![FftBin::ZERO; 34];
assert!(matches!(
fft.forward(&mut signal, &mut spectrum),
Err(Error::FftSignalLen { actual: 65, .. })
));
assert!(matches!(
fft.inverse(&mut spectrum, &mut signal),
Err(Error::FftSpectrumLen { actual: 34, .. })
));
}
#[test]
fn the_check_order_follows_the_argument_order() {
let mut fft = plan();
let mut signal = vec![0.0; 8];
let mut spectrum = vec![FftBin::ZERO; 8];
assert!(
matches!(
fft.forward(&mut signal, &mut spectrum),
Err(Error::FftSignalLen { .. })
),
"forward takes the signal first"
);
assert!(
matches!(
fft.inverse(&mut spectrum, &mut signal),
Err(Error::FftSpectrumLen { .. })
),
"inverse takes the spectrum first"
);
}
}
}