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