1use backend::hal::{
  2    api::{
  3        ScratchAvailable, TakeVecZnxDft, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx,
  4        VecZnxDftToVecZnxBigConsume, VecZnxNormalizeTmpBytes, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, ZnxZero,
  5    },
  6    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  7};
  8
  9use crate::layouts::{GGSWCiphertext, GLWECiphertext, Infos, prepared::GGSWCiphertextPrepared};
 10
 11impl GGSWCiphertext<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        GLWECiphertext::external_product_scratch_space(module, n, basek, k_out, k_in, k_ggsw, digits, rank)
 26    }
 27
 28    pub fn external_product_inplace_scratch_space<B: Backend>(
 29        module: &Module<B>,
 30        n: usize,
 31        basek: usize,
 32        k_out: usize,
 33        k_ggsw: usize,
 34        digits: usize,
 35        rank: usize,
 36    ) -> usize
 37    where
 38        Module<B>: VecZnxDftAllocBytes + VmpApplyTmpBytes + VecZnxNormalizeTmpBytes,
 39    {
 40        GLWECiphertext::external_product_inplace_scratch_space(module, n, basek, k_out, k_ggsw, digits, rank)
 41    }
 42}
 43
 44impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
 45    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
 46        &mut self,
 47        module: &Module<B>,
 48        lhs: &GGSWCiphertext<DataLhs>,
 49        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
 50        scratch: &mut Scratch<B>,
 51    ) where
 52        Module<B>: VecZnxDftAllocBytes
 53            + VmpApplyTmpBytes
 54            + VecZnxNormalizeTmpBytes
 55            + VecZnxDftFromVecZnx<B>
 56            + VmpApply<B>
 57            + VmpApplyAdd<B>
 58            + VecZnxDftToVecZnxBigConsume<B>
 59            + VecZnxBigNormalize<B>,
 60        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 61    {
 62        #[cfg(debug_assertions)]
 63        {
 64            use crate::layouts::Infos;
 65
 66            assert_eq!(lhs.n(), self.n());
 67            assert_eq!(rhs.n(), self.n());
 68
 69            assert_eq!(
 70                self.rank(),
 71                lhs.rank(),
 72                "ggsw_out rank: {} != ggsw_in rank: {}",
 73                self.rank(),
 74                lhs.rank()
 75            );
 76            assert_eq!(
 77                self.rank(),
 78                rhs.rank(),
 79                "ggsw_in rank: {} != ggsw_apply rank: {}",
 80                self.rank(),
 81                rhs.rank()
 82            );
 83
 84            assert!(
 85                scratch.available()
 86                    >= GGSWCiphertext::external_product_scratch_space(
 87                        module,
 88                        self.n(),
 89                        self.basek(),
 90                        self.k(),
 91                        lhs.k(),
 92                        rhs.k(),
 93                        rhs.digits(),
 94                        rhs.rank()
 95                    )
 96            )
 97        }
 98
 99        let min_rows: usize = self.rows().min(lhs.rows());
100
101        (0..self.rank() + 1).for_each(|col_i| {
102            (0..min_rows).for_each(|row_j| {
103                self.at_mut(row_j, col_i)
104                    .external_product(module, &lhs.at(row_j, col_i), rhs, scratch);
105            });
106            (min_rows..self.rows()).for_each(|row_i| {
107                self.at_mut(row_i, col_i).data.zero();
108            });
109        });
110    }
111
112    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
113        &mut self,
114        module: &Module<B>,
115        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
116        scratch: &mut Scratch<B>,
117    ) where
118        Module<B>: VecZnxDftAllocBytes
119            + VmpApplyTmpBytes
120            + VecZnxNormalizeTmpBytes
121            + VecZnxDftFromVecZnx<B>
122            + VmpApply<B>
123            + VmpApplyAdd<B>
124            + VecZnxDftToVecZnxBigConsume<B>
125            + VecZnxBigNormalize<B>,
126        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
127    {
128        #[cfg(debug_assertions)]
129        {
130            assert_eq!(rhs.n(), self.n());
131            assert_eq!(
132                self.rank(),
133                rhs.rank(),
134                "ggsw_out rank: {} != ggsw_apply: {}",
135                self.rank(),
136                rhs.rank()
137            );
138        }
139
140        (0..self.rank() + 1).for_each(|col_i| {
141            (0..self.rows()).for_each(|row_j| {
142                self.at_mut(row_j, col_i)
143                    .external_product_inplace(module, rhs, scratch);
144            });
145        });
146    }
147}