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