use crate::mla::{c_mul_fast, fmla};
use crate::util::DctSample;
use num_complex::Complex;
use std::marker::PhantomData;
use std::sync::{Arc, OnceLock};
pub(crate) trait DctSpectrumMul<T> {
fn mul_spectrum_to_real(&self, a: &[Complex<T>], b: &[Complex<T>], out: &mut [T]);
fn mul_spectrum_to_real_rev(&self, a: &[Complex<T>], b: &[Complex<T>], out: &mut [T]);
fn mul_spectrum_and_half(&self, a: &[T], b: &[Complex<T>], out: &mut [Complex<T>]);
fn mul_spectrum_and_half_rev(&self, a: &[T], b: &[Complex<T>], out: &mut [Complex<T>]);
}
pub(crate) struct FftSpectrumMul<T> {
phantom_data: PhantomData<T>,
}
impl<T: DctSample> DctSpectrumMul<T> for FftSpectrumMul<T> {
fn mul_spectrum_to_real(
&self,
complex_input: &[Complex<T>],
twiddles: &[Complex<T>],
out: &mut [T],
) {
let len = out.len();
let half = out.len() / 2;
let complex_length = complex_input.len();
assert!(twiddles.len() >= len);
assert!(!complex_input.is_empty());
assert!(complex_input.len() >= half);
assert!(out.len() >= len);
unsafe {
*out.get_unchecked_mut(0) =
twiddles.get_unchecked(0).re * complex_input.get_unchecked(0).re;
if len.is_multiple_of(2) {
*out.get_unchecked_mut(half) =
twiddles.get_unchecked(half).re * complex_input.get_unchecked(half).re;
}
}
for i in 1..complex_length {
unsafe {
let twiddle = *twiddles.get_unchecked(i);
let twiddle_rev = *twiddles.get_unchecked(len - i);
let fft_value = *complex_input.get_unchecked(i);
*out.get_unchecked_mut(i) =
fmla(fft_value.re, twiddle.re, -fft_value.im * twiddle.im);
let conj_fft = fft_value.conj();
*out.get_unchecked_mut(len - i) =
fmla(conj_fft.re, twiddle_rev.re, -conj_fft.im * twiddle_rev.im);
}
}
}
fn mul_spectrum_to_real_rev(&self, a: &[Complex<T>], b: &[Complex<T>], out: &mut [T]) {
for ((fft_entry, twiddle), out) in a.iter().zip(b.iter()).zip(out.iter_mut().rev()) {
*out = c_mul_fast(*fft_entry, *twiddle).re;
}
}
fn mul_spectrum_and_half(&self, a: &[T], b: &[Complex<T>], out: &mut [Complex<T>]) {
out[0] = Complex::from(a[0] * T::HALF);
for (((entry, twiddle), c_forward), c_backward) in out
.iter_mut()
.zip(b.iter())
.zip(a.iter())
.skip(1)
.zip(a.iter().rev())
{
let c = Complex {
re: *c_forward,
im: *c_backward,
};
*entry = c_mul_fast(c, *twiddle) * T::HALF;
}
}
fn mul_spectrum_and_half_rev(&self, a: &[T], b: &[Complex<T>], out: &mut [Complex<T>]) {
let len_m1 = a.len() - 1;
out[0] = Complex::from(a[len_m1] * T::HALF);
for (((entry, twiddle), c_forward), c_backward) in out
.iter_mut()
.skip(1)
.zip(b.iter().skip(1))
.zip(a.iter())
.zip(a.iter().rev().skip(1))
{
let c = Complex {
re: *c_backward,
im: *c_forward,
};
*entry = c_mul_fast(c, *twiddle) * T::HALF;
}
}
}
pub(crate) trait FftSpectrumMulFactory<T> {
fn create_mul_spectrum_to_real() -> Arc<dyn DctSpectrumMul<T> + Send + Sync>;
}
impl FftSpectrumMulFactory<f32> for f32 {
fn create_mul_spectrum_to_real() -> Arc<dyn DctSpectrumMul<f32> + Send + Sync> {
static IMPL: OnceLock<Arc<dyn DctSpectrumMul<f32> + Send + Sync>> = OnceLock::new();
IMPL.get_or_init(|| {
#[cfg(all(target_arch = "aarch64", feature = "neon"))]
{
#[cfg(feature = "fcma")]
{
if std::arch::is_aarch64_feature_detected!("fcma") {
use crate::neon::FcmaDctSpectrumMulF32;
return Arc::new(FcmaDctSpectrumMulF32 {});
}
}
use crate::neon::DctSpectrumMulF32;
Arc::new(DctSpectrumMulF32 {})
}
#[cfg(all(target_arch = "x86_64", feature = "avx"))]
{
use crate::util::has_valid_avx;
if has_valid_avx() {
use crate::avx::AvxDctSpectrumMulF32;
return Arc::new(AvxDctSpectrumMulF32::default());
}
}
#[cfg(not(all(target_arch = "aarch64", feature = "neon")))]
Arc::new(FftSpectrumMul {
phantom_data: PhantomData,
})
})
.clone()
}
}
impl FftSpectrumMulFactory<f64> for f64 {
fn create_mul_spectrum_to_real() -> Arc<dyn DctSpectrumMul<f64> + Send + Sync> {
Arc::new(FftSpectrumMul {
phantom_data: PhantomData,
})
}
}