1use backend::hal::{
  2    api::{ScratchAvailable, TakeVecZnxDft, ZnxZero},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::{
  7    layouts::{GLWECiphertext, GGLWESwitchingKey, Infos, prepared::GGSWCiphertextExec},
  8    trait_families::GLWEExternalProductFamily,
  9};
 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>: GLWEExternalProductFamily<B>,
 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>: GLWEExternalProductFamily<B>,
 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: &GGSWCiphertextExec<DataRhs, B>,
 50        scratch: &mut Scratch<B>,
 51    ) where
 52        Module<B>: GLWEExternalProductFamily<B>,
 53        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 54    {
 55        #[cfg(debug_assertions)]
 56        {
 57            assert_eq!(
 58                self.rank_in(),
 59                lhs.rank_in(),
 60                "ksk_out input rank: {} != ksk_in input rank: {}",
 61                self.rank_in(),
 62                lhs.rank_in()
 63            );
 64            assert_eq!(
 65                lhs.rank_out(),
 66                rhs.rank(),
 67                "ksk_in output rank: {} != ggsw rank: {}",
 68                self.rank_out(),
 69                rhs.rank()
 70            );
 71            assert_eq!(
 72                self.rank_out(),
 73                rhs.rank(),
 74                "ksk_out output rank: {} != ggsw rank: {}",
 75                self.rank_out(),
 76                rhs.rank()
 77            );
 78        }
 79
 80        (0..self.rank_in()).for_each(|col_i| {
 81            (0..self.rows()).for_each(|row_j| {
 82                self.at_mut(row_j, col_i)
 83                    .external_product(module, &lhs.at(row_j, col_i), rhs, scratch);
 84            });
 85        });
 86
 87        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
 88            (0..self.rank_in()).for_each(|col_j| {
 89                self.at_mut(row_i, col_j).data.zero();
 90            });
 91        });
 92    }
 93
 94    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
 95        &mut self,
 96        module: &Module<B>,
 97        rhs: &GGSWCiphertextExec<DataRhs, B>,
 98        scratch: &mut Scratch<B>,
 99    ) where
100        Module<B>: GLWEExternalProductFamily<B>,
101        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
102    {
103        #[cfg(debug_assertions)]
104        {
105            assert_eq!(
106                self.rank_out(),
107                rhs.rank(),
108                "ksk_out output rank: {} != ggsw rank: {}",
109                self.rank_out(),
110                rhs.rank()
111            );
112        }
113
114        (0..self.rank_in()).for_each(|col_i| {
115            (0..self.rows()).for_each(|row_j| {
116                self.at_mut(row_j, col_i)
117                    .external_product_inplace(module, rhs, scratch);
118            });
119        });
120    }
121}