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