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