1use base2k::{
  2    FFT64, Module, Scratch, VecZnx, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef,
  3    VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxZero,
  4};
  5
  6use crate::{
  7    elem::{GetRow, Infos, SetRow},
  8    glwe_ciphertext::GLWECiphertext,
  9    glwe_ciphertext_fourier::GLWECiphertextFourier,
 10};
 11
 12pub(crate) trait VecGLWEProductScratchSpace {
 13    fn prod_with_glwe_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize;
 14
 15    fn prod_with_glwe_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
 16        Self::prod_with_glwe_scratch_space(module, res_size, res_size, rhs)
 17    }
 18
 19    fn prod_with_glwe_dft_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
 20        (Self::prod_with_glwe_scratch_space(module, res_size, lhs, rhs) | module.vec_znx_idft_tmp_bytes())
 21            + module.bytes_of_vec_znx(2, lhs)
 22            + module.bytes_of_vec_znx(2, res_size)
 23    }
 24
 25    fn prod_with_glwe_dft_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
 26        (Self::prod_with_glwe_inplace_scratch_space(module, res_size, rhs) | module.vec_znx_idft_tmp_bytes())
 27            + module.bytes_of_vec_znx(2, res_size)
 28    }
 29
 30    fn prod_with_vec_glwe_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
 31        Self::prod_with_glwe_dft_scratch_space(module, res_size, lhs, rhs)
 32            + module.bytes_of_vec_znx_dft(2, lhs)
 33            + module.bytes_of_vec_znx_dft(2, res_size)
 34    }
 35
 36    fn prod_with_vec_glwe_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
 37        Self::prod_with_glwe_dft_inplace_scratch_space(module, res_size, rhs) + module.bytes_of_vec_znx_dft(2, res_size)
 38    }
 39}
 40
 41pub(crate) trait VecGLWEProduct: Infos {
 42    fn prod_with_glwe<MUT, REF>(
 43        &self,
 44        module: &Module<FFT64>,
 45        res: &mut GLWECiphertext<MUT>,
 46        a: &GLWECiphertext<REF>,
 47        scratch: &mut Scratch,
 48    ) where
 49        VecZnx<MUT>: VecZnxToMut,
 50        VecZnx<REF>: VecZnxToRef;
 51
 52    fn prod_with_glwe_inplace<MUT>(&self, module: &Module<FFT64>, res: &mut GLWECiphertext<MUT>, scratch: &mut Scratch)
 53    where
 54        VecZnx<MUT>: VecZnxToMut + VecZnxToRef,
 55    {
 56        unsafe {
 57            let res_ptr: *mut GLWECiphertext<MUT> = res as *mut GLWECiphertext<MUT>; // This is ok because [Self::mul_rlwe] only updates res at the end.
 58            self.prod_with_glwe(&module, &mut *res_ptr, &*res_ptr, scratch);
 59        }
 60    }
 61
 62    fn prod_with_glwe_fourier<MUT, REF>(
 63        &self,
 64        module: &Module<FFT64>,
 65        res: &mut GLWECiphertextFourier<MUT, FFT64>,
 66        a: &GLWECiphertextFourier<REF, FFT64>,
 67        scratch: &mut Scratch,
 68    ) where
 69        VecZnxDft<MUT, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64> + ZnxInfos,
 70        VecZnxDft<REF, FFT64>: VecZnxDftToRef<FFT64> + ZnxInfos,
 71    {
 72        let log_base2k: usize = self.basek();
 73
 74        #[cfg(debug_assertions)]
 75        {
 76            assert_eq!(res.basek(), log_base2k);
 77            assert_eq!(self.n(), module.n());
 78            assert_eq!(res.n(), module.n());
 79        }
 80
 81        let (a_data, scratch_1) = scratch.tmp_vec_znx(module, 2, a.size());
 82
 83        let mut a_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
 84            data: a_data,
 85            basek: a.basek(),
 86            k: a.k(),
 87        };
 88
 89        a.idft(module, &mut a_idft, scratch_1);
 90
 91        let (res_data, scratch_2) = scratch_1.tmp_vec_znx(module, 2, res.size());
 92
 93        let mut res_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
 94            data: res_data,
 95            basek: res.basek(),
 96            k: res.k(),
 97        };
 98
 99        self.prod_with_glwe(module, &mut res_idft, &a_idft, scratch_2);
