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