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