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