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