1use poulpy_hal::{
  2    api::{
  3        TakeSlice, VecZnxAutomorphismInplaceTmpBytes, VecZnxLshTmpBytes, VecZnxMergeRingsTmpBytes,
  4        VecZnxMulXpMinusOneInplaceTmpBytes, VecZnxNormalizeTmpBytes, VecZnxRotateInplaceTmpBytes, VecZnxRshTmpBytes,
  5        VecZnxSplitRingTmpBytes,
  6    },
  7    layouts::{Module, ScalarZnxToRef, Scratch, VecZnxToMut, VecZnxToRef},
  8    oep::{
  9        TakeSliceImpl, VecZnxAddImpl, VecZnxAddInplaceImpl, VecZnxAddNormalImpl, VecZnxAddScalarImpl, VecZnxAddScalarInplaceImpl,
 10        VecZnxAutomorphismImpl, VecZnxAutomorphismInplaceImpl, VecZnxAutomorphismInplaceTmpBytesImpl, VecZnxCopyImpl,
 11        VecZnxFillNormalImpl, VecZnxFillUniformImpl, VecZnxLshImpl, VecZnxLshInplaceImpl, VecZnxLshTmpBytesImpl,
 12        VecZnxMergeRingsImpl, VecZnxMergeRingsTmpBytesImpl, VecZnxMulXpMinusOneImpl, VecZnxMulXpMinusOneInplaceImpl,
 13        VecZnxMulXpMinusOneInplaceTmpBytesImpl, VecZnxNegateImpl, VecZnxNegateInplaceImpl, VecZnxNormalizeImpl,
 14        VecZnxNormalizeInplaceImpl, VecZnxNormalizeTmpBytesImpl, VecZnxRotateImpl, VecZnxRotateInplaceImpl,
 15        VecZnxRotateInplaceTmpBytesImpl, VecZnxRshImpl, VecZnxRshInplaceImpl, VecZnxRshTmpBytesImpl, VecZnxSplitRingImpl,
 16        VecZnxSplitRingTmpBytesImpl, VecZnxSubImpl, VecZnxSubInplaceImpl, VecZnxSubNegateInplaceImpl, VecZnxSubScalarImpl,
 17        VecZnxSubScalarInplaceImpl, VecZnxSwitchRingImpl,
 18    },
 19    reference::vec_znx::{
 20        vec_znx_add, vec_znx_add_inplace, vec_znx_add_normal_ref, vec_znx_add_scalar, vec_znx_add_scalar_inplace,
 21        vec_znx_automorphism, vec_znx_automorphism_inplace, vec_znx_automorphism_inplace_tmp_bytes, vec_znx_copy,
 22        vec_znx_fill_normal_ref, vec_znx_fill_uniform_ref, vec_znx_lsh, vec_znx_lsh_inplace, vec_znx_lsh_tmp_bytes,
 23        vec_znx_merge_rings, vec_znx_merge_rings_tmp_bytes, vec_znx_mul_xp_minus_one, vec_znx_mul_xp_minus_one_inplace,
 24        vec_znx_mul_xp_minus_one_inplace_tmp_bytes, vec_znx_negate, vec_znx_negate_inplace, vec_znx_normalize,
 25        vec_znx_normalize_inplace, vec_znx_normalize_tmp_bytes, vec_znx_rotate, vec_znx_rotate_inplace,
 26        vec_znx_rotate_inplace_tmp_bytes, vec_znx_rsh, vec_znx_rsh_inplace, vec_znx_rsh_tmp_bytes, vec_znx_split_ring,
 27        vec_znx_split_ring_tmp_bytes, vec_znx_sub, vec_znx_sub_inplace, vec_znx_sub_negate_inplace, vec_znx_sub_scalar,
 28        vec_znx_sub_scalar_inplace, vec_znx_switch_ring,
 29    },
 30    source::Source,
 31};
 32
 33use crate::cpu_fft64_ref::FFT64Ref;
 34
 35unsafe impl VecZnxNormalizeTmpBytesImpl<Self> for FFT64Ref {
 36    fn vec_znx_normalize_tmp_bytes_impl(module: &Module<Self>) -> usize {
 37        vec_znx_normalize_tmp_bytes(module.n())
 38    }
 39}
 40
 41unsafe impl VecZnxNormalizeImpl<Self> for FFT64Ref
 42where
 43    Self: TakeSliceImpl<Self> + VecZnxNormalizeTmpBytesImpl<Self>,
 44{
 45    fn vec_znx_normalize_impl<R, A>(
 46        module: &Module<Self>,
 47        res_basek: usize,
 48        res: &mut R,
 49        res_col: usize,
 50        a_basek: usize,
 51        a: &A,
 52        a_col: usize,
 53        scratch: &mut Scratch<Self>,
 54    ) where
 55        R: VecZnxToMut,
 56        A: VecZnxToRef,
 57    {
 58        let (carry, _) = scratch.take_slice(module.vec_znx_normalize_tmp_bytes() / size_of::<i64>());
 59        vec_znx_normalize::<R, A, Self>(res_basek, res, res_col, a_basek, a, a_col, carry);
 60    }
 61}
 62
 63unsafe impl VecZnxNormalizeInplaceImpl<Self> for FFT64Ref
 64where
 65    Self: TakeSliceImpl<Self> + VecZnxNormalizeTmpBytesImpl<Self>,
 66{
 67    fn vec_znx_normalize_inplace_impl<R>(
 68        module: &Module<Self>,
 69        base2k: usize,
 70        res: &mut R,
 71        res_col: usize,
 72        scratch: &mut Scratch<Self>,
 73    ) where
 74        R: VecZnxToMut,
 75    {
 76        let (carry, _) = scratch.take_slice(module.vec_znx_normalize_tmp_bytes() / size_of::<i64>());
 77        vec_znx_normalize_inplace::<R, Self>(base2k, res, res_col, carry);
 78    }
 79}
 80
 81unsafe impl VecZnxAddImpl<Self> for FFT64Ref {
 82    fn vec_znx_add_impl<R, A, B>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
 83    where
 84        R: VecZnxToMut,
 85        A: VecZnxToRef,
 86        B: VecZnxToRef,
 87    {
 88        vec_znx_add::<R, A, B, Self>(res, res_col, a, a_col, b, b_col);
 89    }
 90}
 91
 92unsafe impl VecZnxAddInplaceImpl<Self> for FFT64Ref {
 93    fn vec_znx_add_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
 94    where
 95        R: VecZnxToMut,
 96        A: VecZnxToRef,
 97    {
 98        vec_znx_add_inplace::<R, A, Self>(res, res_col, a, a_col);
 99    }
100}
101
102unsafe impl VecZnxAddScalarInplaceImpl<Self> for FFT64Ref {
103    fn vec_znx_add_scalar_inplace_impl<R, A>(
104        _module: &Module<Self>,
105        res: &mut R,
106        res_col: usize,
107        res_limb: usize,
108        a: &A,
109        a_col: usize,
110    ) where
111        R: VecZnxToMut,
112        A: ScalarZnxToRef,
113    {
114        vec_znx_add_scalar_inplace::<R, A, Self>(res, res_col, res_limb, a, a_col);
115    }
116}
117
118unsafe impl VecZnxAddScalarImpl<Self> for FFT64Ref {
119    fn vec_znx_add_scalar_impl<R, A, B>(
120        _module: &Module<Self>,
121        res: &mut R,
122        res_col: usize,
123        a: &A,
124        a_col: usize,
125        b: &B,
126        b_col: usize,
127        b_limb: usize,
128    ) where
129        R: VecZnxToMut,
130        A: ScalarZnxToRef,
131        B: VecZnxToRef,
132    {
133        vec_znx_add_scalar::<R, A, B, Self>(res, res_col, a, a_col, b, b_col, b_limb);
134    }
135}
136
137unsafe impl VecZnxSubImpl<Self> for FFT64Ref {
138    fn vec_znx_sub_impl<R, A, B>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &B, b_col: usize)
139    where
140        R: VecZnxToMut,
141        A: VecZnxToRef,
142        B: VecZnxToRef,
143    {
144        vec_znx_sub::<R, A, B, Self>(res, res_col, a, a_col, b, b_col);
145    }
146}
147
148unsafe impl VecZnxSubInplaceImpl<Self> for FFT64Ref {
149    fn vec_znx_sub_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
150    where
151        R: VecZnxToMut,
152        A: VecZnxToRef,
153    {
154        vec_znx_sub_inplace::<R, A, Self>(res, res_col, a, a_col);
155    }
156}
157
158unsafe impl VecZnxSubNegateInplaceImpl<Self> for FFT64Ref {
159    fn vec_znx_sub_negate_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
160    where
161        R: VecZnxToMut,
162        A: VecZnxToRef,
163    {
164        vec_znx_sub_negate_inplace::<R, A, Self>(res, res_col, a, a_col);
165    }
166}
167
168unsafe impl VecZnxSubScalarImpl<Self> for FFT64Ref {
169    fn vec_znx_sub_scalar_impl<R, A, B>(
170        _module: &Module<Self>,
171        res: &mut R,
172        res_col: usize,
173        a: &A,
174        a_col: usize,
175        b: &B,
176        b_col: usize,
177        b_limb: usize,
178    ) where
179        R: VecZnxToMut,
180        A: ScalarZnxToRef,
181        B: VecZnxToRef,
182    {
183        vec_znx_sub_scalar::<R, A, B, Self>(res, res_col, a, a_col, b, b_col, b_limb);
184    }
185}
186
187unsafe impl VecZnxSubScalarInplaceImpl<Self> for FFT64Ref {
188    fn vec_znx_sub_scalar_inplace_impl<R, A>(
189        _module: &Module<Self>,
190        res: &mut R,
191        res_col: usize,
192        res_limb: usize,
193        a: &A,
194        a_col: usize,
195    ) where
196        R: VecZnxToMut,
197        A: ScalarZnxToRef,
198    {
199        vec_znx_sub_scalar_inplace::<R, A, Self>(res, res_col, res_limb, a, a_col);
200    }
201}
202
203unsafe impl VecZnxNegateImpl<Self> for FFT64Ref {
204    fn vec_znx_negate_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
205    where
206        R: VecZnxToMut,
207        A: VecZnxToRef,
208    {
209        vec_znx_negate::<R, A, Self>(res, res_col, a, a_col);
210    }
211}
212
213unsafe impl VecZnxNegateInplaceImpl<Self> for FFT64Ref {
214    fn vec_znx_negate_inplace_impl<R>(_module: &Module<Self>, res: &mut R, res_col: usize)
215    where
216        R: VecZnxToMut,
217    {
218        vec_znx_negate_inplace::<R, Self>(res, res_col);
219    }
220}
221
222unsafe impl VecZnxLshTmpBytesImpl<Self> for FFT64Ref {
223    fn vec_znx_lsh_tmp_bytes_impl(module: &Module<Self>) -> usize {
224        vec_znx_lsh_tmp_bytes(module.n())
225    }
226}
227
228unsafe impl VecZnxRshTmpBytesImpl<Self> for FFT64Ref {
229    fn vec_znx_rsh_tmp_bytes_impl(module: &Module<Self>) -> usize {
230        vec_znx_rsh_tmp_bytes(module.n())
231    }
232}
233
234unsafe impl VecZnxLshImpl<Self> for FFT64Ref
235where
236    Module<Self>: VecZnxNormalizeTmpBytes,
237    Scratch<Self>: TakeSlice,
238{
239    fn vec_znx_lsh_impl<R, A>(
240        module: &Module<Self>,
241        base2k: usize,
242        k: usize,
243        res: &mut R,
244        res_col: usize,
245        a: &A,
246        a_col: usize,
247        scratch: &mut Scratch<Self>,
248    ) where
249        R: VecZnxToMut,
250        A: VecZnxToRef,
251    {
252        let (carry, _) = scratch.take_slice(module.vec_znx_lsh_tmp_bytes() / size_of::<i64>());
253        vec_znx_lsh::<_, _, Self>(base2k, k, res, res_col, a, a_col, carry);
254    }
255}
256
257unsafe impl VecZnxLshInplaceImpl<Self> for FFT64Ref
258where
259    Module<Self>: VecZnxNormalizeTmpBytes,
260    Scratch<Self>: TakeSlice,
261{
262    fn vec_znx_lsh_inplace_impl<A>(
263        module: &Module<Self>,
264        base2k: usize,
265        k: usize,
266        a: &mut A,
267        a_col: usize,
268        scratch: &mut Scratch<Self>,
269    ) where
270        A: VecZnxToMut,
271    {
272        let (carry, _) = scratch.take_slice(module.vec_znx_lsh_tmp_bytes() / size_of::<i64>());
273        vec_znx_lsh_inplace::<_, Self>(base2k, k, a, a_col, carry);
274    }
275}
276
277unsafe impl VecZnxRshImpl<Self> for FFT64Ref
278where
279    Module<Self>: VecZnxNormalizeTmpBytes,
280    Scratch<Self>: TakeSlice,
281{
282    fn vec_znx_rsh_impl<R, A>(
283        module: &Module<Self>,
284        base2k: usize,
285        k: usize,
286        res: &mut R,
287        res_col: usize,
288        a: &A,
289        a_col: usize,
290        scratch: &mut Scratch<Self>,
291    ) where
292        R: VecZnxToMut,
293        A: VecZnxToRef,
294    {
295        let (carry, _) = scratch.take_slice(module.vec_znx_rsh_tmp_bytes() / size_of::<i64>());
296        vec_znx_rsh::<_, _, Self>(base2k, k, res, res_col, a, a_col, carry);
297    }
298}
299
300unsafe impl VecZnxRshInplaceImpl<Self> for FFT64Ref
301where
302    Module<Self>: VecZnxNormalizeTmpBytes,
303    Scratch<Self>: TakeSlice,
304{
305    fn vec_znx_rsh_inplace_impl<A>(
306        module: &Module<Self>,
307        base2k: usize,
308        k: usize,
309        a: &mut A,
310        a_col: usize,
311        scratch: &mut Scratch<Self>,
312    ) where
313        A: VecZnxToMut,
314    {
315        let (carry, _) = scratch.take_slice(module.vec_znx_rsh_tmp_bytes() / size_of::<i64>());
316        vec_znx_rsh_inplace::<_, Self>(base2k, k, a, a_col, carry);
317    }
318}
319
320unsafe impl VecZnxRotateImpl<Self> for FFT64Ref {
321    fn vec_znx_rotate_impl<R, A>(_module: &Module<Self>, p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
322    where
323        R: VecZnxToMut,
324        A: VecZnxToRef,
325    {
326        vec_znx_rotate::<R, A, Self>(p, res, res_col, a, a_col);
327    }
328}
329
330unsafe impl VecZnxRotateInplaceTmpBytesImpl<Self> for FFT64Ref
331where
332    Scratch<Self>: TakeSlice,
333{
334    fn vec_znx_rotate_inplace_tmp_bytes_impl(module: &Module<Self>) -> usize {
335        vec_znx_rotate_inplace_tmp_bytes(module.n())
336    }
337}
338
339unsafe impl VecZnxRotateInplaceImpl<Self> for FFT64Ref
340where
341    Scratch<Self>: TakeSlice,
342    Self: VecZnxRotateInplaceTmpBytesImpl<Self>,
343{
344    fn vec_znx_rotate_inplace_impl<R>(module: &Module<Self>, p: i64, res: &mut R, res_col: usize, scratch: &mut Scratch<Self>)
345    where
346        R: VecZnxToMut,
347    {
348        let (tmp, _) = scratch.take_slice(module.vec_znx_rotate_inplace_tmp_bytes() / size_of::<i64>());
349        vec_znx_rotate_inplace::<R, Self>(p, res, res_col, tmp);
350    }
351}
352
353unsafe impl VecZnxAutomorphismImpl<Self> for FFT64Ref {
354    fn vec_znx_automorphism_impl<R, A>(_module: &Module<Self>, p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
355    where
356        R: VecZnxToMut,
357        A: VecZnxToRef,
358    {
359        vec_znx_automorphism::<R, A, Self>(p, res, res_col, a, a_col);
360    }
361}
362
363unsafe impl VecZnxAutomorphismInplaceTmpBytesImpl<Self> for FFT64Ref {
364    fn vec_znx_automorphism_inplace_tmp_bytes_impl(module: &Module<Self>) -> usize {
365        vec_znx_automorphism_inplace_tmp_bytes(module.n())
366    }
367}
368
369unsafe impl VecZnxAutomorphismInplaceImpl<Self> for FFT64Ref
370where
371    Scratch<Self>: TakeSlice,
372    Self: VecZnxAutomorphismInplaceTmpBytesImpl<Self>,
373{
374    fn vec_znx_automorphism_inplace_impl<R>(
375        module: &Module<Self>,
376        p: i64,
377        res: &mut R,
378        res_col: usize,
379        scratch: &mut Scratch<Self>,
380    ) where
381        R: VecZnxToMut,
382    {
383        let (tmp, _) = scratch.take_slice(module.vec_znx_automorphism_inplace_tmp_bytes() / size_of::<i64>());
384        vec_znx_automorphism_inplace::<R, Self>(p, res, res_col, tmp);
385    }
386}
387
388unsafe impl VecZnxMulXpMinusOneImpl<Self> for FFT64Ref {
389    fn vec_znx_mul_xp_minus_one_impl<R, A>(_module: &Module<Self>, p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
390    where
391        R: VecZnxToMut,
392        A: VecZnxToRef,
393    {
394        vec_znx_mul_xp_minus_one::<R, A, Self>(p, res, res_col, a, a_col);
395    }
396}
397
398unsafe impl VecZnxMulXpMinusOneInplaceTmpBytesImpl<Self> for FFT64Ref
399where
400    Scratch<Self>: TakeSlice,
401    Self: VecZnxMulXpMinusOneImpl<Self>,
402{
403    fn vec_znx_mul_xp_minus_one_inplace_tmp_bytes_impl(module: &Module<Self>) -> usize {
404        vec_znx_mul_xp_minus_one_inplace_tmp_bytes(module.n())
405    }
406}
407
408unsafe impl VecZnxMulXpMinusOneInplaceImpl<Self> for FFT64Ref {
409    fn vec_znx_mul_xp_minus_one_inplace_impl<R>(
410        module: &Module<Self>,
411        p: i64,
412        res: &mut R,
413        res_col: usize,
414        scratch: &mut Scratch<Self>,
415    ) where
416        R: VecZnxToMut,
417    {
418        let (tmp, _) = scratch.take_slice(module.vec_znx_mul_xp_minus_one_inplace_tmp_bytes() / size_of::<i64>());
419        vec_znx_mul_xp_minus_one_inplace::<R, Self>(p, res, res_col, tmp);
420    }
421}
422
423unsafe impl VecZnxSplitRingTmpBytesImpl<Self> for FFT64Ref {
424    fn vec_znx_split_ring_tmp_bytes_impl(module: &Module<Self>) -> usize {
425        vec_znx_split_ring_tmp_bytes(module.n())
426    }
427}
428
429unsafe impl VecZnxSplitRingImpl<Self> for FFT64Ref
430where
431    Module<Self>: VecZnxSplitRingTmpBytes,
432    Scratch<Self>: TakeSlice,
433{
434    fn vec_znx_split_ring_impl<R, A>(
435        module: &Module<Self>,
436        res: &mut [R],
437        res_col: usize,
438        a: &A,
439        a_col: usize,
440        scratch: &mut Scratch<Self>,
441    ) where
442        R: VecZnxToMut,
443        A: VecZnxToRef,
444    {
445        let (tmp, _) = scratch.take_slice(module.vec_znx_split_ring_tmp_bytes() / size_of::<i64>());
446        vec_znx_split_ring::<R, A, Self>(res, res_col, a, a_col, tmp);
447    }
448}
449
450unsafe impl VecZnxMergeRingsTmpBytesImpl<Self> for FFT64Ref {
451    fn vec_znx_merge_rings_tmp_bytes_impl(module: &Module<Self>) -> usize {
452        vec_znx_merge_rings_tmp_bytes(module.n())
453    }
454}
455
456unsafe impl VecZnxMergeRingsImpl<Self> for FFT64Ref
457where
458    Module<Self>: VecZnxMergeRingsTmpBytes,
459{
460    fn vec_znx_merge_rings_impl<R, A>(
461        module: &Module<Self>,
462        res: &mut R,
463        res_col: usize,
464        a: &[A],
465        a_col: usize,
466        scratch: &mut Scratch<Self>,
467    ) where
468        R: VecZnxToMut,
469        A: VecZnxToRef,
470    {
471        let (tmp, _) = scratch.take_slice(module.vec_znx_merge_rings_tmp_bytes() / size_of::<i64>());
472        vec_znx_merge_rings::<R, A, Self>(res, res_col, a, a_col, tmp);
473    }
474}
475
476unsafe impl VecZnxSwitchRingImpl<Self> for FFT64Ref
477where
478    Self: VecZnxCopyImpl<Self>,
479{
480    fn vec_znx_switch_ring_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
481    where
482        R: VecZnxToMut,
483        A: VecZnxToRef,
484    {
485        vec_znx_switch_ring::<R, A, Self>(res, res_col, a, a_col);
486    }
487}
488
489unsafe impl VecZnxCopyImpl<Self> for FFT64Ref {
490    fn vec_znx_copy_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
491    where
492        R: VecZnxToMut,
493        A: VecZnxToRef,
494    {
495        vec_znx_copy::<R, A, Self>(res, res_col, a, a_col)
496    }
497}
498
499unsafe impl VecZnxFillUniformImpl<Self> for FFT64Ref {
500    fn vec_znx_fill_uniform_impl<R>(_module: &Module<Self>, base2k: usize, res: &mut R, res_col: usize, source: &mut Source)
501    where
502        R: VecZnxToMut,
503    {
504        vec_znx_fill_uniform_ref(base2k, res, res_col, source)
505    }
506}
507
508unsafe impl VecZnxFillNormalImpl<Self> for FFT64Ref {
509    fn vec_znx_fill_normal_impl<R>(
510        _module: &Module<Self>,
511        base2k: usize,
512        res: &mut R,
513        res_col: usize,
514        k: usize,
515        source: &mut Source,
516        sigma: f64,
517        bound: f64,
518    ) where
519        R: VecZnxToMut,
520    {
521        vec_znx_fill_normal_ref(base2k, res, res_col, k, sigma, bound, source);
522    }
523}
524
525unsafe impl VecZnxAddNormalImpl<Self> for FFT64Ref {
526    fn vec_znx_add_normal_impl<R>(
527        _module: &Module<Self>,
528        base2k: usize,
529        res: &mut R,
530        res_col: usize,
531        k: usize,
532        source: &mut Source,
533        sigma: f64,
534        bound: f64,
535    ) where
536        R: VecZnxToMut,
537    {
538        vec_znx_add_normal_ref(base2k, res, res_col, k, sigma, bound, source);
539    }
540}