1use backend::hal::{
  2    api::{ScratchAvailable, TakeVecZnxDft, ZnxZero},
  3    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::{AutomorphismKey, GGSWCiphertextExec, GLWECiphertext, GLWEExternalProductFamily, GLWESwitchingKey, Infos};
  7
  8impl GLWESwitchingKey<Vec<u8>> {
  9    pub fn external_product_scratch_space<B: Backend>(
 10        module: &Module<B>,
 11        n: usize,
 12        basek: usize,
 13        k_out: usize,
 14        k_in: usize,
 15        k_ggsw: usize,
 16        digits: usize,
 17        rank: usize,
 18    ) -> usize
 19    where
 20        Module<B>: GLWEExternalProductFamily<B>,
 21    {
 22        GLWECiphertext::external_product_scratch_space(module, n, basek, k_out, k_in, k_ggsw, digits, rank)
 23    }
 24
 25    pub fn external_product_inplace_scratch_space<B: Backend>(
 26        module: &Module<B>,
 27        n: usize,
 28        basek: usize,
 29        k_out: usize,
 30        k_ggsw: usize,
 31        digits: usize,
 32        rank: usize,
 33    ) -> usize
 34    where
 35        Module<B>: GLWEExternalProductFamily<B>,
 36    {
 37        GLWECiphertext::external_product_inplace_scratch_space(module, n, basek, k_out, k_ggsw, digits, rank)
 38    }
 39}
 40
 41impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
 42    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
 43        &mut self,
 44        module: &Module<B>,
 45        lhs: &GLWESwitchingKey<DataLhs>,
 46        rhs: &GGSWCiphertextExec<DataRhs, B>,
 47        scratch: &mut Scratch<B>,
 48    ) where
 49        Module<B>: GLWEExternalProductFamily<B>,
 50        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 51    {
 52        #[cfg(debug_assertions)]
 53        {
 54            assert_eq!(
 55                self.rank_in(),
 56                lhs.rank_in(),
 57                "ksk_out input rank: {} != ksk_in input rank: {}",
 58                self.rank_in(),
 59                lhs.rank_in()
 60            );
 61            assert_eq!(
 62                lhs.rank_out(),
 63                rhs.rank(),
 64                "ksk_in output rank: {} != ggsw rank: {}",
 65                self.rank_out(),
 66                rhs.rank()
 67            );
 68            assert_eq!(
 69                self.rank_out(),
 70                rhs.rank(),
 71                "ksk_out output rank: {} != ggsw rank: {}",
 72                self.rank_out(),
 73                rhs.rank()
 74            );
 75        }
 76
 77        (0..self.rank_in()).for_each(|col_i| {
 78            (0..self.rows()).for_each(|row_j| {
 79                self.at_mut(row_j, col_i)
 80                    .external_product(module, &lhs.at(row_j, col_i), rhs, scratch);
 81            });
 82        });
 83
 84        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
 85            (0..self.rank_in()).for_each(|col_j| {
 86                self.at_mut(row_i, col_j).data.zero();
 87            });
 88        });
 89    }
 90
 91    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
 92        &mut self,
 93        module: &Module<B>,
 94        rhs: &GGSWCiphertextExec<DataRhs, B>,
 95        scratch: &mut Scratch<B>,
 96    ) where
 97        Module<B>: GLWEExternalProductFamily<B>,
 98        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 99    {
100        #[cfg(debug_assertions)]
101        {
102            assert_eq!(
103                self.rank_out(),
104                rhs.rank(),
105                "ksk_out output rank: {} != ggsw rank: {}",
106                self.rank_out(),
107                rhs.rank()
108            );
109        }
110
111        (0..self.rank_in()).for_each(|col_i| {
112            (0..self.rows()).for_each(|row_j| {
113                self.at_mut(row_j, col_i)
114                    .external_product_inplace(module, rhs, scratch);
115            });
116        });
117    }
118}
119
120impl AutomorphismKey<Vec<u8>> {
121    pub fn external_product_scratch_space<B: Backend>(
122        module: &Module<B>,
123        n: usize,
124        basek: usize,
125        k_out: usize,
126        k_in: usize,
127        ggsw_k: usize,
128        digits: usize,
129        rank: usize,
130    ) -> usize
131    where
132        Module<B>: GLWEExternalProductFamily<B>,
133    {
134        GLWESwitchingKey::external_product_scratch_space(module, n, basek, k_out, k_in, ggsw_k, digits, rank)
135    }
136
137    pub fn external_product_inplace_scratch_space<B: Backend>(
138        module: &Module<B>,
139        n: usize,
140        basek: usize,
141        k_out: usize,
142        ggsw_k: usize,
143        digits: usize,
144        rank: usize,
145    ) -> usize
146    where
147        Module<B>: GLWEExternalProductFamily<B>,
148    {
149        GLWESwitchingKey::external_product_inplace_scratch_space(module, n, basek, k_out, ggsw_k, digits, rank)
150    }
151}
152
153impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
154    pub fn external_product<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
155        &mut self,
156        module: &Module<B>,
157        lhs: &AutomorphismKey<DataLhs>,
158        rhs: &GGSWCiphertextExec<DataRhs, B>,
159        scratch: &mut Scratch<B>,
160    ) where
161        Module<B>: GLWEExternalProductFamily<B>,
162        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
163    {
164        self.key.external_product(module, &lhs.key, rhs, scratch);
165    }
166
167    pub fn external_product_inplace<DataRhs: DataRef, B: Backend>(
168        &mut self,
169        module: &Module<B>,
170        rhs: &GGSWCiphertextExec<DataRhs, B>,
171        scratch: &mut Scratch<B>,
172    ) where
173        Module<B>: GLWEExternalProductFamily<B>,
174        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
175    {
176        self.key.external_product_inplace(module, rhs, scratch);
177    }
178}