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