1use poulpy_hal::{
2 api::{
3 ScratchTakeBasic, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize,
4 VecZnxBigSubSmallInplace, VecZnxBigSubSmallNegateInplace, VecZnxNormalize,
5 },
6 layouts::{Backend, DataMut, Module, Scratch, VecZnxBig},
7};
8
9use crate::{
10 GLWEKeySwitchInternal, GLWEKeyswitch, GLWENormalize, ScratchTakeCore,
11 layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWELayout, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos},
12};
13
14impl GLWE<Vec<u8>> {
15 pub fn automorphism_tmp_bytes<M, R, A, K, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
16 where
17 R: GLWEInfos,
18 A: GLWEInfos,
19 K: GGLWEInfos,
20 M: GLWEAutomorphism<BE>,
21 {
22 module.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos)
23 }
24}
25
26impl<DataSelf: DataMut> GLWE<DataSelf> {
27 pub fn automorphism<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
28 where
29 M: GLWEAutomorphism<BE>,
30 A: GLWEToRef + GLWEInfos,
31 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
32 Scratch<BE>: ScratchTakeCore<BE>,
33 {
34 module.glwe_automorphism(self, a, key, scratch);
35 }
36
37 pub fn automorphism_add<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
38 where
39 M: GLWEAutomorphism<BE>,
40 A: GLWEToRef + GLWEInfos,
41 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
42 Scratch<BE>: ScratchTakeCore<BE>,
43 {
44 module.glwe_automorphism_add(self, a, key, scratch);
45 }
46
47 pub fn automorphism_sub<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
48 where
49 M: GLWEAutomorphism<BE>,
50 A: GLWEToRef + GLWEInfos,
51 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
52 Scratch<BE>: ScratchTakeCore<BE>,
53 {
54 module.glwe_automorphism_sub(self, a, key, scratch);
55 }
56
57 pub fn automorphism_sub_negate<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
58 where
59 M: GLWEAutomorphism<BE>,
60 A: GLWEToRef + GLWEInfos,
61 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
62 Scratch<BE>: ScratchTakeCore<BE>,
63 {
64 module.glwe_automorphism_sub_negate(self, a, key, scratch);
65 }
66
67 pub fn automorphism_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
68 where
69 M: GLWEAutomorphism<BE>,
70 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
71 Scratch<BE>: ScratchTakeCore<BE>,
72 {
73 module.glwe_automorphism_inplace(self, key, scratch);
74 }
75
76 pub fn automorphism_add_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
77 where
78 M: GLWEAutomorphism<BE>,
79 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
80 Scratch<BE>: ScratchTakeCore<BE>,
81 {
82 module.glwe_automorphism_add_inplace(self, key, scratch);
83 }
84
85 pub fn automorphism_sub_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
86 where
87 M: GLWEAutomorphism<BE>,
88 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
89 Scratch<BE>: ScratchTakeCore<BE>,
90 {
91 module.glwe_automorphism_sub_inplace(self, key, scratch);
92 }
93
94 pub fn automorphism_sub_negate_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
95 where
96 M: GLWEAutomorphism<BE>,
97 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
98 Scratch<BE>: ScratchTakeCore<BE>,
99 {
100 module.glwe_automorphism_sub_negate_inplace(self, key, scratch);
101 }
102}
103
104pub trait GLWEAutomorphism<BE: Backend> {
105 fn glwe_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
106 where
107 R: GLWEInfos,
108 A: GLWEInfos,
109 K: GGLWEInfos;
110
111 fn glwe_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
112 where
113 R: GLWEToMut + GLWEInfos,
114 A: GLWEToRef + GLWEInfos,
115 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
116
117 fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
118 where
119 R: GLWEToMut + GLWEInfos,
120 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
121
122 fn glwe_automorphism_add<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
123 where
124 R: GLWEToMut + GLWEInfos,
125 A: GLWEToRef + GLWEInfos,
126 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
127
128 fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
129 where
130 R: GLWEToMut + GLWEInfos,
131 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
132
133 fn glwe_automorphism_sub<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
134 where
135 R: GLWEToMut + GLWEInfos,
136 A: GLWEToRef + GLWEInfos,
137 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
138
139 fn glwe_automorphism_sub_negate<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
140 where
141 R: GLWEToMut + GLWEInfos,
142 A: GLWEToRef + GLWEInfos,
143 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
144
145 fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
146 where
147 R: GLWEToMut + GLWEInfos,
148 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
149
150 fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
151 where
152 R: GLWEToMut + GLWEInfos,
153 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
154}
155
156impl<BE: Backend> GLWEAutomorphism<BE> for Module<BE>
157where
158 Self: Sized
159 + GLWEKeyswitch<BE>
160 + GLWEKeySwitchInternal<BE>
161 + VecZnxNormalize<BE>
162 + VecZnxAutomorphismInplace<BE>
163 + VecZnxBigAutomorphismInplace<BE>
164 + VecZnxBigSubSmallInplace<BE>
165 + VecZnxBigSubSmallNegateInplace<BE>
166 + VecZnxBigAddSmallInplace<BE>
167 + VecZnxBigNormalize<BE>
168 + GLWENormalize<BE>,
169 Scratch<BE>: ScratchTakeCore<BE>,
170{
171 fn glwe_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
172 where
173 R: GLWEInfos,
174 A: GLWEInfos,
175 K: GGLWEInfos,
176 {
177 self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
178 }
179
180 fn glwe_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
181 where
182 R: GLWEToMut + GLWEInfos,
183 A: GLWEToRef + GLWEInfos,
184 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
185 Scratch<BE>: ScratchTakeCore<BE>,
186 {
187 self.glwe_keyswitch(res, a, key, scratch);
188
189 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
190
191 for i in 0..res.rank().as_usize() + 1 {
192 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
193 }
194 }
195
196 fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
197 where
198 R: GLWEToMut + GLWEInfos,
199 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
200 Scratch<BE>: ScratchTakeCore<BE>,
201 {
202 self.glwe_keyswitch_inplace(res, key, scratch);
203
204 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
205
206 for i in 0..res.rank().as_usize() + 1 {
207 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
208 }
209 }
210
211 fn glwe_automorphism_add<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
212 where
213 R: GLWEToMut + GLWEInfos,
214 A: GLWEToRef + GLWEInfos,
215 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
216 Scratch<BE>: ScratchTakeCore<BE>,
217 {
218 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
219 let a: &GLWE<&[u8]> = &a.to_ref();
220
221 let base2k_a: usize = a.base2k().into();
222 let base2k_key: usize = key.base2k().into();
223 let base2k_res: usize = res.base2k().into();
224
225 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); if base2k_a != base2k_key {
228 let (mut a_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
229 n: a.n(),
230 base2k: key.base2k(),
231 k: a.k(),
232 rank: a.rank(),
233 });
234 self.glwe_normalize(&mut a_conv, a, scratch_2);
235 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &a_conv, key, scratch_2);
236 for i in 0..res.rank().as_usize() + 1 {
237 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
238 self.vec_znx_big_add_small_inplace(&mut res_big, i, a_conv.data(), i);
239 self.vec_znx_big_normalize(
240 base2k_res,
241 res.data_mut(),
242 i,
243 base2k_key,
244 &res_big,
245 i,
246 scratch_2,
247 );
248 }
249 } else {
250 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
251 for i in 0..res.rank().as_usize() + 1 {
252 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
253 self.vec_znx_big_add_small_inplace(&mut res_big, i, a.data(), i);
254 self.vec_znx_big_normalize(
255 base2k_res,
256 res.data_mut(),
257 i,
258 base2k_key,
259 &res_big,
260 i,
261 scratch_1,
262 );
263 }
264 };
265 }
266
267 fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
268 where
269 R: GLWEToMut,
270 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
271 Scratch<BE>: ScratchTakeCore<BE>,
272 {
273 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
274
275 let base2k_key: usize = key.base2k().into();
276 let base2k_res: usize = res.base2k().into();
277
278 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); if base2k_res != base2k_key {
281 let (mut res_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
282 n: res.n(),
283 base2k: key.base2k(),
284 k: res.k(),
285 rank: res.rank(),
286 });
287 self.glwe_normalize(&mut res_conv, res, scratch_2);
288 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &res_conv, key, scratch_2);
289 for i in 0..res.rank().as_usize() + 1 {
290 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
291 self.vec_znx_big_add_small_inplace(&mut res_big, i, res_conv.data(), i);
292 self.vec_znx_big_normalize(
293 base2k_res,
294 res.data_mut(),
295 i,
296 base2k_key,
297 &res_big,
298 i,
299 scratch_2,
300 );
301 }
302 } else {
303 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
304 for i in 0..res.rank().as_usize() + 1 {
305 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
306 self.vec_znx_big_add_small_inplace(&mut res_big, i, res.data(), i);
307 self.vec_znx_big_normalize(
308 base2k_res,
309 res.data_mut(),
310 i,
311 base2k_key,
312 &res_big,
313 i,
314 scratch_1,
315 );
316 }
317 };
318 }
319
320 fn glwe_automorphism_sub<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
321 where
322 R: GLWEToMut,
323 A: GLWEToRef,
324 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
325 Scratch<BE>: ScratchTakeCore<BE>,
326 {
327 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
328 let a: &GLWE<&[u8]> = &a.to_ref();
329
330 let base2k_a: usize = a.base2k().into();
331 let base2k_key: usize = key.base2k().into();
332 let base2k_res: usize = res.base2k().into();
333
334 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); if base2k_a != base2k_key {
337 let (mut a_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
338 n: a.n(),
339 base2k: key.base2k(),
340 k: a.k(),
341 rank: a.rank(),
342 });
343 self.glwe_normalize(&mut a_conv, a, scratch_2);
344 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &a_conv, key, scratch_2);
345 for i in 0..res.rank().as_usize() + 1 {
346 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
347 self.vec_znx_big_sub_small_inplace(&mut res_big, i, a_conv.data(), i);
348 self.vec_znx_big_normalize(
349 base2k_res,
350 res.data_mut(),
351 i,
352 base2k_key,
353 &res_big,
354 i,
355 scratch_2,
356 );
357 }
358 } else {
359 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
360 for i in 0..res.rank().as_usize() + 1 {
361 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
362 self.vec_znx_big_sub_small_inplace(&mut res_big, i, a.data(), i);
363 self.vec_znx_big_normalize(
364 base2k_res,
365 res.data_mut(),
366 i,
367 base2k_key,
368 &res_big,
369 i,
370 scratch_1,
371 );
372 }
373 };
374 }
375
376 fn glwe_automorphism_sub_negate<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
377 where
378 R: GLWEToMut,
379 A: GLWEToRef,
380 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
381 Scratch<BE>: ScratchTakeCore<BE>,
382 {
383 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
384 let a: &GLWE<&[u8]> = &a.to_ref();
385
386 let base2k_a: usize = a.base2k().into();
387 let base2k_key: usize = key.base2k().into();
388 let base2k_res: usize = res.base2k().into();
389
390 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); if base2k_a != base2k_key {
393 let (mut a_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
394 n: a.n(),
395 base2k: key.base2k(),
396 k: a.k(),
397 rank: a.rank(),
398 });
399 self.glwe_normalize(&mut a_conv, a, scratch_2);
400 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &a_conv, key, scratch_2);
401 for i in 0..res.rank().as_usize() + 1 {
402 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
403 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, a_conv.data(), i);
404 self.vec_znx_big_normalize(
405 base2k_res,
406 res.data_mut(),
407 i,
408 base2k_key,
409 &res_big,
410 i,
411 scratch_2,
412 );
413 }
414 } else {
415 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
416 for i in 0..res.rank().as_usize() + 1 {
417 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
418 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, a.data(), i);
419 self.vec_znx_big_normalize(
420 base2k_res,
421 res.data_mut(),
422 i,
423 base2k_key,
424 &res_big,
425 i,
426 scratch_1,
427 );
428 }
429 };
430 }
431
432 fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
433 where
434 R: GLWEToMut,
435 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
436 Scratch<BE>: ScratchTakeCore<BE>,
437 {
438 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
439
440 let base2k_key: usize = key.base2k().into();
441 let base2k_res: usize = res.base2k().into();
442
443 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); if base2k_res != base2k_key {
446 let (mut res_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
447 n: res.n(),
448 base2k: key.base2k(),
449 k: res.k(),
450 rank: res.rank(),
451 });
452 self.glwe_normalize(&mut res_conv, res, scratch_2);
453 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &res_conv, key, scratch_2);
454 for i in 0..res.rank().as_usize() + 1 {
455 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
456 self.vec_znx_big_sub_small_inplace(&mut res_big, i, res_conv.data(), i);
457 self.vec_znx_big_normalize(
458 base2k_res,
459 res.data_mut(),
460 i,
461 base2k_key,
462 &res_big,
463 i,
464 scratch_2,
465 );
466 }
467 } else {
468 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
469 for i in 0..res.rank().as_usize() + 1 {
470 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
471 self.vec_znx_big_sub_small_inplace(&mut res_big, i, res.data(), i);
472 self.vec_znx_big_normalize(
473 base2k_res,
474 res.data_mut(),
475 i,
476 base2k_key,
477 &res_big,
478 i,
479 scratch_1,
480 );
481 }
482 };
483 }
484
485 fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
486 where
487 R: GLWEToMut,
488 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
489 Scratch<BE>: ScratchTakeCore<BE>,
490 {
491 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
492
493 let base2k_key: usize = key.base2k().into();
494 let base2k_res: usize = res.base2k().into();
495
496 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); if base2k_res != base2k_key {
499 let (mut res_conv, scratch_2) = scratch_1.take_glwe(&GLWELayout {
500 n: res.n(),
501 base2k: key.base2k(),
502 k: res.k(),
503 rank: res.rank(),
504 });
505 self.glwe_normalize(&mut res_conv, res, scratch_2);
506 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, &res_conv, key, scratch_2);
507 for i in 0..res.rank().as_usize() + 1 {
508 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_2);
509 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, res_conv.data(), i);
510 self.vec_znx_big_normalize(
511 base2k_res,
512 res.data_mut(),
513 i,
514 base2k_key,
515 &res_big,
516 i,
517 scratch_2,
518 );
519 }
520 } else {
521 let mut res_big: VecZnxBig<&mut [u8], BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
522 for i in 0..res.rank().as_usize() + 1 {
523 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
524 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, res.data(), i);
525 self.vec_znx_big_normalize(
526 base2k_res,
527 res.data_mut(),
528 i,
529 base2k_key,
530 &res_big,
531 i,
532 scratch_1,
533 );
534 }
535 };
536 }
537}