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