1use crate::{common::FftNum, FftDirection};
2use num_complex::Complex;
3
4use strength_reduce::{StrengthReducedU128, StrengthReducedU64};
5
6pub fn compute_twiddle<T: FftNum>(
7 index: usize,
8 fft_len: usize,
9 direction: FftDirection,
10) -> Complex<T> {
11 let constant = -2f64 * std::f64::consts::PI / fft_len as f64;
12 let angle = constant * index as f64;
13
14 let result = Complex {
15 re: T::from_f64(angle.cos()).unwrap(),
16 im: T::from_f64(angle.sin()).unwrap(),
17 };
18
19 match direction {
20 FftDirection::Forward => result,
21 FftDirection::Inverse => result.conj(),
22 }
23}
24
25pub fn fill_bluesteins_twiddles<T: FftNum>(
26 destination: &mut [Complex<T>],
27 direction: FftDirection,
28) {
29 let twice_len = destination.len() * 2;
30
31 if destination.len() < std::u32::MAX as usize {
38 let twice_len_reduced = StrengthReducedU64::new(twice_len as u64);
39
40 for (i, e) in destination.iter_mut().enumerate() {
41 let i_squared = i as u64 * i as u64;
42 let i_mod = i_squared % twice_len_reduced;
43 *e = compute_twiddle(i_mod as usize, twice_len, direction);
44 }
45 } else {
46 let twice_len_reduced = StrengthReducedU128::new(twice_len as u128);
48
49 for (i, e) in destination.iter_mut().enumerate() {
50 let i_squared = i as u128 * i as u128;
53 let i_mod = i_squared % twice_len_reduced;
54 *e = compute_twiddle(i_mod as usize, twice_len, direction);
55 }
56 }
57}
58
59pub fn rotate_90<T: FftNum>(value: Complex<T>, direction: FftDirection) -> Complex<T> {
60 match direction {
61 FftDirection::Forward => Complex {
62 re: value.im,
63 im: -value.re,
64 },
65 FftDirection::Inverse => Complex {
66 re: -value.im,
67 im: value.re,
68 },
69 }
70}
71
72#[cfg(test)]
73mod unit_tests {
74 use super::*;
75
76 #[test]
77 fn test_rotate() {
78 let value = Complex { re: 9.1, im: 2.2 };
80 let rotated_forward = rotate_90(value, FftDirection::Forward);
81 let twiddled_forward = value * compute_twiddle(1, 4, FftDirection::Forward);
82
83 assert_eq!(value.re, -rotated_forward.im);
84 assert_eq!(value.im, rotated_forward.re);
85
86 assert!(value.re + twiddled_forward.im < 0.0001);
87 assert!(value.im - rotated_forward.re < 0.0001);
88
89 let rotated_forward = rotate_90(value, FftDirection::Inverse);
91 let twiddled_forward = value * compute_twiddle(1, 4, FftDirection::Inverse);
92
93 assert_eq!(value.re, rotated_forward.im);
94 assert_eq!(value.im, -rotated_forward.re);
95
96 assert!(value.re - twiddled_forward.im < 0.0001);
97 assert!(value.im + rotated_forward.re < 0.0001);
98 }
99}