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