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