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