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