1use poulpy_hal::{
  2    api::{
  3        ScratchAvailable, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply,
  4        VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
  5        VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
  6    },
  7    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxZero},
  8};
  9
 10use crate::layouts::{
 11    AutomorphismKey, GGLWEInfos, GLWE, GLWEInfos, GLWESwitchingKey,
 12    prepared::{AutomorphismKeyPrepared, GLWESwitchingKeyPrepared},
 13};
 14
 15impl AutomorphismKey<Vec<u8>> {
 16    pub fn keyswitch_tmp_bytes<B: Backend, OUT, IN, KEY>(
 17        module: &Module<B>,
 18        out_infos: &OUT,
 19        in_infos: &IN,
 20        key_infos: &KEY,
 21    ) -> usize
 22    where
 23        OUT: GGLWEInfos,
 24        IN: GGLWEInfos,
 25        KEY: GGLWEInfos,
 26        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
 27    {
 28        GLWESwitchingKey::keyswitch_tmp_bytes(module, out_infos, in_infos, key_infos)
 29    }
 30
 31    pub fn keyswitch_inplace_tmp_bytes<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
 32    where
 33        OUT: GGLWEInfos,
 34        KEY: GGLWEInfos,
 35        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
 36    {
 37        GLWESwitchingKey::keyswitch_inplace_tmp_bytes(module, out_infos, key_infos)
 38    }
 39}
 40
 41impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
 42    pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
 43        &mut self,
 44        module: &Module<B>,
 45        lhs: &AutomorphismKey<DataLhs>,
 46        rhs: &GLWESwitchingKeyPrepared<DataRhs, B>,
 47        scratch: &mut Scratch<B>,
 48    ) where
 49        Module<B>: VecZnxDftBytesOf
 50            + VmpApplyDftToDftTmpBytes
 51            + VecZnxBigNormalizeTmpBytes
 52            + VmpApplyDftToDft<B>
 53            + VmpApplyDftToDftAdd<B>
 54            + VecZnxDftApply<B>
 55            + VecZnxIdftApplyConsume<B>
 56            + VecZnxBigAddSmallInplace<B>
 57            + VecZnxBigNormalize<B>
 58            + VecZnxNormalize<B>
 59            + VecZnxNormalizeTmpBytes,
 60        Scratch<B>: ScratchAvailable,
 61    {
 62        self.key.keyswitch(module, &lhs.key, rhs, scratch);
 63    }
 64
 65    pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
 66        &mut self,
 67        module: &Module<B>,
 68        rhs: &AutomorphismKeyPrepared<DataRhs, B>,
 69        scratch: &mut Scratch<B>,
 70    ) where
 71        Module<B>: VecZnxDftBytesOf
 72            + VmpApplyDftToDftTmpBytes
 73            + VecZnxBigNormalizeTmpBytes
 74            + VmpApplyDftToDft<B>
 75            + VmpApplyDftToDftAdd<B>
 76            + VecZnxDftApply<B>
 77            + VecZnxIdftApplyConsume<B>
 78            + VecZnxBigAddSmallInplace<B>
 79            + VecZnxBigNormalize<B>
 80            + VecZnxNormalize<B>
 81            + VecZnxNormalizeTmpBytes,
 82        Scratch<B>: ScratchAvailable,
 83    {
 84        self.key.keyswitch_inplace(module, &rhs.key, scratch);
 85    }
 86}
 87
 88impl GLWESwitchingKey<Vec<u8>> {
 89    pub fn keyswitch_tmp_bytes<B: Backend, OUT, IN, KEY>(
 90        module: &Module<B>,
 91        out_infos: &OUT,
 92        in_infos: &IN,
 93        key_apply: &KEY,
 94    ) -> usize
 95    where
 96        OUT: GGLWEInfos,
 97        IN: GGLWEInfos,
 98        KEY: GGLWEInfos,
 99        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
100    {
101        GLWE::keyswitch_tmp_bytes(module, out_infos, in_infos, key_apply)
102    }
103
104    pub fn keyswitch_inplace_tmp_bytes<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_apply: &KEY) -> usize
105    where
106        OUT: GGLWEInfos + GLWEInfos,
107        KEY: GGLWEInfos + GLWEInfos,
108        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
109    {
110        GLWE::keyswitch_inplace_tmp_bytes(module, out_infos, key_apply)
111    }
112}
113
114impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
115    pub fn keyswitch<DataLhs: DataRef, DataRhs: DataRef, B: Backend>(
116        &mut self,
117        module: &Module<B>,
118        lhs: &GLWESwitchingKey<DataLhs>,
119        rhs: &GLWESwitchingKeyPrepared<DataRhs, B>,
120        scratch: &mut Scratch<B>,
121    ) where
122        Module<B>: VecZnxDftBytesOf
123            + VmpApplyDftToDftTmpBytes
124            + VecZnxBigNormalizeTmpBytes
125            + VmpApplyDftToDft<B>
126            + VmpApplyDftToDftAdd<B>
127            + VecZnxDftApply<B>
128            + VecZnxIdftApplyConsume<B>
129            + VecZnxBigAddSmallInplace<B>
130            + VecZnxBigNormalize<B>
131            + VecZnxNormalize<B>
132            + VecZnxNormalizeTmpBytes,
133        Scratch<B>: ScratchAvailable,
134    {
135        #[cfg(debug_assertions)]
136        {
137            assert_eq!(
138                self.rank_in(),
139                lhs.rank_in(),
140                "ksk_out input rank: {} != ksk_in input rank: {}",
141                self.rank_in(),
142                lhs.rank_in()
143            );
144            assert_eq!(
145                lhs.rank_out(),
146                rhs.rank_in(),
147                "ksk_in output rank: {} != ksk_apply input rank: {}",
148                self.rank_out(),
149                rhs.rank_in()
150            );
151            assert_eq!(
152                self.rank_out(),
153                rhs.rank_out(),
154                "ksk_out output rank: {} != ksk_apply output rank: {}",
155                self.rank_out(),
156                rhs.rank_out()
157            );
158            assert!(
159                self.dnum() <= lhs.dnum(),
160                "self.dnum()={} > lhs.dnum()={}",
161                self.dnum(),
162                lhs.dnum()
163            );
164            assert_eq!(
165                self.dsize(),
166                lhs.dsize(),
167                "ksk_out dsize: {} != ksk_in dsize: {}",
168                self.dsize(),
169                lhs.dsize()
170            )
171        }
172
173        (0..self.rank_in().into()).for_each(|col_i| {
174            (0..self.dnum().into()).for_each(|row_j| {
175                self.at_mut(row_j, col_i)
176                    .keyswitch(module, &lhs.at(row_j, col_i), rhs, scratch);
177            });
178        });
179
180        (self.dnum().min(lhs.dnum()).into()..self.dnum().into()).for_each(|row_i| {
181            (0..self.rank_in().into()).for_each(|col_j| {
182                self.at_mut(row_i, col_j).data.zero();
183            });
184        });
185    }
186
187    pub fn keyswitch_inplace<DataRhs: DataRef, B: Backend>(
188        &mut self,
189        module: &Module<B>,
190        rhs: &GLWESwitchingKeyPrepared<DataRhs, B>,
191        scratch: &mut Scratch<B>,
192    ) where
193        Module<B>: VecZnxDftBytesOf
194            + VmpApplyDftToDftTmpBytes
195            + VecZnxBigNormalizeTmpBytes
196            + VmpApplyDftToDft<B>
197            + VmpApplyDftToDftAdd<B>
198            + VecZnxDftApply<B>
199            + VecZnxIdftApplyConsume<B>
200            + VecZnxBigAddSmallInplace<B>
201            + VecZnxBigNormalize<B>
202            + VecZnxNormalize<B>
203            + VecZnxNormalizeTmpBytes,
204        Scratch<B>: ScratchAvailable,
205    {
206        #[cfg(debug_assertions)]
207        {
208            assert_eq!(
209                self.rank_out(),
210                rhs.rank_out(),
211                "ksk_out output rank: {} != ksk_apply output rank: {}",
212                self.rank_out(),
213                rhs.rank_out()
214            );
215        }
216
217        (0..self.rank_in().into()).for_each(|col_i| {
218            (0..self.dnum().into()).for_each(|row_j| {
219                self.at_mut(row_j, col_i)
220                    .keyswitch_inplace(module, rhs, scratch)
221            });
222        });
223    }
224}