1use backend::{
  2    FFT64, MatZnxDftOps, MatZnxDftScratch, Module, Scratch, VecZnxBig, VecZnxBigOps, VecZnxDftAlloc, VecZnxDftOps, VecZnxScratch,
  3};
  4
  5use crate::{FourierGLWECiphertext, GGSWCiphertext, GLWECiphertext, Infos};
  6
  7impl GLWECiphertext<Vec<u8>> {
  8    pub fn external_product_scratch_space(
  9        module: &Module<FFT64>,
 10        basek: usize,
 11        k_out: usize,
 12        k_in: usize,
 13        k_ggsw: usize,
 14        digits: usize,
 15        rank: usize,
 16    ) -> usize {
 17        let res_dft: usize = FourierGLWECiphertext::bytes_of(module, basek, k_out, rank);
 18        let in_size: usize = k_in.div_ceil(basek).div_ceil(digits);
 19        let out_size: usize = k_out.div_ceil(basek);
 20        let ggsw_size: usize = k_ggsw.div_ceil(basek);
 21        let vmp: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size)
 22            + module.vmp_apply_tmp_bytes(
 23                out_size,
 24                in_size,
 25                in_size,  // rows
 26                rank + 1, // cols in
 27                rank + 1, // cols out
 28                ggsw_size,
 29            );
 30        let normalize: usize = module.vec_znx_normalize_tmp_bytes();
 31        res_dft + (vmp | normalize)
 32    }
 33
 34    pub fn external_product_inplace_scratch_space(
 35        module: &Module<FFT64>,
 36        basek: usize,
 37        k_out: usize,
 38        k_ggsw: usize,
 39        digits: usize,
 40        rank: usize,
 41    ) -> usize {
 42        Self::external_product_scratch_space(module, basek, k_out, k_out, k_ggsw, digits, rank)
 43    }
 44}
 45
 46impl<DataSelf: AsRef<[u8]> + AsMut<[u8]>> GLWECiphertext<DataSelf> {
 47    pub fn external_product<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
 48        &mut self,
 49        module: &Module<FFT64>,
 50        lhs: &GLWECiphertext<DataLhs>,
 51        rhs: &GGSWCiphertext<DataRhs, FFT64>,
 52        scratch: &mut Scratch,
 53    ) {
 54        let basek: usize = self.basek();
 55
 56        #[cfg(debug_assertions)]
 57        {
 58            assert_eq!(rhs.rank(), lhs.rank());
 59            assert_eq!(rhs.rank(), self.rank());
 60            assert_eq!(self.basek(), basek);
 61            assert_eq!(lhs.basek(), basek);
 62            assert_eq!(rhs.n(), module.n());
 63            assert_eq!(self.n(), module.n());
 64            assert_eq!(lhs.n(), module.n());
 65            assert!(
 66                scratch.available()
 67                    >= GLWECiphertext::external_product_scratch_space(
 68                        module,
 69                        self.basek(),
 70                        self.k(),
 71                        lhs.k(),
 72                        rhs.k(),
 73                        rhs.digits(),
 74                        rhs.rank(),
 75                    )
 76            );
 77        }
 78
 79        let cols: usize = rhs.rank() + 1;
 80        let digits: usize = rhs.digits();
 81
 82        let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size()); // Todo optimise
 83        let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, (lhs.size() + digits - 1) / digits);
 84
 85        {
 86            (0..digits).for_each(|di| {
 87                // (lhs.size() + di) / digits = (a - (digit - di - 1) + digit - 1) / digits
 88                a_dft.set_size((lhs.size() + di) / digits);
 89
 90                // Small optimization for digits > 2
 91                // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
 92                // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
 93                // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
 94                // It is possible to further ignore the last digits-1 limbs, but this introduce
 95                // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
 96                // noise is kept with respect to the ideal functionality.
 97                res_dft.set_size(rhs.size() - ((digits - di) as isize - 2).max(0) as usize);
 98
 99                (0..cols).for_each(|col_i| {
100                    module.vec_znx_dft(digits, digits - 1 - di, &mut a_dft, col_i, &lhs.data, col_i);
101                });
102
103                if di == 0 {
104                    module.vmp_apply(&mut res_dft, &a_dft, &rhs.data, scratch2);
105                } else {
106                    module.vmp_apply_add(&mut res_dft, &a_dft, &rhs.data, di, scratch2);
107                }
108            });
109        }
110
111        let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
112
113        (0..cols).for_each(|i| {
114            module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
115        });
116    }
117
118    pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
119        &mut self,
120        module: &Module<FFT64>,
121        rhs: &GGSWCiphertext<DataRhs, FFT64>,
122        scratch: &mut Scratch,
123    ) {
124        unsafe {
125            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
126            self.external_product(&module, &*self_ptr, rhs, scratch);
127        }
128    }
129}