1use std::fmt::Debug;
  2
  3use rand_distr::num_traits::{Float, FloatConst};
  4
  5use crate::reference::fft64::reim::{as_arr, as_arr_mut};
  6
  7#[inline(always)]
  8pub fn fft_ref<R: Float + FloatConst + Debug>(m: usize, omg: &[R], data: &mut [R]) {
  9    assert!(data.len() == 2 * m);
 10    let (re, im) = data.split_at_mut(m);
 11
 12    if m <= 16 {
 13        match m {
 14            1 => {}
 15            2 => fft2_ref(as_arr_mut::<2, R>(re), as_arr_mut::<2, R>(im), *as_arr::<2, R>(omg)),
 16            4 => fft4_ref(as_arr_mut::<4, R>(re), as_arr_mut::<4, R>(im), *as_arr::<4, R>(omg)),
 17            8 => fft8_ref(as_arr_mut::<8, R>(re), as_arr_mut::<8, R>(im), *as_arr::<8, R>(omg)),
 18            16 => fft16_ref(as_arr_mut::<16, R>(re), as_arr_mut::<16, R>(im), *as_arr::<16, R>(omg)),
 19            _ => {}
 20        }
 21    } else if m <= 2048 {
 22        fft_bfs_16_ref(m, re, im, omg, 0);
 23    } else {
 24        fft_rec_16_ref(m, re, im, omg, 0);
 25    }
 26}
 27
 28#[inline(always)]
 29fn fft_rec_16_ref<R: Float + FloatConst + Debug>(m: usize, re: &mut [R], im: &mut [R], omg: &[R], mut pos: usize) -> usize {
 30    if m <= 2048 {
 31        return fft_bfs_16_ref(m, re, im, omg, pos);
 32    };
 33
 34    let h = m >> 1;
 35    twiddle_fft_ref(h, re, im, as_arr::<2, R>(&omg[pos..]));
 36    pos += 2;
 37    pos = fft_rec_16_ref(h, re, im, omg, pos);
 38    pos = fft_rec_16_ref(h, &mut re[h..], &mut im[h..], omg, pos);
 39    pos
 40}
 41
 42#[inline(always)]
 43fn cplx_twiddle<R: Float + FloatConst>(ra: &mut R, ia: &mut R, rb: &mut R, ib: &mut R, omg_re: R, omg_im: R) {
 44    let dr: R = *rb * omg_re - *ib * omg_im;
 45    let di: R = *rb * omg_im + *ib * omg_re;
 46    *rb = *ra - dr;
 47    *ib = *ia - di;
 48    *ra = *ra + dr;
 49    *ia = *ia + di;
 50}
 51
 52#[inline(always)]
 53fn cplx_i_twiddle<R: Float + FloatConst>(ra: &mut R, ia: &mut R, rb: &mut R, ib: &mut R, omg_re: R, omg_im: R) {
 54    let dr: R = *rb * omg_im + *ib * omg_re;
 55    let di: R = *rb * omg_re - *ib * omg_im;
 56    *rb = *ra + dr;
 57    *ib = *ia - di;
 58    *ra = *ra - dr;
 59    *ia = *ia + di;
 60}
 61
 62#[inline(always)]
 63fn fft2_ref<R: Float + FloatConst>(re: &mut [R; 2], im: &mut [R; 2], omg: [R; 2]) {
 64    let [ra, rb] = re;
 65    let [ia, ib] = im;
 66    let [romg, iomg] = omg;
 67    cplx_twiddle(ra, ia, rb, ib, romg, iomg);
 68}
 69
 70#[inline(always)]
 71fn fft4_ref<R: Float + FloatConst>(re: &mut [R; 4], im: &mut [R; 4], omg: [R; 4]) {
 72    let [re_0, re_1, re_2, re_3] = re;
 73    let [im_0, im_1, im_2, im_3] = im;
 74
 75    {
 76        let omg_0 = omg[0];
 77        let omg_1 = omg[1];
 78        cplx_twiddle(re_0, im_0, re_2, im_2, omg_0, omg_1);
 79        cplx_twiddle(re_1, im_1, re_3, im_3, omg_0, omg_1);
 80    }
 81
 82    {
 83        let omg_0 = omg[2];
 84        let omg_1 = omg[3];
 85        cplx_twiddle(re_0, im_0, re_1, im_1, omg_0, omg_1);
 86        cplx_i_twiddle(re_2, im_2, re_3, im_3, omg_0, omg_1);
 87    }
 88}
 89
 90#[inline(always)]
 91fn fft8_ref<R: Float + FloatConst>(re: &mut [R; 8], im: &mut [R; 8], omg: [R; 8]) {
 92    let [re_0, re_1, re_2, re_3, re_4, re_5, re_6, re_7] = re;
 93    let [im_0, im_1, im_2, im_3, im_4, im_5, im_6, im_7] = im;
 94
 95    {
 96        let omg_0 = omg[0];
 97        let omg_1 = omg[1];
 98        cplx_twiddle(re_0, im_0, re_4, im_4, omg_0, omg_1);
 99        cplx_twiddle(re_1, im_1, re_5, im_5, omg_0, omg_1);
100        cplx_twiddle(re_2, im_2, re_6, im_6, omg_0, omg_1);
101        cplx_twiddle(re_3, im_3, re_7, im_7, omg_0, omg_1);
102    }
103
104    {
105        let omg_2 = omg[2];
106        let omg_3 = omg[3];
107        cplx_twiddle(re_0, im_0, re_2, im_2, omg_2, omg_3);
108        cplx_twiddle(re_1, im_1, re_3, im_3, omg_2, omg_3);
109        cplx_i_twiddle(re_4, im_4, re_6, im_6, omg_2, omg_3);
110        cplx_i_twiddle(re_5, im_5, re_7, im_7, omg_2, omg_3);
111    }
112
113    {
114        let omg_4 = omg[4];
115        let omg_5 = omg[5];
116        let omg_6 = omg[6];
117        let omg_7 = omg[7];
118        cplx_twiddle(re_0, im_0, re_1, im_1, omg_4, omg_6);
119        cplx_i_twiddle(re_2, im_2, re_3, im_3, omg_4, omg_6);
120        cplx_twiddle(re_4, im_4, re_5, im_5, omg_5, omg_7);
121        cplx_i_twiddle(re_6, im_6, re_7, im_7, omg_5, omg_7);
122    }
123}
124
125#[inline(always)]
126fn fft16_ref<R: Float + FloatConst + Debug>(re: &mut [R; 16], im: &mut [R; 16], omg: [R; 16]) {
127    let [
128        re_0,
129        re_1,
130        re_2,
131        re_3,
132        re_4,
133        re_5,
134        re_6,
135        re_7,
136        re_8,
137        re_9,
138        re_10,
139        re_11,
140        re_12,
141        re_13,
142        re_14,
143        re_15,
144    ] = re;
145    let [
146        im_0,
147        im_1,
148        im_2,
149        im_3,
150        im_4,
151        im_5,
152        im_6,
153        im_7,
154        im_8,
155        im_9,
156        im_10,
157        im_11,
158        im_12,
159        im_13,
160        im_14,
161        im_15,
162    ] = im;
163
164    {
165        let omg_0: R = omg[0];
166        let omg_1: R = omg[1];
167        cplx_twiddle(re_0, im_0, re_8, im_8, omg_0, omg_1);
168        cplx_twiddle(re_1, im_1, re_9, im_9, omg_0, omg_1);
169        cplx_twiddle(re_2, im_2, re_10, im_10, omg_0, omg_1);
170        cplx_twiddle(re_3, im_3, re_11, im_11, omg_0, omg_1);
171
172        cplx_twiddle(re_4, im_4, re_12, im_12, omg_0, omg_1);
173        cplx_twiddle(re_5, im_5, re_13, im_13, omg_0, omg_1);
174        cplx_twiddle(re_6, im_6, re_14, im_14, omg_0, omg_1);
175        cplx_twiddle(re_7, im_7, re_15, im_15, omg_0, omg_1);
176    }
177
178    {
179        let omg_2: R = omg[2];
180        let omg_3: R = omg[3];
181        cplx_twiddle(re_0, im_0, re_4, im_4, omg_2, omg_3);
182        cplx_twiddle(re_1, im_1, re_5, im_5, omg_2, omg_3);
183        cplx_twiddle(re_2, im_2, re_6, im_6, omg_2, omg_3);
184        cplx_twiddle(re_3, im_3, re_7, im_7, omg_2, omg_3);
185
186        cplx_i_twiddle(re_8, im_8, re_12, im_12, omg_2, omg_3);
187        cplx_i_twiddle(re_9, im_9, re_13, im_13, omg_2, omg_3);
188        cplx_i_twiddle(re_10, im_10, re_14, im_14, omg_2, omg_3);
189        cplx_i_twiddle(re_11, im_11, re_15, im_15, omg_2, omg_3);
190    }
191
192    {
193        let omg_0: R = omg[4];
194        let omg_1: R = omg[5];
195        let omg_2: R = omg[6];
196        let omg_3: R = omg[7];
197        cplx_twiddle(re_0, im_0, re_2, im_2, omg_0, omg_1);
198        cplx_twiddle(re_1, im_1, re_3, im_3, omg_0, omg_1);
199        cplx_twiddle(re_8, im_8, re_10, im_10, omg_2, omg_3);
200        cplx_twiddle(re_9, im_9, re_11, im_11, omg_2, omg_3);
201
202        cplx_i_twiddle(re_4, im_4, re_6, im_6, omg_0, omg_1);
203        cplx_i_twiddle(re_5, im_5, re_7, im_7, omg_0, omg_1);
204        cplx_i_twiddle(re_12, im_12, re_14, im_14, omg_2, omg_3);
205        cplx_i_twiddle(re_13, im_13, re_15, im_15, omg_2, omg_3);
206    }
207
208    {
209        let omg_0: R = omg[8];
210        let omg_1: R = omg[9];
211        let omg_2: R = omg[10];
212        let omg_3: R = omg[11];
213        let omg_4: R = omg[12];
214        let omg_5: R = omg[13];
215        let omg_6: R = omg[14];
216        let omg_7: R = omg[15];
217        cplx_twiddle(re_0, im_0, re_1, im_1, omg_0, omg_4);
218        cplx_twiddle(re_4, im_4, re_5, im_5, omg_1, omg_5);
219        cplx_twiddle(re_8, im_8, re_9, im_9, omg_2, omg_6);
220        cplx_twiddle(re_12, im_12, re_13, im_13, omg_3, omg_7);
221
222        cplx_i_twiddle(re_2, im_2, re_3, im_3, omg_0, omg_4);
223        cplx_i_twiddle(re_6, im_6, re_7, im_7, omg_1, omg_5);
224        cplx_i_twiddle(re_10, im_10, re_11, im_11, omg_2, omg_6);
225        cplx_i_twiddle(re_14, im_14, re_15, im_15, omg_3, omg_7);
226    }
227}
228
229#[inline(always)]
230fn fft_bfs_16_ref<R: Float + FloatConst + Debug>(m: usize, re: &mut [R], im: &mut [R], omg: &[R], mut pos: usize) -> usize {
231    let log_m: usize = (usize::BITS - (m - 1).leading_zeros()) as usize;
232    let mut mm: usize = m;
233
234    if !log_m.is_multiple_of(2) {
235        let h: usize = mm >> 1;
236        twiddle_fft_ref(h, re, im, as_arr::<2, R>(&omg[pos..]));
237        pos += 2;
238        mm = h
239    }
240
241    while mm > 16 {
242        let h: usize = mm >> 2;
243        for off in (0..m).step_by(mm) {
244            bitwiddle_fft_ref(h, &mut re[off..], &mut im[off..], as_arr::<4, R>(&omg[pos..]));
245            pos += 4;
246        }
247        mm = h
248    }
249
250    for off in (0..m).step_by(16) {
251        fft16_ref(
252            as_arr_mut::<16, R>(&mut re[off..]),
253            as_arr_mut::<16, R>(&mut im[off..]),
254            *as_arr::<16, R>(&omg[pos..]),
255        );
256        pos += 16;
257    }
258
259    pos
260}
261
262#[inline(always)]
263fn twiddle_fft_ref<R: Float + FloatConst>(h: usize, re: &mut [R], im: &mut [R], omg: &[R; 2]) {
264    let romg = omg[0];
265    let iomg = omg[1];
266
267    let (re_lhs, re_rhs) = re.split_at_mut(h);
268    let (im_lhs, im_rhs) = im.split_at_mut(h);
269
270    for i in 0..h {
271        cplx_twiddle(&mut re_lhs[i], &mut im_lhs[i], &mut re_rhs[i], &mut im_rhs[i], romg, iomg);
272    }
273}
274
275#[inline(always)]
276fn bitwiddle_fft_ref<R: Float + FloatConst>(h: usize, re: &mut [R], im: &mut [R], omg: &[R; 4]) {
277    let (r0, r2) = re.split_at_mut(2 * h);
278    let (r0, r1) = r0.split_at_mut(h);
279    let (r2, r3) = r2.split_at_mut(h);
280
281    let (i0, i2) = im.split_at_mut(2 * h);
282    let (i0, i1) = i0.split_at_mut(h);
283    let (i2, i3) = i2.split_at_mut(h);
284
285    let omg_0: R = omg[0];
286    let omg_1: R = omg[1];
287    let omg_2: R = omg[2];
288    let omg_3: R = omg[3];
289
290    for i in 0..h {
291        cplx_twiddle(&mut r0[i], &mut i0[i], &mut r2[i], &mut i2[i], omg_0, omg_1);
292        cplx_twiddle(&mut r1[i], &mut i1[i], &mut r3[i], &mut i3[i], omg_0, omg_1);
293    }
294
295    for i in 0..h {
296        cplx_twiddle(&mut r0[i], &mut i0[i], &mut r1[i], &mut i1[i], omg_2, omg_3);
297        cplx_i_twiddle(&mut r2[i], &mut i2[i], &mut r3[i], &mut i3[i], omg_2, omg_3);
298    }
299}