1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftApply,
  4        VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
  5        VmpApplyDftToDftTmpBytes,
  6    },
  7    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxZero},
  8};
  9
 10use crate::layouts::{GGSWCiphertext, GGSWInfos, GLWECiphertext, GLWEInfos, prepared::GGSWCiphertextPrepared};
 11
 12impl GGSWCiphertext<Vec<u8>> {
 13    #[allow(clippy::too_many_arguments)]
 14    pub fn external_product_scratch_space<B: Backend, OUT, IN, GGSW>(
 15        module: &Module<B>,
 16        out_infos: &OUT,
 17        in_infos: &IN,
 18        apply_infos: &GGSW,
 19    ) -> usize
 20    where
 21        OUT: GGSWInfos,
 22        IN: GGSWInfos,
 23        GGSW: GGSWInfos,
 24        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxNormalizeTmpBytes,
 25    {
 26        GLWECiphertext::external_product_scratch_space(
 27            module,
 28            &out_infos.glwe_layout(),
 29            &in_infos.glwe_layout(),
 30            apply_infos,
 31        )
 32    }
 33
 34    pub fn external_product_inplace_scratch_space<B: Backend, OUT, GGSW>(
 35        module: &Module<B>,
 36        out_infos: &OUT,
 37        apply_infos: &GGSW,
 38    ) -> usize
 39    where
 40        OUT: GGSWInfos,
 41        GGSW: GGSWInfos,
 42        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxNormalizeTmpBytes,
 43    {
 44        GLWECiphertext::external_product_inplace_scratch_space(module, &out_infos.glwe_layout(), apply_infos)
 45    }
 46}
 47
 48impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
 49    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
 50        &mut self,
 51        module: &Module<B>,
 52        lhs: &GGSWCiphertext<DataLhs>,
 53        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
 54        scratch: &mut Scratch<B>,
 55    ) where
 56        Module<B>: VecZnxDftAllocBytes
 57            + VmpApplyDftToDftTmpBytes
 58            + VecZnxNormalizeTmpBytes
 59            + VecZnxDftApply<B>
 60            + VmpApplyDftToDft<B>
 61            + VmpApplyDftToDftAdd<B>
 62            + VecZnxIdftApplyConsume<B>
 63            + VecZnxBigNormalize<B>
 64            + VecZnxNormalize<B>,
 65        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
 66    {
 67        #[cfg(debug_assertions)]
 68        {
 69            use crate::layouts::LWEInfos;
 70
 71            assert_eq!(lhs.n(), self.n());
 72            assert_eq!(rhs.n(), self.n());
 73
 74            assert_eq!(
 75                self.rank(),
 76                lhs.rank(),
 77                "ggsw_out rank: {} != ggsw_in rank: {}",
 78                self.rank(),
 79                lhs.rank()
 80            );
 81            assert_eq!(
 82                self.rank(),
 83                rhs.rank(),
 84                "ggsw_in rank: {} != ggsw_apply rank: {}",
 85                self.rank(),
 86                rhs.rank()
 87            );
 88
 89            assert!(scratch.available() >= GGSWCiphertext::external_product_scratch_space(module, self, lhs, rhs))
 90        }
 91
 92        let min_dnum: usize = self.dnum().min(lhs.dnum()).into();
 93
 94        (0..(self.rank() + 1).into()).for_each(|col_i| {
 95            (0..min_dnum).for_each(|row_j| {
 96                self.at_mut(row_j, col_i)
 97                    .external_product(module, &lhs.at(row_j, col_i), rhs, scratch);
 98            });
 99            (min_dnum..self.dnum().into()).for_each(|row_i| {
100                self.at_mut(row_i, col_i).data.zero();
101            });
102        });
103    }
104
105    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
106        &mut self,
107        module: &Module<B>,
108        rhs: &GGSWCiphertextPrepared<DataRhs, B>,
109        scratch: &mut Scratch<B>,
110    ) where
111        Module<B>: VecZnxDftAllocBytes
112            + VmpApplyDftToDftTmpBytes
113            + VecZnxNormalizeTmpBytes
114            + VecZnxDftApply<B>
115            + VmpApplyDftToDft<B>
116            + VmpApplyDftToDftAdd<B>
117            + VecZnxIdftApplyConsume<B>
118            + VecZnxBigNormalize<B>
119            + VecZnxNormalize<B>,
120        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
121    {
122        #[cfg(debug_assertions)]
123        {
124            use crate::layouts::LWEInfos;
125
126            assert_eq!(rhs.n(), self.n());
127            assert_eq!(
128                self.rank(),
129                rhs.rank(),
130                "ggsw_out rank: {} != ggsw_apply: {}",
131                self.rank(),
132                rhs.rank()
133            );
134        }
135
136        (0..(self.rank() + 1).into()).for_each(|col_i| {
137            (0..self.dnum().into()).for_each(|row_j| {
138                self.at_mut(row_j, col_i)
139                    .external_product_inplace(module, rhs, scratch);
140            });
141        });
142    }
143}