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