poulpy/core/src/automorphism.rs
core/src/automorphism.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDftOps, ScalarZnxOps,
3 ScalarZnxToRef, Scratch, VecZnx, VecZnxDft, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps, ZnxZero,
4};
5use sampling::source::Source;
6
7use crate::{
8 elem::{GetRow, Infos, SetRow},
9 gglwe_ciphertext::GGLWECiphertext,
10 ggsw_ciphertext::GGSWCiphertext,
11 glwe_ciphertext::GLWECiphertext,
12 glwe_ciphertext_fourier::GLWECiphertextFourier,
13 keys::{SecretKey, SecretKeyFourier},
14 keyswitch_key::GLWESwitchingKey,
15};
16
17pub struct AutomorphismKey<Data, B: Backend> {
18 pub(crate) key: GLWESwitchingKey<Data, B>,
19 pub(crate) p: i64,
20}
21
22impl AutomorphismKey<Vec<u8>, FFT64> {
23 pub fn new(module: &Module<FFT64>, basek: usize, p: i64, k: usize, rows: usize, rank: usize) -> Self {
24 AutomorphismKey {
25 key: GLWESwitchingKey::new(module, basek, k, rows, rank, rank),
26 p: p,
27 }
28 }
29}
30
31impl<T, B: Backend> Infos for AutomorphismKey<T, B> {
32 type Inner = MatZnxDft<T, B>;
33
34 fn inner(&self) -> &Self::Inner {
35 &self.key.inner()
36 }
37
38 fn basek(&self) -> usize {
39 self.key.basek()
40 }
41
42 fn k(&self) -> usize {
43 self.key.k()
44 }
45}
46
47impl<T, B: Backend> AutomorphismKey<T, B> {
48 pub fn p(&self) -> i64 {
49 self.p
50 }
51
52 pub fn rank(&self) -> usize {
53 self.key.rank()
54 }
55
56 pub fn rank_in(&self) -> usize {
57 self.key.rank_in()
58 }
59
60 pub fn rank_out(&self) -> usize {
61 self.key.rank_out()
62 }
63}
64
65impl<DataSelf, B: Backend> MatZnxDftToMut<B> for AutomorphismKey<DataSelf, B>
66where
67 MatZnxDft<DataSelf, B>: MatZnxDftToMut<B>,
68{
69 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
70 self.key.to_mut()
71 }
72}
73
74impl<DataSelf, B: Backend> MatZnxDftToRef<B> for AutomorphismKey<DataSelf, B>
75where
76 MatZnxDft<DataSelf, B>: MatZnxDftToRef<B>,
77{
78 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
79 self.key.to_ref()
80 }
81}
82
83impl<C> GetRow<FFT64> for AutomorphismKey<C, FFT64>
84where
85 MatZnxDft<C, FFT64>: MatZnxDftToRef<FFT64>,
86{
87 fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut GLWECiphertextFourier<R, FFT64>)
88 where
89 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64>,
90 {
91 module.vmp_extract_row(res, self, row_i, col_j);
92 }
93}
94
95impl<C> SetRow<FFT64> for AutomorphismKey<C, FFT64>
96where
97 MatZnxDft<C, FFT64>: MatZnxDftToMut<FFT64>,
98{
99 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &GLWECiphertextFourier<R, FFT64>)
100 where
101 VecZnxDft<R, FFT64>: VecZnxDftToRef<FFT64>,
102 {
103 module.vmp_prepare_row(self, row_i, col_j, a);
104 }
105}
106
107impl AutomorphismKey<Vec<u8>, FFT64> {
108 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
109 GGLWECiphertext::encrypt_sk_scratch_space(module, rank, size)
110 }
111
112 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, rank: usize, pk_size: usize) -> usize {
113 GGLWECiphertext::encrypt_pk_scratch_space(module, rank, pk_size)
114 }
115
116 pub fn keyswitch_scratch_space(
117 module: &Module<FFT64>,
118 out_size: usize,
119 in_size: usize,
120 ksk_size: usize,
121 rank: usize,
122 ) -> usize {
123 GLWESwitchingKey::keyswitch_scratch_space(module, out_size, rank, in_size, rank, ksk_size)
124 }
125
126 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
127 GLWESwitchingKey::keyswitch_inplace_scratch_space(module, out_size, out_rank, ksk_size)
128 }
129
130 pub fn external_product_scratch_space(
131 module: &Module<FFT64>,
132 out_size: usize,
133 in_size: usize,
134 ggsw_size: usize,
135 rank: usize,
136 ) -> usize {
137 GLWESwitchingKey::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank)
138 }
139
140 pub fn external_product_inplace_scratch_space(
141 module: &Module<FFT64>,
142 out_size: usize,
143 ggsw_size: usize,
144 rank: usize,
145 ) -> usize {
146 GLWESwitchingKey::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank)
147 }
148}
149
150impl<DataSelf> AutomorphismKey<DataSelf, FFT64>
151where
152 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
153{
154 pub fn encrypt_sk<DataSk>(
155 &mut self,
156 module: &Module<FFT64>,
157 p: i64,
158 sk: &SecretKey<DataSk>,
159 source_xa: &mut Source,
160 source_xe: &mut Source,
161 sigma: f64,
162 scratch: &mut Scratch,
163 ) where
164 ScalarZnx<DataSk>: ScalarZnxToRef,
165 {
166 #[cfg(debug_assertions)]
167 {
168 assert_eq!(self.n(), module.n());
169 assert_eq!(sk.n(), module.n());
170 assert_eq!(self.rank_out(), self.rank_in());
171 assert_eq!(sk.rank(), self.rank());
172 }
173
174 let (sk_out_dft_data, scratch_1) = scratch.tmp_scalar_znx_dft(module, sk.rank());
175
176 let mut sk_out_dft: SecretKeyFourier<&mut [u8], FFT64> = SecretKeyFourier {
177 data: sk_out_dft_data,
178 dist: sk.dist,
179 };
180
181 {
182 (0..self.rank()).for_each(|i| {
183 let (mut sk_inv_auto, _) = scratch_1.tmp_scalar_znx(module, 1);
184 module.scalar_znx_automorphism(module.galois_element_inv(p), &mut sk_inv_auto, 0, sk, i);
185 module.svp_prepare(&mut sk_out_dft, i, &sk_inv_auto, 0);
186 });
187 }
188
189 self.key.encrypt_sk(
190 module,
191 &sk,
192 &sk_out_dft,
193 source_xa,
194 source_xe,
195 sigma,
196 scratch_1,
197 );
198
199 self.p = p;
200 }
201}
202
203impl<DataSelf> AutomorphismKey<DataSelf, FFT64>
204where
205 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
206{
207 pub fn automorphism<DataLhs, DataRhs>(
208 &mut self,
209 module: &Module<FFT64>,
210 lhs: &AutomorphismKey<DataLhs, FFT64>,
211 rhs: &AutomorphismKey<DataRhs, FFT64>,
212 scratch: &mut base2k::Scratch,
213 ) where
214 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
215 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
216 {
217 #[cfg(debug_assertions)]
218 {
219 assert_eq!(
220 self.rank_in(),
221 lhs.rank_in(),
222 "ksk_out input rank: {} != ksk_in input rank: {}",
223 self.rank_in(),
224 lhs.rank_in()
225 );
226 assert_eq!(
227 lhs.rank_out(),
228 rhs.rank_in(),
229 "ksk_in output rank: {} != ksk_apply input rank: {}",
230 self.rank_out(),
231 rhs.rank_in()
232 );
233 assert_eq!(
234 self.rank_out(),
235 rhs.rank_out(),
236 "ksk_out output rank: {} != ksk_apply output rank: {}",
237 self.rank_out(),
238 rhs.rank_out()
239 );
240 }
241
242 let cols_out: usize = rhs.rank_out() + 1;
243
244 let (tmp_dft_data, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, lhs.size());
245
246 let mut tmp_dft: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
247 data: tmp_dft_data,
248 basek: lhs.basek(),
249 k: lhs.k(),
250 };
251
252 (0..self.rank_in()).for_each(|col_i| {
253 (0..self.rows()).for_each(|row_j| {
254 // Extracts relevant row
255 lhs.get_row(module, row_j, col_i, &mut tmp_dft);
256
257 // Get a VecZnxBig from scratch space
258 let (mut tmp_idft_data, scratch2) = scratch1.tmp_vec_znx_big(module, cols_out, self.size());
259
260 // Switches input outside of DFT
261 (0..cols_out).for_each(|i| {
262 module.vec_znx_idft(&mut tmp_idft_data, i, &tmp_dft.data, i, scratch2);
263 });
264
265 // Consumes to small vec znx
266 let mut tmp_idft_small_data: VecZnx<&mut [u8]> = tmp_idft_data.to_vec_znx_small();
267
268 // Reverts the automorphis key from (-pi^{-1}_{k}(s)a + s, a) to (-sa + pi_{k}(s), a)
269 (0..cols_out).for_each(|i| {
270 module.vec_znx_automorphism_inplace(self.p(), &mut tmp_idft_small_data, i);
271 });
272
273 // Wraps into ciphertext
274 let mut tmp_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
275 data: tmp_idft_small_data,
276 basek: self.basek(),
277 k: self.k(),
278 };
279
280 // Key-switch (-sa + pi_{k}(s), a) to (-pi^{-1}_{k'}(s)a + pi_{k}(s), a)
281 tmp_idft.keyswitch_inplace(module, &rhs.key, scratch2);
282
283 // Applies back the automorphism X^{k}: (-pi^{-1}_{k'}(s)a + pi_{k}(s), a) -> (-pi^{-1}_{k'+k}(s)a + s, a)
284 // and switches back to DFT domain
285 (0..self.rank_out() + 1).for_each(|i| {
286 module.vec_znx_automorphism_inplace(rhs.p(), &mut tmp_idft, i);
287 module.vec_znx_dft(&mut tmp_dft, i, &tmp_idft, i);
288 });
289
290 // Sets back the relevant row
291 self.set_row(module, row_j, col_i, &tmp_dft);
292 });
293 });
294
295 tmp_dft.data.zero();
296
297 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
298 (0..self.rank_in()).for_each(|col_j| {
299 self.set_row(module, row_i, col_j, &tmp_dft);
300 });
301 });
302 }
303
304 pub fn keyswitch<DataLhs, DataRhs>(
305 &mut self,
306 module: &Module<FFT64>,
307 lhs: &AutomorphismKey<DataLhs, FFT64>,
308 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
309 scratch: &mut base2k::Scratch,
310 ) where
311 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
312 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
313 {
314 self.key.keyswitch(module, &lhs.key, rhs, scratch);
315 }
316
317 pub fn keyswitch_inplace<DataRhs>(
318 &mut self,
319 module: &Module<FFT64>,
320 rhs: &AutomorphismKey<DataRhs, FFT64>,
321 scratch: &mut base2k::Scratch,
322 ) where
323 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
324 {
325 self.key.keyswitch_inplace(module, &rhs.key, scratch);
326 }
327
328 pub fn external_product<DataLhs, DataRhs>(
329 &mut self,
330 module: &Module<FFT64>,
331 lhs: &AutomorphismKey<DataLhs, FFT64>,
332 rhs: &GGSWCiphertext<DataRhs, FFT64>,
333 scratch: &mut Scratch,
334 ) where
335 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
336 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
337 {
338 self.key.external_product(module, &lhs.key, rhs, scratch);
339 }
340
341 pub fn external_product_inplace<DataRhs>(
342 &mut self,
343 module: &Module<FFT64>,
344 rhs: &GGSWCiphertext<DataRhs, FFT64>,
345 scratch: &mut Scratch,
346 ) where
347 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
348 {
349 self.key.external_product_inplace(module, rhs, scratch);
350 }
351}