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