1use poulpy_hal::{
  2    api::{TakeSlice, VecZnxIDFTTmpBytes},
  3    layouts::{
  4        Backend, Data, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigToMut, VecZnxDft, VecZnxDftOwned, VecZnxDftToMut,
  5        VecZnxDftToRef, VecZnxToRef, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero,
  6    },
  7    oep::{
  8        DFTImpl, IDFTConsumeImpl, IDFTImpl, IDFTTmpAImpl, VecZnxDftAddImpl, VecZnxDftAddInplaceImpl, VecZnxDftAllocBytesImpl,
  9        VecZnxDftAllocImpl, VecZnxDftCopyImpl, VecZnxDftFromBytesImpl, VecZnxDftSubABInplaceImpl, VecZnxDftSubBAInplaceImpl,
 10        VecZnxDftSubImpl, VecZnxDftZeroImpl, VecZnxIDFTTmpBytesImpl,
 11    },
 12};
 13
 14use crate::cpu_spqlios::{
 15    FFT64,
 16    ffi::{vec_znx_big, vec_znx_dft},
 17};
 18
 19unsafe impl VecZnxDftFromBytesImpl<Self> for FFT64 {
 20    fn vec_znx_dft_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxDftOwned<Self> {
 21        VecZnxDft::<Vec<u8>, FFT64>::from_bytes(n, cols, size, bytes)
 22    }
 23}
 24
 25unsafe impl VecZnxDftAllocBytesImpl<Self> for FFT64 {
 26    fn vec_znx_dft_alloc_bytes_impl(n: usize, cols: usize, size: usize) -> usize {
 27        FFT64::layout_prep_word_count() * n * cols * size * size_of::<<FFT64 as Backend>::ScalarPrep>()
 28    }
 29}
 30
 31unsafe impl VecZnxDftAllocImpl<Self> for FFT64 {
 32    fn vec_znx_dft_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxDftOwned<Self> {
 33        VecZnxDftOwned::alloc(n, cols, size)
 34    }
 35}
 36
 37unsafe impl VecZnxIDFTTmpBytesImpl<Self> for FFT64 {
 38    fn vec_znx_idft_tmp_bytes_impl(module: &Module<Self>) -> usize {
 39        unsafe { vec_znx_dft::vec_znx_idft_tmp_bytes(module.ptr()) as usize }
 40    }
 41}
 42
 43unsafe impl IDFTImpl<Self> for FFT64 {
 44    fn idft_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, scratch: &mut Scratch<Self>)
 45    where
 46        R: VecZnxBigToMut<Self>,
 47        A: VecZnxDftToRef<Self>,
 48    {
 49        let mut res: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
 50        let a: VecZnxDft<&[u8], FFT64> = a.to_ref();
 51
 52        #[cfg(debug_assertions)]
 53        {
 54            assert_eq!(res.n(), a.n())
 55        }
 56
 57        let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_idft_tmp_bytes());
 58
 59        let min_size: usize = res.size().min(a.size());
 60
 61        unsafe {
 62            (0..min_size).for_each(|j| {
 63                vec_znx_dft::vec_znx_idft(
 64                    module.ptr(),
 65                    res.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
 66                    1_u64,
 67                    a.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
 68                    1_u64,
 69                    tmp_bytes.as_mut_ptr(),
 70                )
 71            });
 72            (min_size..res.size()).for_each(|j| {
 73                res.zero_at(res_col, j);
 74            });
 75        }
 76    }
 77}
 78
 79unsafe impl IDFTTmpAImpl<Self> for FFT64 {
 80    fn idft_tmp_a_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &mut A, a_col: usize)
 81    where
 82        R: VecZnxBigToMut<Self>,
 83        A: VecZnxDftToMut<Self>,
 84    {
 85        let mut res_mut: VecZnxBig<&mut [u8], FFT64> = res.to_mut();
 86        let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
 87
 88        let min_size: usize = res_mut.size().min(a_mut.size());
 89
 90        unsafe {
 91            (0..min_size).for_each(|j| {
 92                vec_znx_dft::vec_znx_idft_tmp_a(
 93                    module.ptr(),
 94                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_big::vec_znx_big_t,
 95                    1_u64,
 96                    a_mut.at_mut_ptr(a_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
 97                    1_u64,
 98                )
 99            });
100            (min_size..res_mut.size()).for_each(|j| {
101                res_mut.zero_at(res_col, j);
102            })
103        }
104    }
105}
106
107unsafe impl IDFTConsumeImpl<Self> for FFT64 {
108    fn idft_consume_impl<D: Data>(module: &Module<Self>, mut a: VecZnxDft<D, FFT64>) -> VecZnxBig<D, FFT64>
109    where
110        VecZnxDft<D, FFT64>: VecZnxDftToMut<Self>,
111    {
112        let mut a_mut: VecZnxDft<&mut [u8], FFT64> = a.to_mut();
113
114        unsafe {
115            // Rev col and rows because ZnxDft.sl() >= ZnxBig.sl()
116            (0..a_mut.size()).for_each(|j| {
117                (0..a_mut.cols()).for_each(|i| {
118                    vec_znx_dft::vec_znx_idft_tmp_a(
119                        module.ptr(),
120                        a_mut.at_mut_ptr(i, j) as *mut vec_znx_big::vec_znx_big_t,
121                        1_u64,
122                        a_mut.at_mut_ptr(i, j) as *mut vec_znx_dft::vec_znx_dft_t,
123                        1_u64,
124                    )
125                });
126            });
127        }
128
129        a.into_big()
130    }
131}
132
133unsafe impl DFTImpl<Self> for FFT64 {
134    fn dft_impl<R, A>(module: &Module<Self>, step: usize, offset: usize, res: &mut R, res_col: usize, a: &A, a_col: usize)
135    where
136        R: VecZnxDftToMut<Self>,
137        A: VecZnxToRef,
138    {
139        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
140        let a_ref: VecZnx<&[u8]> = a.to_ref();
141        let steps: usize = a_ref.size().div_ceil(step);
142        let min_steps: usize = res_mut.size().min(steps);
143        unsafe {
144            (0..min_steps).for_each(|j| {
145                let limb: usize = offset + j * step;
146                if limb < a_ref.size() {
147                    vec_znx_dft::vec_znx_dft(
148                        module.ptr(),
149                        res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
150                        1_u64,
151                        a_ref.at_ptr(a_col, limb),
152                        1_u64,
153                        a_ref.sl() as u64,
154                    )
155                }
156            });
157            (min_steps..res_mut.size()).for_each(|j| {
158                res_mut.zero_at(res_col, j);
159            });
160        }
161    }
162}
163
164unsafe impl VecZnxDftAddImpl<Self> for FFT64 {
165    fn vec_znx_dft_add_impl<R, A, D>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
166    where
167        R: VecZnxDftToMut<Self>,
168        A: VecZnxDftToRef<Self>,
169        D: VecZnxDftToRef<Self>,
170    {
171        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
172        let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
173        let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
174
175        let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
176
177        unsafe {
178            (0..min_size).for_each(|j| {
179                vec_znx_dft::vec_dft_add(
180                    module.ptr(),
181                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
182                    1,
183                    a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
184                    1,
185                    b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
186                    1,
187                );
188            });
189        }
190        (min_size..res_mut.size()).for_each(|j| {
191            res_mut.zero_at(res_col, j);
192        })
193    }
194}
195
196unsafe impl VecZnxDftAddInplaceImpl<Self> for FFT64 {
197    fn vec_znx_dft_add_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
198    where
199        R: VecZnxDftToMut<Self>,
200        A: VecZnxDftToRef<Self>,
201    {
202        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
203        let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
204
205        let min_size: usize = res_mut.size().min(a_ref.size());
206
207        unsafe {
208            (0..min_size).for_each(|j| {
209                vec_znx_dft::vec_dft_add(
210                    module.ptr(),
211                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
212                    1,
213                    res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
214                    1,
215                    a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
216                    1,
217                );
218            });
219        }
220    }
221}
222
223unsafe impl VecZnxDftSubImpl<Self> for FFT64 {
224    fn vec_znx_dft_sub_impl<R, A, D>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &D, b_col: usize)
225    where
226        R: VecZnxDftToMut<Self>,
227        A: VecZnxDftToRef<Self>,
228        D: VecZnxDftToRef<Self>,
229    {
230        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
231        let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
232        let b_ref: VecZnxDft<&[u8], FFT64> = b.to_ref();
233
234        let min_size: usize = res_mut.size().min(a_ref.size()).min(b_ref.size());
235
236        unsafe {
237            (0..min_size).for_each(|j| {
238                vec_znx_dft::vec_dft_sub(
239                    module.ptr(),
240                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
241                    1,
242                    a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
243                    1,
244                    b_ref.at_ptr(b_col, j) as *const vec_znx_dft::vec_znx_dft_t,
245                    1,
246                );
247            });
248        }
249        (min_size..res_mut.size()).for_each(|j| {
250            res_mut.zero_at(res_col, j);
251        })
252    }
253}
254
255unsafe impl VecZnxDftSubABInplaceImpl<Self> for FFT64 {
256    fn vec_znx_dft_sub_ab_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
257    where
258        R: VecZnxDftToMut<Self>,
259        A: VecZnxDftToRef<Self>,
260    {
261        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
262        let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
263
264        let min_size: usize = res_mut.size().min(a_ref.size());
265
266        unsafe {
267            (0..min_size).for_each(|j| {
268                vec_znx_dft::vec_dft_sub(
269                    module.ptr(),
270                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
271                    1,
272                    res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
273                    1,
274                    a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
275                    1,
276                );
277            });
278        }
279    }
280}
281
282unsafe impl VecZnxDftSubBAInplaceImpl<Self> for FFT64 {
283    fn vec_znx_dft_sub_ba_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
284    where
285        R: VecZnxDftToMut<Self>,
286        A: VecZnxDftToRef<Self>,
287    {
288        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
289        let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
290
291        let min_size: usize = res_mut.size().min(a_ref.size());
292
293        unsafe {
294            (0..min_size).for_each(|j| {
295                vec_znx_dft::vec_dft_sub(
296                    module.ptr(),
297                    res_mut.at_mut_ptr(res_col, j) as *mut vec_znx_dft::vec_znx_dft_t,
298                    1,
299                    a_ref.at_ptr(a_col, j) as *const vec_znx_dft::vec_znx_dft_t,
300                    1,
301                    res_mut.at_ptr(res_col, j) as *const vec_znx_dft::vec_znx_dft_t,
302                    1,
303                );
304            });
305        }
306    }
307}
308
309unsafe impl VecZnxDftCopyImpl<Self> for FFT64 {
310    fn vec_znx_dft_copy_impl<R, A>(
311        _module: &Module<Self>,
312        step: usize,
313        offset: usize,
314        res: &mut R,
315        res_col: usize,
316        a: &A,
317        a_col: usize,
318    ) where
319        R: VecZnxDftToMut<Self>,
320        A: VecZnxDftToRef<Self>,
321    {
322        let mut res_mut: VecZnxDft<&mut [u8], FFT64> = res.to_mut();
323        let a_ref: VecZnxDft<&[u8], FFT64> = a.to_ref();
324
325        let steps: usize = a_ref.size().div_ceil(step);
326        let min_steps: usize = res_mut.size().min(steps);
327
328        (0..min_steps).for_each(|j| {
329            let limb: usize = offset + j * step;
330            if limb < a_ref.size() {
331                res_mut
332                    .at_mut(res_col, j)
333                    .copy_from_slice(a_ref.at(a_col, limb));
334            }
335        });
336        (min_steps..res_mut.size()).for_each(|j| {
337            res_mut.zero_at(res_col, j);
338        })
339    }
340}
341
342unsafe impl VecZnxDftZeroImpl<Self> for FFT64 {
343    fn vec_znx_dft_zero_impl<R>(_module: &Module<Self>, res: &mut R)
344    where
345        R: VecZnxDftToMut<Self>,
346    {
347        res.to_mut().data.fill(0);
348    }
349}