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