1use crate::FFT64Avx;
  2use poulpy_hal::{
  3    api::{TakeSlice, VecZnxBigAutomorphismInplaceTmpBytes, VecZnxBigNormalizeTmpBytes},
  4    layouts::{
  5        Backend, Module, Scratch, VecZnx, VecZnxBig, VecZnxBigOwned, VecZnxBigToMut, VecZnxBigToRef, VecZnxToMut, VecZnxToRef,
  6        ZnxInfos, ZnxView, ZnxViewMut,
  7    },
  8    oep::{
  9        TakeSliceImpl, VecZnxBigAddImpl, VecZnxBigAddInplaceImpl, VecZnxBigAddNormalImpl, VecZnxBigAddSmallImpl,
 10        VecZnxBigAddSmallInplaceImpl, VecZnxBigAllocBytesImpl, VecZnxBigAllocImpl, VecZnxBigAutomorphismImpl,
 11        VecZnxBigAutomorphismInplaceImpl, VecZnxBigAutomorphismInplaceTmpBytesImpl, VecZnxBigFromBytesImpl,
 12        VecZnxBigFromSmallImpl, VecZnxBigNegateImpl, VecZnxBigNegateInplaceImpl, VecZnxBigNormalizeImpl,
 13        VecZnxBigNormalizeTmpBytesImpl, VecZnxBigSubImpl, VecZnxBigSubInplaceImpl, VecZnxBigSubNegateInplaceImpl,
 14        VecZnxBigSubSmallAImpl, VecZnxBigSubSmallBImpl, VecZnxBigSubSmallInplaceImpl, VecZnxBigSubSmallNegateInplaceImpl,
 15    },
 16    reference::{
 17        fft64::vec_znx_big::{
 18            vec_znx_big_add, vec_znx_big_add_inplace, vec_znx_big_add_normal_ref, vec_znx_big_add_small,
 19            vec_znx_big_add_small_inplace, vec_znx_big_automorphism, vec_znx_big_automorphism_inplace,
 20            vec_znx_big_automorphism_inplace_tmp_bytes, vec_znx_big_negate, vec_znx_big_negate_inplace, vec_znx_big_normalize,
 21            vec_znx_big_normalize_tmp_bytes, vec_znx_big_sub, vec_znx_big_sub_inplace, vec_znx_big_sub_negate_inplace,
 22            vec_znx_big_sub_small_a, vec_znx_big_sub_small_a_inplace, vec_znx_big_sub_small_b, vec_znx_big_sub_small_b_inplace,
 23        },
 24        znx::{znx_copy_ref, znx_zero_ref},
 25    },
 26    source::Source,
 27};
 28
 29unsafe impl VecZnxBigAllocBytesImpl<Self> for FFT64Avx {
 30    fn vec_znx_big_bytes_of_impl(n: usize, cols: usize, size: usize) -> usize {
 31        Self::layout_big_word_count() * n * cols * size * size_of::<f64>()
 32    }
 33}
 34
 35unsafe impl VecZnxBigAllocImpl<Self> for FFT64Avx {
 36    fn vec_znx_big_alloc_impl(n: usize, cols: usize, size: usize) -> VecZnxBigOwned<Self> {
 37        VecZnxBig::alloc(n, cols, size)
 38    }
 39}
 40
 41unsafe impl VecZnxBigFromBytesImpl<Self> for FFT64Avx {
 42    fn vec_znx_big_from_bytes_impl(n: usize, cols: usize, size: usize, bytes: Vec<u8>) -> VecZnxBigOwned<Self> {
 43        VecZnxBig::from_bytes(n, cols, size, bytes)
 44    }
 45}
 46
 47unsafe impl VecZnxBigFromSmallImpl<Self> for FFT64Avx {
 48    fn vec_znx_big_from_small_impl<R, A>(res: &mut R, res_col: usize, a: &A, a_col: usize)
 49    where
 50        R: VecZnxBigToMut<Self>,
 51        A: VecZnxToRef,
 52    {
 53        let mut res: VecZnxBig<&mut [u8], FFT64Avx> = res.to_mut();
 54        let a: VecZnx<&[u8]> = a.to_ref();
 55
 56        #[cfg(debug_assertions)]
 57        {
 58            assert_eq!(res.n(), a.n());
 59        }
 60
 61        let res_size: usize = res.size();
 62        let a_size: usize = a.size();
 63
 64        let min_size: usize = res_size.min(a_size);
 65
 66        for j in 0..min_size {
 67            znx_copy_ref(res.at_mut(res_col, j), a.at(a_col, j));
 68        }
 69
 70        for j in min_size..res_size {
 71            znx_zero_ref(res.at_mut(res_col, j));
 72        }
 73    }
 74}
 75
 76unsafe impl VecZnxBigAddNormalImpl<Self> for FFT64Avx {
 77    fn add_normal_impl<R: VecZnxBigToMut<Self>>(
 78        _module: &Module<Self>,
 79        base2k: usize,
 80        res: &mut R,
 81        res_col: usize,
 82        k: usize,
 83        source: &mut Source,
 84        sigma: f64,
 85        bound: f64,
 86    ) {
 87        vec_znx_big_add_normal_ref(base2k, res, res_col, k, sigma, bound, source);
 88    }
 89}
 90
 91unsafe impl VecZnxBigAddImpl<Self> for FFT64Avx {
 92    /// Adds `a` to `b` and stores the result on `c`.
 93    fn vec_znx_big_add_impl<R, A, B>(
 94        _module: &Module<Self>,
 95        res: &mut R,
 96        res_col: usize,
 97        a: &A,
 98        a_col: usize,
 99        b: &B,
100        b_col: usize,
101    ) where
102        R: VecZnxBigToMut<Self>,
103        A: VecZnxBigToRef<Self>,
104        B: VecZnxBigToRef<Self>,
105    {
106        vec_znx_big_add(res, res_col, a, a_col, b, b_col);
107    }
108}
109
110unsafe impl VecZnxBigAddInplaceImpl<Self> for FFT64Avx {
111    /// Adds `a` to `b` and stores the result on `b`.
112    fn vec_znx_big_add_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
113    where
114        R: VecZnxBigToMut<Self>,
115        A: VecZnxBigToRef<Self>,
116    {
117        vec_znx_big_add_inplace(res, res_col, a, a_col);
118    }
119}
120
121unsafe impl VecZnxBigAddSmallImpl<Self> for FFT64Avx {
122    /// Adds `a` to `b` and stores the result on `c`.
123    fn vec_znx_big_add_small_impl<R, A, B>(
124        _module: &Module<Self>,
125        res: &mut R,
126        res_col: usize,
127        a: &A,
128        a_col: usize,
129        b: &B,
130        b_col: usize,
131    ) where
132        R: VecZnxBigToMut<Self>,
133        A: VecZnxBigToRef<Self>,
134        B: VecZnxToRef,
135    {
136        vec_znx_big_add_small(res, res_col, a, a_col, b, b_col);
137    }
138}
139
140unsafe impl VecZnxBigAddSmallInplaceImpl<Self> for FFT64Avx {
141    /// Adds `a` to `b` and stores the result on `b`.
142    fn vec_znx_big_add_small_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
143    where
144        R: VecZnxBigToMut<Self>,
145        A: VecZnxToRef,
146    {
147        vec_znx_big_add_small_inplace(res, res_col, a, a_col);
148    }
149}
150
151unsafe impl VecZnxBigSubImpl<Self> for FFT64Avx {
152    /// Subtracts `a` to `b` and stores the result on `c`.
153    fn vec_znx_big_sub_impl<R, A, B>(
154        _module: &Module<Self>,
155        res: &mut R,
156        res_col: usize,
157        a: &A,
158        a_col: usize,
159        b: &B,
160        b_col: usize,
161    ) where
162        R: VecZnxBigToMut<Self>,
163        A: VecZnxBigToRef<Self>,
164        B: VecZnxBigToRef<Self>,
165    {
166        vec_znx_big_sub(res, res_col, a, a_col, b, b_col);
167    }
168}
169
170unsafe impl VecZnxBigSubInplaceImpl<Self> for FFT64Avx {
171    /// Subtracts `a` from `b` and stores the result on `b`.
172    fn vec_znx_big_sub_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
173    where
174        R: VecZnxBigToMut<Self>,
175        A: VecZnxBigToRef<Self>,
176    {
177        vec_znx_big_sub_inplace(res, res_col, a, a_col);
178    }
179}
180
181unsafe impl VecZnxBigSubNegateInplaceImpl<Self> for FFT64Avx {
182    /// Subtracts `b` from `a` and stores the result on `b`.
183    fn vec_znx_big_sub_negate_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
184    where
185        R: VecZnxBigToMut<Self>,
186        A: VecZnxBigToRef<Self>,
187    {
188        vec_znx_big_sub_negate_inplace(res, res_col, a, a_col);
189    }
190}
191
192unsafe impl VecZnxBigSubSmallAImpl<Self> for FFT64Avx {
193    /// Subtracts `b` from `a` and stores the result on `c`.
194    fn vec_znx_big_sub_small_a_impl<R, A, B>(
195        _module: &Module<Self>,
196        res: &mut R,
197        res_col: usize,
198        a: &A,
199        a_col: usize,
200        b: &B,
201        b_col: usize,
202    ) where
203        R: VecZnxBigToMut<Self>,
204        A: VecZnxToRef,
205        B: VecZnxBigToRef<Self>,
206    {
207        vec_znx_big_sub_small_a(res, res_col, a, a_col, b, b_col);
208    }
209}
210
211unsafe impl VecZnxBigSubSmallInplaceImpl<Self> for FFT64Avx {
212    /// Subtracts `a` from `res` and stores the result on `res`.
213    fn vec_znx_big_sub_small_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
214    where
215        R: VecZnxBigToMut<Self>,
216        A: VecZnxToRef,
217    {
218        vec_znx_big_sub_small_a_inplace(res, res_col, a, a_col);
219    }
220}
221
222unsafe impl VecZnxBigSubSmallBImpl<Self> for FFT64Avx {
223    /// Subtracts `b` from `a` and stores the result on `c`.
224    fn vec_znx_big_sub_small_b_impl<R, A, B>(
225        _module: &Module<Self>,
226        res: &mut R,
227        res_col: usize,
228        a: &A,
229        a_col: usize,
230        b: &B,
231        b_col: usize,
232    ) where
233        R: VecZnxBigToMut<Self>,
234        A: VecZnxBigToRef<Self>,
235        B: VecZnxToRef,
236    {
237        vec_znx_big_sub_small_b(res, res_col, a, a_col, b, b_col);
238    }
239}
240
241unsafe impl VecZnxBigSubSmallNegateInplaceImpl<Self> for FFT64Avx {
242    /// Subtracts `res` from `a` and stores the result on `res`.
243    fn vec_znx_big_sub_small_negate_inplace_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
244    where
245        R: VecZnxBigToMut<Self>,
246        A: VecZnxToRef,
247    {
248        vec_znx_big_sub_small_b_inplace(res, res_col, a, a_col);
249    }
250}
251
252unsafe impl VecZnxBigNegateImpl<Self> for FFT64Avx {
253    fn vec_znx_big_negate_impl<R, A>(_module: &Module<Self>, res: &mut R, res_col: usize, a: &A, a_col: usize)
254    where
255        R: VecZnxBigToMut<Self>,
256        A: VecZnxBigToRef<Self>,
257    {
258        vec_znx_big_negate(res, res_col, a, a_col);
259    }
260}
261
262unsafe impl VecZnxBigNegateInplaceImpl<Self> for FFT64Avx {
263    fn vec_znx_big_negate_inplace_impl<R>(_module: &Module<Self>, res: &mut R, res_col: usize)
264    where
265        R: VecZnxBigToMut<Self>,
266    {
267        vec_znx_big_negate_inplace(res, res_col);
268    }
269}
270
271unsafe impl VecZnxBigNormalizeTmpBytesImpl<Self> for FFT64Avx {
272    fn vec_znx_big_normalize_tmp_bytes_impl(module: &Module<Self>) -> usize {
273        vec_znx_big_normalize_tmp_bytes(module.n())
274    }
275}
276
277unsafe impl VecZnxBigNormalizeImpl<Self> for FFT64Avx
278where
279    Self: TakeSliceImpl<Self>,
280{
281    fn vec_znx_big_normalize_impl<R, A>(
282        module: &Module<Self>,
283        res_basek: usize,
284        res: &mut R,
285        res_col: usize,
286        a_basek: usize,
287        a: &A,
288        a_col: usize,
289        scratch: &mut Scratch<Self>,
290    ) where
291        R: VecZnxToMut,
292        A: VecZnxBigToRef<Self>,
293    {
294        let (carry, _) = scratch.take_slice(module.vec_znx_big_normalize_tmp_bytes() / size_of::<i64>());
295        vec_znx_big_normalize(res_basek, res, res_col, a_basek, a, a_col, carry);
296    }
297}
298
299unsafe impl VecZnxBigAutomorphismImpl<Self> for FFT64Avx {
300    /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
301    fn vec_znx_big_automorphism_impl<R, A>(_module: &Module<Self>, p: i64, res: &mut R, res_col: usize, a: &A, a_col: usize)
302    where
303        R: VecZnxBigToMut<Self>,
304        A: VecZnxBigToRef<Self>,
305    {
306        vec_znx_big_automorphism(p, res, res_col, a, a_col);
307    }
308}
309
310unsafe impl VecZnxBigAutomorphismInplaceTmpBytesImpl<Self> for FFT64Avx {
311    fn vec_znx_big_automorphism_inplace_tmp_bytes_impl(module: &Module<Self>) -> usize {
312        vec_znx_big_automorphism_inplace_tmp_bytes(module.n())
313    }
314}
315
316unsafe impl VecZnxBigAutomorphismInplaceImpl<Self> for FFT64Avx
317where
318    Module<Self>: VecZnxBigAutomorphismInplaceTmpBytes,
319{
320    /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
321    fn vec_znx_big_automorphism_inplace_impl<R>(
322        module: &Module<Self>,
323        p: i64,
324        res: &mut R,
325        res_col: usize,
326        scratch: &mut Scratch<Self>,
327    ) where
328        R: VecZnxBigToMut<Self>,
329    {
330        let (tmp, _) = scratch.take_slice(module.vec_znx_big_automorphism_inplace_tmp_bytes() / size_of::<i64>());
331        vec_znx_big_automorphism_inplace(p, res, res_col, tmp);
332    }
333}