1use itertools::izip;
  2use rand_distr::Normal;
  3
  4use poulpy_hal::{
  5    api::{
  6        TakeSlice, TakeVecZnx, VecZnxAddDistF64, VecZnxCopy, VecZnxFillDistF64, VecZnxNormalizeTmpBytes, VecZnxRotate,
  7        VecZnxRotateInplace, VecZnxSwithcDegree, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero,
  8    },
  9    layouts::{Backend, Module, ScalarZnx, ScalarZnxToRef, Scratch, VecZnx, VecZnxToMut, VecZnxToRef},
 10    oep::{
 11        TakeSliceImpl, TakeVecZnxImpl, VecZnxAddDistF64Impl, VecZnxAddImpl, VecZnxAddInplaceImpl, VecZnxAddNormalImpl,
 12        VecZnxAddScalarInplaceImpl, VecZnxAutomorphismImpl, VecZnxAutomorphismInplaceImpl, VecZnxCopyImpl, VecZnxFillDistF64Impl,
 13        VecZnxFillNormalImpl, VecZnxFillUniformImpl, VecZnxLshInplaceImpl, VecZnxMergeImpl, VecZnxMulXpMinusOneImpl,
 14        VecZnxMulXpMinusOneInplaceImpl, VecZnxNegateImpl, VecZnxNegateInplaceImpl, VecZnxNormalizeImpl,
 15        VecZnxNormalizeInplaceImpl, VecZnxNormalizeTmpBytesImpl, VecZnxRotateImpl, VecZnxRotateInplaceImpl, VecZnxRshInplaceImpl,
 16        VecZnxSplitImpl, VecZnxSubABInplaceImpl, VecZnxSubBAInplaceImpl, VecZnxSubImpl, VecZnxSubScalarInplaceImpl,
 17        VecZnxSwithcDegreeImpl,
 18    },
 19    source::Source,
 20};
 21
 22use crate::cpu_spqlios::{
 23    FFT64,
 24    ffi::{module::module_info_t, vec_znx, znx},
 25};
 26
 27unsafe impl VecZnxNormalizeTmpBytesImpl<Self> for FFT64 {
 28    fn vec_znx_normalize_tmp_bytes_impl(module: &Module<Self>) -> usize {
 29        unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(module.ptr() as *const module_info_t) as usize }
 30    }
 31}
 32
 33unsafe impl VecZnxNormalizeImpl<Self> for FFT64
 34where
 35    Self: TakeSliceImpl<Self> + VecZnxNormalizeTmpBytesImpl<Self>,
 36{
 37    fn vec_znx_normalize_impl<R, A>(
 38        module: &Module<Self>,
 39        basek: usize,
 40        res: &mut R,
 41        res_col: usize,
 42        a: &A,
 43        a_col: usize,
 44        scratch: &mut Scratch<Self>,
 45    ) where
 46        R: VecZnxToMut,
 47        A: VecZnxToRef,
 48    {
 49        let a: VecZnx<&[u8]> = a.to_ref();
 50        let mut res: VecZnx<&mut [u8]> = res.to_mut();
 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_normalize_tmp_bytes());
 58
 59        unsafe {
 60            vec_znx::vec_znx_normalize_base2k(
 61                module.ptr() as *const module_info_t,
 62                basek as u64,
 63                res.at_mut_ptr(res_col, 0),
 64                res.size() as u64,
 65                res.sl() as u64,
 66                a.at_ptr(a_col, 0),
 67                a.size() as u64,
 68                a.sl() as u64,
 69                tmp_bytes.as_mut_ptr(),
 70            );
 71        }
 72    }
 73}
 74
 75unsafe impl VecZnxNormalizeInplaceImpl<Self> for FFT64
 76where
 77    Self: TakeSliceImpl<Self> + VecZnxNormalizeTmpBytesImpl<Self>,
 78{
 79    fn vec_znx_normalize_inplace_impl<A>(
 80        module: &Module<Self>,
 81        basek: usize,
 82        a: &mut A,
 83        a_col: usize,
 84        scratch: &mut Scratch<Self>,
 85    ) where
 86        A: VecZnxToMut,
 87    {
 88        let mut a: VecZnx<&mut [u8]> = a.to_mut();
 89
 90        let (tmp_bytes, _) = scratch.take_slice(module.vec_znx_normalize_tmp_bytes());
 91
 92        unsafe {
 93            vec_znx::vec_znx_normalize_base2k(
 94                module.ptr() as *const module_info_t,
 95                basek as u64,
 96                a.at_mut_ptr(a_col, 0),
 97                a.size() as u64,
 98                a.sl() as u64,
 99                a.at_ptr(a_col, 0),
100                a.size() as u64,
101                a.sl() as u64,
102                tmp_bytes.as_mut_ptr(),
103            );
104        }
105    }
106}
107
108unsafe impl VecZnxAddImpl<Self> for FFT64 {
109    fn vec_znx_add_impl<R, A, C>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
110    where
111        R: VecZnxToMut,
112        A: VecZnxToRef,
113        C: VecZnxToRef,
114    {
115        let a: VecZnx<&[u8]> = a.to_ref();
116        let b: VecZnx<&[u8]> = b.to_ref();
117        let mut res: VecZnx<&mut [u8]> = res.to_mut();
118
119        #[cfg(debug_assertions)]
120        {
121            assert_eq!(a.n(), res.n());
122            assert_eq!(b.n(), res.n());
123            assert_ne!(a.as_ptr(), b.as_ptr());
124        }
125        unsafe {
126            vec_znx::vec_znx_add(
127                module.ptr() as *const module_info_t,
128                res.at_mut_ptr(res_col, 0),
129                res.size() as u64,
130                res.sl() as u64,
131                a.at_ptr(a_col, 0),
132                a.size() as u64,
133                a.sl() as u64,
134                b.at_ptr(b_col, 0),
135                b.size() as u64,
136                b.sl() as u64,
137            )
138        }
139    }
140}
141
142unsafe impl VecZnxAddInplaceImpl<Self> for FFT64 {
143    fn vec_znx_add_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
144    where
145        R: VecZnxToMut,
146        A: VecZnxToRef,
147    {
148        let a: VecZnx<&[u8]> = a.to_ref();
149        let mut res: VecZnx<&mut [u8]> = res.to_mut();
150
151        #[cfg(debug_assertions)]
152        {
153            assert_eq!(a.n(), res.n());
154        }
155        unsafe {
156            vec_znx::vec_znx_add(
157                module.ptr() as *const module_info_t,
158                res.at_mut_ptr(res_col, 0),
159                res.size() as u64,
160                res.sl() as u64,
161                a.at_ptr(a_col, 0),
162                a.size() as u64,
163                a.sl() as u64,
164                res.at_ptr(res_col, 0),
165                res.size() as u64,
166                res.sl() as u64,
167            )
168        }
169    }
170}
171
172unsafe impl VecZnxAddScalarInplaceImpl<Self> for FFT64 {
173    fn vec_znx_add_scalar_inplace_impl<R, A>(
174        module: &Module<Self>,
175        res: &mut R,
176        res_col: usize,
177        res_limb: usize,
178        a: &A,
179        a_col: usize,
180    ) where
181        R: VecZnxToMut,
182        A: ScalarZnxToRef,
183    {
184        let mut res: VecZnx<&mut [u8]> = res.to_mut();
185        let a: ScalarZnx<&[u8]> = a.to_ref();
186
187        #[cfg(debug_assertions)]
188        {
189            assert_eq!(a.n(), res.n());
190        }
191
192        unsafe {
193            vec_znx::vec_znx_add(
194                module.ptr() as *const module_info_t,
195                res.at_mut_ptr(res_col, res_limb),
196                1_u64,
197                res.sl() as u64,
198                a.at_ptr(a_col, 0),
199                a.size() as u64,
200                a.sl() as u64,
201                res.at_ptr(res_col, res_limb),
202                1_u64,
203                res.sl() as u64,
204            )
205        }
206    }
207}
208
209unsafe impl VecZnxSubImpl<Self> for FFT64 {
210    fn vec_znx_sub_impl<R, A, C>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize, b: &C, b_col: usize)
211    where
212        R: VecZnxToMut,
213        A: VecZnxToRef,
214        C: VecZnxToRef,
215    {
216        let a: VecZnx<&[u8]> = a.to_ref();
217        let b: VecZnx<&[u8]> = b.to_ref();
218        let mut res: VecZnx<&mut [u8]> = res.to_mut();
219
220        #[cfg(debug_assertions)]
221        {
222            assert_eq!(a.n(), res.n());
223            assert_eq!(b.n(), res.n());
224            assert_ne!(a.as_ptr(), b.as_ptr());
225        }
226        unsafe {
227            vec_znx::vec_znx_sub(
228                module.ptr() as *const module_info_t,
229                res.at_mut_ptr(res_col, 0),
230                res.size() as u64,
231                res.sl() as u64,
232                a.at_ptr(a_col, 0),
233                a.size() as u64,
234                a.sl() as u64,
235                b.at_ptr(b_col, 0),
236                b.size() as u64,
237                b.sl() as u64,
238            )
239        }
240    }
241}
242
243unsafe impl VecZnxSubABInplaceImpl<Self> for FFT64 {
244    fn vec_znx_sub_ab_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
245    where
246        R: VecZnxToMut,
247        A: VecZnxToRef,
248    {
249        let a: VecZnx<&[u8]> = a.to_ref();
250        let mut res: VecZnx<&mut [u8]> = res.to_mut();
251        #[cfg(debug_assertions)]
252        {
253            assert_eq!(a.n(), res.n());
254        }
255        unsafe {
256            vec_znx::vec_znx_sub(
257                module.ptr() as *const module_info_t,
258                res.at_mut_ptr(res_col, 0),
259                res.size() as u64,
260                res.sl() as u64,
261                res.at_ptr(res_col, 0),
262                res.size() as u64,
263                res.sl() as u64,
264                a.at_ptr(a_col, 0),
265                a.size() as u64,
266                a.sl() as u64,
267            )
268        }
269    }
270}
271
272unsafe impl VecZnxSubBAInplaceImpl<Self> for FFT64 {
273    fn vec_znx_sub_ba_inplace_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
274    where
275        R: VecZnxToMut,
276        A: VecZnxToRef,
277    {
278        let a: VecZnx<&[u8]> = a.to_ref();
279        let mut res: VecZnx<&mut [u8]> = res.to_mut();
280        #[cfg(debug_assertions)]
281        {
282            assert_eq!(a.n(), res.n());
283        }
284        unsafe {
285            vec_znx::vec_znx_sub(
286                module.ptr() as *const module_info_t,
287                res.at_mut_ptr(res_col, 0),
288                res.size() as u64,
289                res.sl() as u64,
290                a.at_ptr(a_col, 0),
291                a.size() as u64,
292                a.sl() as u64,
293                res.at_ptr(res_col, 0),
294                res.size() as u64,
295                res.sl() as u64,
296            )
297        }
298    }
299}
300
301unsafe impl VecZnxSubScalarInplaceImpl<Self> for FFT64 {
302    fn vec_znx_sub_scalar_inplace_impl<R, A>(
303        module: &Module<Self>,
304        res: &mut R,
305        res_col: usize,
306        res_limb: usize,
307        a: &A,
308        a_col: usize,
309    ) where
310        R: VecZnxToMut,
311        A: ScalarZnxToRef,
312    {
313        let mut res: VecZnx<&mut [u8]> = res.to_mut();
314        let a: ScalarZnx<&[u8]> = a.to_ref();
315
316        #[cfg(debug_assertions)]
317        {
318            assert_eq!(a.n(), res.n());
319        }
320
321        unsafe {
322            vec_znx::vec_znx_sub(
323                module.ptr() as *const module_info_t,
324                res.at_mut_ptr(res_col, res_limb),
325                1_u64,
326                res.sl() as u64,
327                a.at_ptr(a_col, 0),
328                a.size() as u64,
329                a.sl() as u64,
330                res.at_ptr(res_col, res_limb),
331                1_u64,
332                res.sl() as u64,
333            )
334        }
335    }
336}
337
338unsafe impl VecZnxNegateImpl<Self> for FFT64 {
339    fn vec_znx_negate_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
340    where
341        R: VecZnxToMut,
342        A: VecZnxToRef,
343    {
344        let a: VecZnx<&[u8]> = a.to_ref();
345        let mut res: VecZnx<&mut [u8]> = res.to_mut();
346        #[cfg(debug_assertions)]
347        {
348            assert_eq!(a.n(), res.n());
349        }
350        unsafe {
351            vec_znx::vec_znx_negate(
352                module.ptr() as *const module_info_t,
353                res.at_mut_ptr(res_col, 0),
354                res.size() as u64,
355                res.sl() as u64,
356                a.at_ptr(a_col, 0),
357                a.size() as u64,
358                a.sl() as u64,
359            )
360        }
361    }
362}
363
364unsafe impl VecZnxNegateInplaceImpl<Self> for FFT64 {
365    fn vec_znx_negate_inplace_impl<A>(module: &Module<Self>, a: &mut A, a_col: usize)
366    where
367        A: VecZnxToMut,
368    {
369        let mut a: VecZnx<&mut [u8]> = a.to_mut();
370        unsafe {
371            vec_znx::vec_znx_negate(
372                module.ptr() as *const module_info_t,
373                a.at_mut_ptr(a_col, 0),
374                a.size() as u64,
375                a.sl() as u64,
376                a.at_ptr(a_col, 0),
377                a.size() as u64,
378                a.sl() as u64,
379            )
380        }
381    }
382}
383
384unsafe impl VecZnxLshInplaceImpl<Self> for FFT64 {
385    fn vec_znx_lsh_inplace_impl<A>(_module: &Module<Self>, basek: usize, k: usize, a: &mut A)
386    where
387        A: VecZnxToMut,
388    {
389        vec_znx_lsh_inplace_ref(basek, k, a)
390    }
391}
392
393pub fn vec_znx_lsh_inplace_ref<A>(basek: usize, k: usize, a: &mut A)
394where
395    A: VecZnxToMut,
396{
397    let mut a: VecZnx<&mut [u8]> = a.to_mut();
398
399    let n: usize = a.n();
400    let cols: usize = a.cols();
401    let size: usize = a.size();
402    let steps: usize = k / basek;
403
404    a.raw_mut().rotate_left(n * steps * cols);
405    (0..cols).for_each(|i| {
406        (size - steps..size).for_each(|j| {
407            a.zero_at(i, j);
408        })
409    });
410
411    let k_rem: usize = k % basek;
412
413    if k_rem != 0 {
414        let shift: usize = i64::BITS as usize - k_rem;
415        (0..cols).for_each(|i| {
416            (0..steps).for_each(|j| {
417                a.at_mut(i, j).iter_mut().for_each(|xi| {
418                    *xi <<= shift;
419                });
420            });
421        });
422    }
423}
424
425unsafe impl VecZnxRshInplaceImpl<Self> for FFT64 {
426    fn vec_znx_rsh_inplace_impl<A>(_module: &Module<Self>, basek: usize, k: usize, a: &mut A)
427    where
428        A: VecZnxToMut,
429    {
430        vec_znx_rsh_inplace_ref(basek, k, a)
431    }
432}
433
434pub fn vec_znx_rsh_inplace_ref<A>(basek: usize, k: usize, a: &mut A)
435where
436    A: VecZnxToMut,
437{
438    let mut a: VecZnx<&mut [u8]> = a.to_mut();
439    let n: usize = a.n();
440    let cols: usize = a.cols();
441    let size: usize = a.size();
442    let steps: usize = k / basek;
443
444    a.raw_mut().rotate_right(n * steps * cols);
445    (0..cols).for_each(|i| {
446        (0..steps).for_each(|j| {
447            a.zero_at(i, j);
448        })
449    });
450
451    let k_rem: usize = k % basek;
452
453    if k_rem != 0 {
454        let mut carry: Vec<i64> = vec![0i64; n]; // ALLOC (but small so OK)
455        let shift: usize = i64::BITS as usize - k_rem;
456        (0..cols).for_each(|i| {
457            carry.fill(0);
458            (steps..size).for_each(|j| {
459                izip!(carry.iter_mut(), a.at_mut(i, j).iter_mut()).for_each(|(ci, xi)| {
460                    *xi += *ci << basek;
461                    *ci = (*xi << shift) >> shift;
462                    *xi = (*xi - *ci) >> k_rem;
463                });
464            });
465        })
466    }
467}
468
469unsafe impl VecZnxRotateImpl<Self> for FFT64 {
470    fn vec_znx_rotate_impl<R, A>(_module: &Module<Self>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
471    where
472        R: VecZnxToMut,
473        A: VecZnxToRef,
474    {
475        let a: VecZnx<&[u8]> = a.to_ref();
476        let mut res: VecZnx<&mut [u8]> = res.to_mut();
477        #[cfg(debug_assertions)]
478        {
479            assert_eq!(res.n(), a.n());
480        }
481        unsafe {
482            (0..a.size()).for_each(|j| {
483                znx::znx_rotate_i64(
484                    a.n() as u64,
485                    k,
486                    res.at_mut_ptr(res_col, j),
487                    a.at_ptr(a_col, j),
488                );
489            });
490        }
491    }
492}
493
494unsafe impl VecZnxRotateInplaceImpl<Self> for FFT64 {
495    fn vec_znx_rotate_inplace_impl<A>(_module: &Module<Self>, k: i64, a: &mut A, a_col: usize)
496    where
497        A: VecZnxToMut,
498    {
499        let mut a: VecZnx<&mut [u8]> = a.to_mut();
500        unsafe {
501            (0..a.size()).for_each(|j| {
502                znx::znx_rotate_inplace_i64(a.n() as u64, k, a.at_mut_ptr(a_col, j));
503            });
504        }
505    }
506}
507
508unsafe impl VecZnxAutomorphismImpl<Self> for FFT64 {
509    fn vec_znx_automorphism_impl<R, A>(module: &Module<Self>, k: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
510    where
511        R: VecZnxToMut,
512        A: VecZnxToRef,
513    {
514        let a: VecZnx<&[u8]> = a.to_ref();
515        let mut res: VecZnx<&mut [u8]> = res.to_mut();
516        #[cfg(debug_assertions)]
517        {
518            assert_eq!(a.n(), res.n());
519        }
520        unsafe {
521            vec_znx::vec_znx_automorphism(
522                module.ptr() as *const module_info_t,
523                k,
524                res.at_mut_ptr(res_col, 0),
525                res.size() as u64,
526                res.sl() as u64,
527                a.at_ptr(a_col, 0),
528                a.size() as u64,
529                a.sl() as u64,
530            )
531        }
532    }
533}
534
535unsafe impl VecZnxAutomorphismInplaceImpl<Self> for FFT64 {
536    fn vec_znx_automorphism_inplace_impl<A>(module: &Module<Self>, k: i64, a: &mut A, a_col: usize)
537    where
538        A: VecZnxToMut,
539    {
540        let mut a: VecZnx<&mut [u8]> = a.to_mut();
541        #[cfg(debug_assertions)]
542        {
543            assert!(
544                k & 1 != 0,
545                "invalid galois element: must be odd but is {}",
546                k
547            );
548        }
549        unsafe {
550            vec_znx::vec_znx_automorphism(
551                module.ptr() as *const module_info_t,
552                k,
553                a.at_mut_ptr(a_col, 0),
554                a.size() as u64,
555                a.sl() as u64,
556                a.at_ptr(a_col, 0),
557                a.size() as u64,
558                a.sl() as u64,
559            )
560        }
561    }
562}
563
564unsafe impl VecZnxMulXpMinusOneImpl<Self> for FFT64 {
565    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)
566    where
567        R: VecZnxToMut,
568        A: VecZnxToRef,
569    {
570        let a: VecZnx<&[u8]> = a.to_ref();
571        let mut res: VecZnx<&mut [u8]> = res.to_mut();
572        #[cfg(debug_assertions)]
573        {
574            assert_eq!(a.n(), res.n());
575            assert_eq!(res.n(), res.n());
576        }
577        unsafe {
578            vec_znx::vec_znx_mul_xp_minus_one(
579                module.ptr() as *const module_info_t,
580                p,
581                res.at_mut_ptr(res_col, 0),
582                res.size() as u64,
583                res.sl() as u64,
584                a.at_ptr(a_col, 0),
585                a.size() as u64,
586                a.sl() as u64,
587            )
588        }
589    }
590}
591
592unsafe impl VecZnxMulXpMinusOneInplaceImpl<Self> for FFT64 {
593    fn vec_znx_mul_xp_minus_one_inplace_impl<R>(module: &Module<Self>, p: i64, res: &mut R, res_col: usize)
594    where
595        R: VecZnxToMut,
596    {
597        let mut res: VecZnx<&mut [u8]> = res.to_mut();
598        #[cfg(debug_assertions)]
599        {
600            assert_eq!(res.n(), res.n());
601        }
602        unsafe {
603            vec_znx::vec_znx_mul_xp_minus_one(
604                module.ptr() as *const module_info_t,
605                p,
606                res.at_mut_ptr(res_col, 0),
607                res.size() as u64,
608                res.sl() as u64,
609                res.at_ptr(res_col, 0),
610                res.size() as u64,
611                res.sl() as u64,
612            )
613        }
614    }
615}
616
617unsafe impl VecZnxSplitImpl<Self> for FFT64
618where
619    Self: TakeVecZnxImpl<Self>
620        + TakeVecZnxImpl<Self>
621        + VecZnxSwithcDegreeImpl<Self>
622        + VecZnxRotateImpl<Self>
623        + VecZnxRotateInplaceImpl<Self>,
624{
625    fn vec_znx_split_impl<R, A>(
626        module: &Module<Self>,
627        res: &mut [R],
628        res_col: usize,
629        a: &A,
630        a_col: usize,
631        scratch: &mut Scratch<Self>,
632    ) where
633        R: VecZnxToMut,
634        A: VecZnxToRef,
635    {
636        vec_znx_split_ref(module, res, res_col, a, a_col, scratch)
637    }
638}
639
640pub fn vec_znx_split_ref<R, A, B>(
641    module: &Module<B>,
642    res: &mut [R],
643    res_col: usize,
644    a: &A,
645    a_col: usize,
646    scratch: &mut Scratch<B>,
647) where
648    B: Backend + TakeVecZnxImpl<B> + VecZnxSwithcDegreeImpl<B> + VecZnxRotateImpl<B> + VecZnxRotateInplaceImpl<B>,
649    R: VecZnxToMut,
650    A: VecZnxToRef,
651{
652    let a: VecZnx<&[u8]> = a.to_ref();
653
654    let (n_in, n_out) = (a.n(), res[0].to_mut().n());
655
656    let (mut buf, _) = scratch.take_vec_znx(n_in.max(n_out), 1, a.size());
657
658    debug_assert!(
659        n_out < n_in,
660        "invalid a: output ring degree should be smaller"
661    );
662    res[1..].iter_mut().for_each(|bi| {
663        debug_assert_eq!(
664            bi.to_mut().n(),
665            n_out,
666            "invalid input a: all VecZnx must have the same degree"
667        )
668    });
669
670    res.iter_mut().enumerate().for_each(|(i, bi)| {
671        if i == 0 {
672            module.vec_znx_switch_degree(bi, res_col, &a, a_col);
673            module.vec_znx_rotate(-1, &mut buf, 0, &a, a_col);
674        } else {
675            module.vec_znx_switch_degree(bi, res_col, &buf, a_col);
676            module.vec_znx_rotate_inplace(-1, &mut buf, a_col);
677        }
678    })
679}
680
681unsafe impl VecZnxMergeImpl<Self> for FFT64
682where
683    Self: VecZnxSwithcDegreeImpl<Self> + VecZnxRotateInplaceImpl<Self>,
684{
685    fn vec_znx_merge_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &[A], a_col: usize)
686    where
687        R: VecZnxToMut,
688        A: VecZnxToRef,
689    {
690        vec_znx_merge_ref(module, res, res_col, a, a_col)
691    }
692}
693
694pub fn vec_znx_merge_ref<R, A, B>(module: &Module<B>, res: &mut R, res_col: usize, a: &[A], a_col: usize)
695where
696    B: Backend + VecZnxSwithcDegreeImpl<B> + VecZnxRotateInplaceImpl<B>,
697    R: VecZnxToMut,
698    A: VecZnxToRef,
699{
700    let mut res: VecZnx<&mut [u8]> = res.to_mut();
701
702    let (n_in, n_out) = (res.n(), a[0].to_ref().n());
703
704    debug_assert!(
705        n_out < n_in,
706        "invalid a: output ring degree should be smaller"
707    );
708    a[1..].iter().for_each(|ai| {
709        debug_assert_eq!(
710            ai.to_ref().n(),
711            n_out,
712            "invalid input a: all VecZnx must have the same degree"
713        )
714    });
715
716    a.iter().for_each(|ai| {
717        module.vec_znx_switch_degree(&mut res, res_col, ai, a_col);
718        module.vec_znx_rotate_inplace(-1, &mut res, res_col);
719    });
720
721    module.vec_znx_rotate_inplace(a.len() as i64, &mut res, res_col);
722}
723
724unsafe impl VecZnxSwithcDegreeImpl<Self> for FFT64
725where
726    Self: VecZnxCopyImpl<Self>,
727{
728    fn vec_znx_switch_degree_impl<R, A>(module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
729    where
730        R: VecZnxToMut,
731        A: VecZnxToRef,
732    {
733        vec_znx_switch_degree_ref(module, res, res_col, a, a_col)
734    }
735}
736
737pub fn vec_znx_switch_degree_ref<R, A, B>(module: &Module<B>, res: &mut R, res_col: usize, a: &A, a_col: usize)
738where
739    B: Backend + VecZnxCopyImpl<B>,
740    R: VecZnxToMut,
741    A: VecZnxToRef,
742{
743    let a: VecZnx<&[u8]> = a.to_ref();
744    let mut res: VecZnx<&mut [u8]> = res.to_mut();
745
746    let (n_in, n_out) = (a.n(), res.n());
747
748    if n_in == n_out {
749        module.vec_znx_copy(&mut res, res_col, &a, a_col);
750        return;
751    }
752
753    let (gap_in, gap_out): (usize, usize);
754    if n_in > n_out {
755        (gap_in, gap_out) = (n_in / n_out, 1)
756    } else {
757        (gap_in, gap_out) = (1, n_out / n_in);
758        res.zero();
759    }
760
761    let size: usize = a.size().min(res.size());
762
763    (0..size).for_each(|i| {
764        izip!(
765            a.at(a_col, i).iter().step_by(gap_in),
766            res.at_mut(res_col, i).iter_mut().step_by(gap_out)
767        )
768        .for_each(|(x_in, x_out)| *x_out = *x_in);
769    });
770}
771
772unsafe impl VecZnxCopyImpl<Self> for FFT64 {
773    fn vec_znx_copy_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
774    where
775        R: VecZnxToMut,
776        A: VecZnxToRef,
777    {
778        vec_znx_copy_ref(res, res_col, a, a_col)
779    }
780}
781
782pub fn vec_znx_copy_ref<R, A>(res: &mut R, res_col: usize, a: &A, a_col: usize)
783where
784    R: VecZnxToMut,
785    A: VecZnxToRef,
786{
787    let mut res_mut: VecZnx<&mut [u8]> = res.to_mut();
788    let a_ref: VecZnx<&[u8]> = a.to_ref();
789
790    let min_size: usize = res_mut.size().min(a_ref.size());
791
792    (0..min_size).for_each(|j| {
793        res_mut
794            .at_mut(res_col, j)
795            .copy_from_slice(a_ref.at(a_col, j));
796    });
797    (min_size..res_mut.size()).for_each(|j| {
798        res_mut.zero_at(res_col, j);
799    })
800}
801
802unsafe impl VecZnxFillUniformImpl<Self> for FFT64 {
803    fn vec_znx_fill_uniform_impl<R>(
804        _module: &Module<Self>,
805        basek: usize,
806        res: &mut R,
807        res_col: usize,
808        k: usize,
809        source: &mut Source,
810    ) where
811        R: VecZnxToMut,
812    {
813        let mut a: VecZnx<&mut [u8]> = res.to_mut();
814        let base2k: u64 = 1 << basek;
815        let mask: u64 = base2k - 1;
816        let base2k_half: i64 = (base2k >> 1) as i64;
817        (0..k.div_ceil(basek)).for_each(|j| {
818            a.at_mut(res_col, j)
819                .iter_mut()
820                .for_each(|x| *x = (source.next_u64n(base2k, mask) as i64) - base2k_half);
821        })
822    }
823}
824
825unsafe impl VecZnxFillDistF64Impl<Self> for FFT64 {
826    fn vec_znx_fill_dist_f64_impl<R, D: rand::prelude::Distribution<f64>>(
827        _module: &Module<Self>,
828        basek: usize,
829        res: &mut R,
830        res_col: usize,
831        k: usize,
832        source: &mut Source,
833        dist: D,
834        bound: f64,
835    ) where
836        R: VecZnxToMut,
837    {
838        let mut a: VecZnx<&mut [u8]> = res.to_mut();
839        assert!(
840            (bound.log2().ceil() as i64) < 64,
841            "invalid bound: ceil(log2(bound))={} > 63",
842            (bound.log2().ceil() as i64)
843        );
844
845        let limb: usize = k.div_ceil(basek) - 1;
846        let basek_rem: usize = (limb + 1) * basek - k;
847
848        if basek_rem != 0 {
849            a.at_mut(res_col, limb).iter_mut().for_each(|a| {
850                let mut dist_f64: f64 = dist.sample(source);
851                while dist_f64.abs() > bound {
852                    dist_f64 = dist.sample(source)
853                }
854                *a = (dist_f64.round() as i64) << basek_rem;
855            });
856        } else {
857            a.at_mut(res_col, limb).iter_mut().for_each(|a| {
858                let mut dist_f64: f64 = dist.sample(source);
859                while dist_f64.abs() > bound {
860                    dist_f64 = dist.sample(source)
861                }
862                *a = dist_f64.round() as i64
863            });
864        }
865    }
866}
867
868unsafe impl VecZnxAddDistF64Impl<Self> for FFT64 {
869    fn vec_znx_add_dist_f64_impl<R, D: rand::prelude::Distribution<f64>>(
870        _module: &Module<Self>,
871        basek: usize,
872        res: &mut R,
873        res_col: usize,
874        k: usize,
875        source: &mut Source,
876        dist: D,
877        bound: f64,
878    ) where
879        R: VecZnxToMut,
880    {
881        let mut a: VecZnx<&mut [u8]> = res.to_mut();
882        assert!(
883            (bound.log2().ceil() as i64) < 64,
884            "invalid bound: ceil(log2(bound))={} > 63",
885            (bound.log2().ceil() as i64)
886        );
887
888        let limb: usize = k.div_ceil(basek) - 1;
889        let basek_rem: usize = (limb + 1) * basek - k;
890
891        if basek_rem != 0 {
892            a.at_mut(res_col, limb).iter_mut().for_each(|a| {
893                let mut dist_f64: f64 = dist.sample(source);
894                while dist_f64.abs() > bound {
895                    dist_f64 = dist.sample(source)
896                }
897                *a += (dist_f64.round() as i64) << basek_rem;
898            });
899        } else {
900            a.at_mut(res_col, limb).iter_mut().for_each(|a| {
901                let mut dist_f64: f64 = dist.sample(source);
902                while dist_f64.abs() > bound {
903                    dist_f64 = dist.sample(source)
904                }
905                *a += dist_f64.round() as i64
906            });
907        }
908    }
909}
910
911unsafe impl VecZnxFillNormalImpl<Self> for FFT64
912where
913    Self: VecZnxFillDistF64Impl<Self>,
914{
915    fn vec_znx_fill_normal_impl<R>(
916        module: &Module<Self>,
917        basek: usize,
918        res: &mut R,
919        res_col: usize,
920        k: usize,
921        source: &mut Source,
922        sigma: f64,
923        bound: f64,
924    ) where
925        R: VecZnxToMut,
926    {
927        module.vec_znx_fill_dist_f64(
928            basek,
929            res,
930            res_col,
931            k,
932            source,
933            Normal::new(0.0, sigma).unwrap(),
934            bound,
935        );
936    }
937}
938
939unsafe impl VecZnxAddNormalImpl<Self> for FFT64
940where
941    Self: VecZnxAddDistF64Impl<Self>,
942{
943    fn vec_znx_add_normal_impl<R>(
944        module: &Module<Self>,
945        basek: usize,
946        res: &mut R,
947        res_col: usize,
948        k: usize,
949        source: &mut Source,
950        sigma: f64,
951        bound: f64,
952    ) where
953        R: VecZnxToMut,
954    {
955        module.vec_znx_add_dist_f64(
956            basek,
957            res,
958            res_col,
959            k,
960            source,
961            Normal::new(0.0, sigma).unwrap(),
962            bound,
963        );
964    }
965}