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