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