1use poulpy_hal::{
  2    api::{
  3        DFT, IDFTConsume, ScratchAvailable, TakeVecZnxDft, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxNormalizeTmpBytes,
  4        VmpApply, VmpApplyAdd, VmpApplyTmpBytes,
  5    },
  6    layouts::{Backend, DataMut, DataRef, DataViewMut, Module, Scratch, VecZnxBig},
  7};
  8
  9use crate::layouts::{GLWECiphertext, Infos, prepared::GGSWCiphertextPrepared};
 10
 11impl GLWECiphertext<Vec<u8>> {
 12    #[allow(clippy::too_many_arguments)]
 13    pub fn external_product_scratch_space<B: Backend>(
 14        module: &Module<B>,
 15        basek: usize,
 16        k_out: usize,
 17        k_in: usize,
 18        k_ggsw: usize,
 19        digits: usize,
 20        rank: usize,
 21    ) -> usize
 22    where
 23        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxNormalizeTmpBytes,
 24    {
 25        let in_size: usize = k_in.div_ceil(basek).div_ceil(digits);
 26        let out_size: usize = k_out.div_ceil(basek);
 27        let ggsw_size: usize = k_ggsw.div_ceil(basek);
 28        let res_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, ggsw_size);
 29        let a_dft: usize = module.vec_znx_dft_alloc_bytes(rank + 1, in_size);
 30        let vmp: usize = module.vmp_apply_tmp_bytes(
 31            out_size,
 32            in_size,
 33            in_size,  // rows
 34            rank + 1, // cols in
 35            rank + 1, // cols out
 36            ggsw_size,
 37        );
 38        let normalize: usize = module.vec_znx_normalize_tmp_bytes();
 39        res_dft + a_dft + (vmp | normalize)
 40    }
 41
 42    pub fn external_product_inplace_scratch_space<B: Backend>(
 43        module: &Module<B>,
 44        basek: usize,
 45        k_out: usize,
 46        k_ggsw: usize,
 47        digits: usize,
 48        rank: usize,
 49    ) -> usize
 50    where
 51        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxNormalizeTmpBytes,
 52    {
 53        Self::external_product_scratch_space(module, basek, k_out, k_out, k_ggsw, digits, rank)
 54    }
 55}
 56
 57impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
 58    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
 59        &mut self,
 60        module: &Module<B>,
 61        lhs: &GLWECiphertext<DataLhs>,
 62        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
 63        scratch: &mut Scratch<B>,
 64    ) where
 65        Module<B>: VecZnxDftAllocBytes
 66            + VmpApplyTmpBytes
 67            + VecZnxNormalizeTmpBytes
 68            + DFT<B>
 69            + VmpApply<B>
 70            + VmpApplyAdd<B>
 71            + IDFTConsume<B>
 72            + VecZnxBigNormalize<B>,
 73        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 74    {
 75        let basek: usize = self.basek();
 76
 77        #[cfg(debug_assertions)]
 78        {
 79            use poulpy_hal::api::ScratchAvailable;
 80
 81            assert_eq!(rhs.rank(), lhs.rank());
 82            assert_eq!(rhs.rank(), self.rank());
 83            assert_eq!(self.basek(), basek);
 84            assert_eq!(lhs.basek(), basek);
 85            assert_eq!(rhs.n(), self.n());
 86            assert_eq!(lhs.n(), self.n());
 87            assert!(
 88                scratch.available()
 89                    >= GLWECiphertext::external_product_scratch_space(
 90                        module,
 91                        self.basek(),
 92                        self.k(),
 93                        lhs.k(),
 94                        rhs.k(),
 95                        rhs.digits(),
 96                        rhs.rank(),
 97                    )
 98            );
 99        }
100
101        let cols: usize = rhs.rank() + 1;
102        let digits: usize = rhs.digits();
103
104        let (mut res_dft, scratch1) = scratch.take_vec_znx_dft(self.n(), cols, rhs.size()); // Todo optimise
105        let (mut a_dft, scratch2) = scratch1.take_vec_znx_dft(self.n(), cols, lhs.size().div_ceil(digits));
106
107        a_dft.data_mut().fill(0);
108
109        {
110            (0..digits).for_each(|di| {
111                // (lhs.size() + di) / digits = (a - (digit - di - 1)).div_ceil(digits)
112                a_dft.set_size((lhs.size() + di) / digits);
113
114                // Small optimization for digits > 2
115                // VMP produce some error e, and since we aggregate vmp * 2^{di * B}, then
116                // we also aggregate ei * 2^{di * B}, with the largest error being ei * 2^{(digits-1) * B}.
117                // As such we can ignore the last digits-2 limbs safely of the sum of vmp products.
118                // It is possible to further ignore the last digits-1 limbs, but this introduce
119                // ~0.5 to 1 bit of additional noise, and thus not chosen here to ensure that the same
120                // noise is kept with respect to the ideal functionality.
121                res_dft.set_size(rhs.size() - ((digits - di) as isize - 2).max(0) as usize);
122
123                (0..cols).for_each(|col_i| {
124                    module.dft(digits, digits - 1 - di, &mut a_dft, col_i, &lhs.data, col_i);
125                });
126
127                if di == 0 {
128                    module.vmp_apply(&mut res_dft, &a_dft, &rhs.data, scratch2);
129                } else {
130                    module.vmp_apply_add(&mut res_dft, &a_dft, &rhs.data, di, scratch2);
131                }
132            });
133        }
134
135        let res_big: VecZnxBig<&mut [u8], B> = module.vec_znx_idft_consume(res_dft);
136
137        (0..cols).for_each(|i| {
138            module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
139        });
140    }
141
142    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
143        &mut self,
144        module: &Module<B>,
145        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
146        scratch: &mut Scratch<B>,
147    ) where
148        Module<B>: VecZnxDftAllocBytes
149            + VmpApplyTmpBytes
150            + VecZnxNormalizeTmpBytes
151            + DFT<B>
152            + VmpApply<B>
153            + VmpApplyAdd<B>
154            + IDFTConsume<B>
155            + VecZnxBigNormalize<B>,
156        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
157    {
158        unsafe {
159            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
160            self.external_product(module, &*self_ptr, rhs, scratch);
161        }
162    }
163}