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