poulpy/core/src/gglwe/test/gglwe_generic.rs
core/src/gglwe/test/gglwe_generic.rsRaw
1use backend::hal::{
2 api::{
3 MatZnxAlloc, ScalarZnxAlloc, ScalarZnxAllocBytes, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace,
4 VecZnxAlloc, VecZnxAllocBytes, VecZnxRotateInplace, VecZnxStd, VecZnxSubScalarInplace, VecZnxSwithcDegree, ZnxViewMut,
5 },
6 layouts::{Backend, Module, ScalarZnx, ScalarZnxToMut, ScratchOwned},
7 oep::{
8 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
9 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
10 },
11};
12use sampling::source::Source;
13
14use crate::{
15 GGLWEEncryptSkFamily, GGLWEExecLayoutFamily, GGSWCiphertext, GGSWCiphertextExec, GGSWLayoutFamily, GLWEDecryptFamily,
16 GLWEExternalProductFamily, GLWEKeyswitchFamily, GLWESecret, GLWESecretExec, GLWESwitchingKey,
17 GLWESwitchingKeyEncryptSkFamily, GLWESwitchingKeyExec,
18 noise::{log2_std_noise_gglwe_product, noise_ggsw_product},
19};
20
21pub(crate) trait TestModuleFamily<B: Backend> = GGLWEEncryptSkFamily<B>
22 + GLWEDecryptFamily<B>
23 + MatZnxAlloc
24 + ScalarZnxAlloc
25 + ScalarZnxAllocBytes
26 + VecZnxAllocBytes
27 + VecZnxSwithcDegree
28 + VecZnxAddScalarInplace
29 + VecZnxStd
30 + VecZnxAlloc
31 + VecZnxSubScalarInplace;
32
33pub(crate) trait TestScratchFamily<B: Backend> = TakeVecZnxDftImpl<B>
34 + TakeVecZnxBigImpl<B>
35 + TakeSvpPPolImpl<B>
36 + ScratchOwnedAllocImpl<B>
37 + ScratchOwnedBorrowImpl<B>
38 + ScratchAvailableImpl<B>
39 + TakeScalarZnxImpl<B>
40 + TakeVecZnxImpl<B>
41 + VecZnxDftAllocBytesImpl<B>
42 + VecZnxBigAllocBytesImpl<B>
43 + TakeSvpPPolImpl<B>;
44
45pub(crate) fn test_encrypt_sk<B: Backend>(
46 module: &Module<B>,
47 basek: usize,
48 k_ksk: usize,
49 digits: usize,
50 rank_in: usize,
51 rank_out: usize,
52 sigma: f64,
53) where
54 Module<B>: TestModuleFamily<B>,
55 B: TestScratchFamily<B>,
56{
57 let rows: usize = (k_ksk - digits * basek) / (digits * basek);
58
59 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc(module, basek, k_ksk, rows, digits, rank_in, rank_out);
60
61 let mut source_xs: Source = Source::new([0u8; 32]);
62 let mut source_xe: Source = Source::new([0u8; 32]);
63 let mut source_xa: Source = Source::new([0u8; 32]);
64
65 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GLWESwitchingKey::encrypt_sk_scratch_space(
66 module, basek, k_ksk, rank_in, rank_out,
67 ));
68
69 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_in);
70 sk_in.fill_ternary_prob(0.5, &mut source_xs);
71
72 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_out);
73 sk_out.fill_ternary_prob(0.5, &mut source_xs);
74 let sk_out_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_out);
75
76 ksk.encrypt_sk(
77 module,
78 &sk_in,
79 &sk_out,
80 &mut source_xa,
81 &mut source_xe,
82 sigma,
83 scratch.borrow(),
84 );
85
86 ksk.key
87 .assert_noise(module, &sk_out_exec, &sk_in.data, sigma);
88}
89
90pub(crate) fn test_keyswitch<B: Backend>(
91 module: &Module<B>,
92 basek: usize,
93 k_out: usize,
94 k_in: usize,
95 k_ksk: usize,
96 digits: usize,
97 rank_in_s0s1: usize,
98 rank_out_s0s1: usize,
99 rank_out_s1s2: usize,
100 sigma: f64,
101) where
102 Module<B>:
103 TestModuleFamily<B> + GGLWEEncryptSkFamily<B> + GLWEDecryptFamily<B> + GLWEKeyswitchFamily<B> + GGLWEExecLayoutFamily<B>,
104 B: TestScratchFamily<B>,
105{
106 let rows: usize = k_in.div_ceil(basek * digits);
107 let digits_in: usize = 1;
108
109 let mut ct_gglwe_s0s1: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc(
110 module,
111 basek,
112 k_in,
113 rows,
114 digits_in,
115 rank_in_s0s1,
116 rank_out_s0s1,
117 );
118 let mut ct_gglwe_s1s2: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc(
119 module,
120 basek,
121 k_ksk,
122 rows,
123 digits,
124 rank_out_s0s1,
125 rank_out_s1s2,
126 );
127 let mut ct_gglwe_s0s2: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc(
128 module,
129 basek,
130 k_out,
131 rows,
132 digits_in,
133 rank_in_s0s1,
134 rank_out_s1s2,
135 );
136
137 let mut source_xs: Source = Source::new([0u8; 32]);
138 let mut source_xe: Source = Source::new([0u8; 32]);
139 let mut source_xa: Source = Source::new([0u8; 32]);
140
141 let mut scratch_enc: ScratchOwned<B> = ScratchOwned::alloc(GLWESwitchingKey::encrypt_sk_scratch_space(
142 module,
143 basek,
144 k_ksk,
145 rank_in_s0s1 | rank_out_s0s1,
146 rank_out_s0s1 | rank_out_s1s2,
147 ));
148 let mut scratch_apply: ScratchOwned<B> = ScratchOwned::alloc(GLWESwitchingKey::keyswitch_scratch_space(
149 module,
150 basek,
151 k_out,
152 k_in,
153 k_ksk,
154 digits,
155 ct_gglwe_s1s2.rank_in(),
156 ct_gglwe_s1s2.rank_out(),
157 ));
158
159 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_in_s0s1);
160 sk0.fill_ternary_prob(0.5, &mut source_xs);
161
162 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_out_s0s1);
163 sk1.fill_ternary_prob(0.5, &mut source_xs);
164
165 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_out_s1s2);
166 sk2.fill_ternary_prob(0.5, &mut source_xs);
167 let sk2_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk2);
168
169 // gglwe_{s1}(s0) = s0 -> s1
170 ct_gglwe_s0s1.encrypt_sk(
171 module,
172 &sk0,
173 &sk1,
174 &mut source_xa,
175 &mut source_xe,
176 sigma,
177 scratch_enc.borrow(),
178 );
179
180 // gglwe_{s2}(s1) -> s1 -> s2
181 ct_gglwe_s1s2.encrypt_sk(
182 module,
183 &sk1,
184 &sk2,
185 &mut source_xa,
186 &mut source_xe,
187 sigma,
188 scratch_enc.borrow(),
189 );
190
191 let ct_gglwe_s1s2_exec: GLWESwitchingKeyExec<Vec<u8>, B> =
192 GLWESwitchingKeyExec::from(module, &ct_gglwe_s1s2, scratch_apply.borrow());
193
194 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
195 ct_gglwe_s0s2.keyswitch(
196 module,
197 &ct_gglwe_s0s1,
198 &ct_gglwe_s1s2_exec,
199 scratch_apply.borrow(),
200 );
201
202 let max_noise: f64 = log2_std_noise_gglwe_product(
203 module.n() as f64,
204 basek * digits,
205 0.5,
206 0.5,
207 0f64,
208 sigma * sigma,
209 0f64,
210 rank_out_s0s1 as f64,
211 k_in,
212 k_ksk,
213 );
214
215 ct_gglwe_s0s2
216 .key
217 .assert_noise(module, &sk2_exec, &sk0.data, max_noise + 0.5);
218}
219
220pub(crate) fn test_keyswitch_inplace<B: Backend>(
221 module: &Module<B>,
222 basek: usize,
223 k_ct: usize,
224 k_ksk: usize,
225 digits: usize,
226 rank_in: usize,
227 rank_out: usize,
228 sigma: f64,
229) where
230 Module<B>: TestModuleFamily<B>
231 + GLWESwitchingKeyEncryptSkFamily<B>
232 + GLWEKeyswitchFamily<B>
233 + GGLWEExecLayoutFamily<B>
234 + GLWEDecryptFamily<B>,
235 B: TestScratchFamily<B>,
236{
237 let rows: usize = k_ct.div_ceil(basek * digits);
238 let digits_in: usize = 1;
239
240 let mut ct_gglwe_s0s1: GLWESwitchingKey<Vec<u8>> =
241 GLWESwitchingKey::alloc(module, basek, k_ct, rows, digits_in, rank_in, rank_out);
242 let mut ct_gglwe_s1s2: GLWESwitchingKey<Vec<u8>> =
243 GLWESwitchingKey::alloc(module, basek, k_ksk, rows, digits, rank_out, rank_out);
244
245 let mut source_xs: Source = Source::new([0u8; 32]);
246 let mut source_xe: Source = Source::new([0u8; 32]);
247 let mut source_xa: Source = Source::new([0u8; 32]);
248
249 let mut scratch_enc: ScratchOwned<B> = ScratchOwned::alloc(GLWESwitchingKey::encrypt_sk_scratch_space(
250 module,
251 basek,
252 k_ksk,
253 rank_in | rank_out,
254 rank_out,
255 ));
256 let mut scratch_apply: ScratchOwned<B> = ScratchOwned::alloc(GLWESwitchingKey::keyswitch_inplace_scratch_space(
257 module, basek, k_ct, k_ksk, digits, rank_out,
258 ));
259
260 let var_xs: f64 = 0.5;
261
262 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_in);
263 sk0.fill_ternary_prob(var_xs, &mut source_xs);
264
265 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_out);
266 sk1.fill_ternary_prob(var_xs, &mut source_xs);
267
268 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_out);
269 sk2.fill_ternary_prob(var_xs, &mut source_xs);
270 let sk2_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk2);
271
272 // gglwe_{s1}(s0) = s0 -> s1
273 ct_gglwe_s0s1.encrypt_sk(
274 module,
275 &sk0,
276 &sk1,
277 &mut source_xa,
278 &mut source_xe,
279 sigma,
280 scratch_enc.borrow(),
281 );
282
283 // gglwe_{s2}(s1) -> s1 -> s2
284 ct_gglwe_s1s2.encrypt_sk(
285 module,
286 &sk1,
287 &sk2,
288 &mut source_xa,
289 &mut source_xe,
290 sigma,
291 scratch_enc.borrow(),
292 );
293
294 let ct_gglwe_s1s2_exec: GLWESwitchingKeyExec<Vec<u8>, B> =
295 GLWESwitchingKeyExec::from(module, &ct_gglwe_s1s2, scratch_apply.borrow());
296
297 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
298 ct_gglwe_s0s1.keyswitch_inplace(module, &ct_gglwe_s1s2_exec, scratch_apply.borrow());
299
300 let ct_gglwe_s0s2: GLWESwitchingKey<Vec<u8>> = ct_gglwe_s0s1;
301
302 let max_noise: f64 = log2_std_noise_gglwe_product(
303 module.n() as f64,
304 basek * digits,
305 var_xs,
306 var_xs,
307 0f64,
308 sigma * sigma,
309 0f64,
310 rank_out as f64,
311 k_ct,
312 k_ksk,
313 );
314
315 ct_gglwe_s0s2
316 .key
317 .assert_noise(module, &sk2_exec, &sk0.data, max_noise + 0.5);
318}
319
320pub(crate) fn test_external_product<B: Backend>(
321 module: &Module<B>,
322 basek: usize,
323 k_out: usize,
324 k_in: usize,
325 k_ggsw: usize,
326 digits: usize,
327 rank_in: usize,
328 rank_out: usize,
329 sigma: f64,
330) where
331 Module<B>: TestModuleFamily<B>
332 + GLWESwitchingKeyEncryptSkFamily<B>
333 + GLWEExternalProductFamily<B>
334 + GGSWLayoutFamily<B>
335 + GLWEDecryptFamily<B>
336 + VecZnxRotateInplace,
337 B: TestScratchFamily<B>,
338{
339 let rows: usize = k_in.div_ceil(basek * digits);
340 let digits_in: usize = 1;
341
342 let mut ct_gglwe_in: GLWESwitchingKey<Vec<u8>> =
343 GLWESwitchingKey::alloc(module, basek, k_in, rows, digits_in, rank_in, rank_out);
344 let mut ct_gglwe_out: GLWESwitchingKey<Vec<u8>> =
345 GLWESwitchingKey::alloc(module, basek, k_out, rows, digits_in, rank_in, rank_out);
346 let mut ct_rgsw: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(module, basek, k_ggsw, rows, digits, rank_out);
347
348 let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.scalar_znx_alloc(1);
349
350 let mut source_xs: Source = Source::new([0u8; 32]);
351 let mut source_xe: Source = Source::new([0u8; 32]);
352 let mut source_xa: Source = Source::new([0u8; 32]);
353
354 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
355 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_in, rank_in, rank_out)
356 | GLWESwitchingKey::external_product_scratch_space(module, basek, k_out, k_in, k_ggsw, digits, rank_out)
357 | GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k_ggsw, rank_out),
358 );
359
360 let r: usize = 1;
361
362 pt_rgsw.to_mut().raw_mut()[r] = 1; // X^{r}
363
364 let var_xs: f64 = 0.5;
365
366 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_in);
367 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
368
369 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_out);
370 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
371 let sk_out_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_out);
372
373 // gglwe_{s1}(s0) = s0 -> s1
374 ct_gglwe_in.encrypt_sk(
375 module,
376 &sk_in,
377 &sk_out,
378 &mut source_xa,
379 &mut source_xe,
380 sigma,
381 scratch.borrow(),
382 );
383
384 ct_rgsw.encrypt_sk(
385 module,
386 &pt_rgsw,
387 &sk_out_exec,
388 &mut source_xa,
389 &mut source_xe,
390 sigma,
391 scratch.borrow(),
392 );
393
394 let mut ct_rgsw_exec: GGSWCiphertextExec<Vec<u8>, B> =
395 GGSWCiphertextExec::alloc(module, basek, k_ggsw, rows, digits, rank_out);
396
397 ct_rgsw_exec.prepare(module, &ct_rgsw, scratch.borrow());
398
399 // gglwe_(m) (x) RGSW_(X^k) = gglwe_(m * X^k)
400 ct_gglwe_out.external_product(module, &ct_gglwe_in, &ct_rgsw_exec, scratch.borrow());
401
402 (0..rank_in).for_each(|i| {
403 module.vec_znx_rotate_inplace(r as i64, &mut sk_in.data, i); // * X^{r}
404 });
405
406 let var_gct_err_lhs: f64 = sigma * sigma;
407 let var_gct_err_rhs: f64 = 0f64;
408
409 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
410 let var_a0_err: f64 = sigma * sigma;
411 let var_a1_err: f64 = 1f64 / 12f64;
412
413 let max_noise: f64 = noise_ggsw_product(
414 module.n() as f64,
415 basek * digits,
416 var_xs,
417 var_msg,
418 var_a0_err,
419 var_a1_err,
420 var_gct_err_lhs,
421 var_gct_err_rhs,
422 rank_out as f64,
423 k_in,
424 k_ggsw,
425 );
426
427 ct_gglwe_out
428 .key
429 .assert_noise(module, &sk_out_exec, &sk_in.data, max_noise + 0.5);
430}
431
432pub(crate) fn test_external_product_inplace<B: Backend>(
433 module: &Module<B>,
434 basek: usize,
435 k_ct: usize,
436 k_ggsw: usize,
437 digits: usize,
438 rank_in: usize,
439 rank_out: usize,
440 sigma: f64,
441) where
442 Module<B>: TestModuleFamily<B>
443 + GLWESwitchingKeyEncryptSkFamily<B>
444 + GLWEExternalProductFamily<B>
445 + GGSWLayoutFamily<B>
446 + GLWEDecryptFamily<B>
447 + VecZnxRotateInplace,
448 B: TestScratchFamily<B>,
449{
450 let rows: usize = k_ct.div_ceil(basek * digits);
451
452 let digits_in: usize = 1;
453
454 let mut ct_gglwe: GLWESwitchingKey<Vec<u8>> =
455 GLWESwitchingKey::alloc(module, basek, k_ct, rows, digits_in, rank_in, rank_out);
456 let mut ct_rgsw: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(module, basek, k_ggsw, rows, digits, rank_out);
457
458 let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.scalar_znx_alloc(1);
459
460 let mut source_xs: Source = Source::new([0u8; 32]);
461 let mut source_xe: Source = Source::new([0u8; 32]);
462 let mut source_xa: Source = Source::new([0u8; 32]);
463
464 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
465 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_ct, rank_in, rank_out)
466 | GLWESwitchingKey::external_product_inplace_scratch_space(module, basek, k_ct, k_ggsw, digits, rank_out)
467 | GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k_ggsw, rank_out),
468 );
469
470 let r: usize = 1;
471
472 pt_rgsw.to_mut().raw_mut()[r] = 1; // X^{r}
473
474 let var_xs: f64 = 0.5;
475
476 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_in);
477 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
478
479 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank_out);
480 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
481 let sk_out_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_out);
482
483 // gglwe_{s1}(s0) = s0 -> s1
484 ct_gglwe.encrypt_sk(
485 module,
486 &sk_in,
487 &sk_out,
488 &mut source_xa,
489 &mut source_xe,
490 sigma,
491 scratch.borrow(),
492 );
493
494 ct_rgsw.encrypt_sk(
495 module,
496 &pt_rgsw,
497 &sk_out_exec,
498 &mut source_xa,
499 &mut source_xe,
500 sigma,
501 scratch.borrow(),
502 );
503
504 let mut ct_rgsw_exec: GGSWCiphertextExec<Vec<u8>, B> =
505 GGSWCiphertextExec::alloc(module, basek, k_ggsw, rows, digits, rank_out);
506
507 ct_rgsw_exec.prepare(module, &ct_rgsw, scratch.borrow());
508
509 // gglwe_(m) (x) RGSW_(X^k) = gglwe_(m * X^k)
510 ct_gglwe.external_product_inplace(module, &ct_rgsw_exec, scratch.borrow());
511
512 (0..rank_in).for_each(|i| {
513 module.vec_znx_rotate_inplace(r as i64, &mut sk_in.data, i); // * X^{r}
514 });
515
516 let var_gct_err_lhs: f64 = sigma * sigma;
517 let var_gct_err_rhs: f64 = 0f64;
518
519 let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
520 let var_a0_err: f64 = sigma * sigma;
521 let var_a1_err: f64 = 1f64 / 12f64;
522
523 let max_noise: f64 = noise_ggsw_product(
524 module.n() as f64,
525 basek * digits,
526 var_xs,
527 var_msg,
528 var_a0_err,
529 var_a1_err,
530 var_gct_err_lhs,
531 var_gct_err_rhs,
532 rank_out as f64,
533 k_ct,
534 k_ggsw,
535 );
536
537 ct_gglwe
538 .key
539 .assert_noise(module, &sk_out_exec, &sk_in.data, max_noise + 0.5);
540}