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