100
101        module.vec_znx_dft(res, 0, &res_idft, 0);
102        module.vec_znx_dft(res, 1, &res_idft, 1);
103    }
104
105    fn prod_with_glwe_fourier_inplace<MUT>(
106        &self,
107        module: &Module<FFT64>,
108        res: &mut GLWECiphertextFourier<MUT, FFT64>,
109        scratch: &mut Scratch,
110    ) where
111        VecZnxDft<MUT, FFT64>: VecZnxDftToRef<FFT64> + VecZnxDftToMut<FFT64>,
112    {
113        let log_base2k: usize = self.basek();
114
115        #[cfg(debug_assertions)]
116        {
117            assert_eq!(res.basek(), log_base2k);
118            assert_eq!(self.n(), module.n());
119            assert_eq!(res.n(), module.n());
120        }
121
122        let (res_data, scratch_1) = scratch.tmp_vec_znx(module, 2, res.size());
123
124        let mut res_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
125            data: res_data,
126            basek: res.basek(),
127            k: res.k(),
128        };
129
130        res.idft(module, &mut res_idft, scratch_1);
131
132        self.prod_with_glwe_inplace(module, &mut res_idft, scratch_1);
133
134        module.vec_znx_dft(res, 0, &res_idft, 0);
135        module.vec_znx_dft(res, 1, &res_idft, 1);
136    }
137
138    fn prod_with_vec_glwe<RES, LHS>(&self, module: &Module<FFT64>, res: &mut RES, a: &LHS, scratch: &mut Scratch)
139    where
140        LHS: GetRow<FFT64> + Infos,
141        RES: SetRow<FFT64> + Infos,
142    {
143        let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, a.size());
144
145        let mut tmp_a_row: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
146            data: tmp_row_data,
147            basek: a.basek(),
148            k: a.k(),
149        };
150
151        let (tmp_res_data, scratch2) = scratch1.tmp_vec_znx_dft(module, 2, res.size());
152
153        let mut tmp_res_row: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
154            data: tmp_res_data,
155            basek: res.basek(),
156            k: res.k(),
157        };
158
159        let min_rows: usize = res.rows().min(a.rows());
160
161        (0..res.rows()).for_each(|row_i| {
162            (0..res.cols()).for_each(|col_j| {
163                a.get_row(module, row_i, col_j, &mut tmp_a_row);
164                self.prod_with_glwe_fourier(module, &mut tmp_res_row, &tmp_a_row, scratch2);
165                res.set_row(module, row_i, col_j, &tmp_res_row);
166            });
167        });
168
169        tmp_res_row.data.zero();
170
171        (min_rows..res.rows()).for_each(|row_i| {
172            (0..self.cols()).for_each(|col_j| {
173                res.set_row(module, row_i, col_j, &tmp_res_row);
174            });
175        });
176    }
177
178    fn prod_with_vec_glwe_inplace<RES>(&self, module: &Module<FFT64>, res: &mut RES, scratch: &mut Scratch)
179    where
180        RES: GetRow<FFT64> + SetRow<FFT64> + Infos,
181    {
182        let (tmp_row_data, scratch1) = scratch.tmp_vec_znx_dft(module, 2, res.size());
183
184        let mut tmp_row: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
185            data: tmp_row_data,
186            basek: res.basek(),
187            k: res.k(),
188        };
189
190        (0..res.rows()).for_each(|row_i| {
191            (0..res.cols()).for_each(|col_j| {
192                res.get_row(module, row_i, col_j, &mut tmp_row);
193                self.prod_with_glwe_fourier_inplace(module, &mut tmp_row, scratch1);
194                res.set_row(module, row_i, col_j, &tmp_row);
195            });
196        });
197    }
198}