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