poulpy/rlwe/src/elem_rlwe.rs
rlwe/src/elem_rlwe.rsRaw
1use base2k::{
2 AddNormal, Backend, FFT64, FillUniform, Module, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps,
3 ScalarZnxDftToRef, Scratch, VecZnx, VecZnxAlloc, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc,
4 VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps, VecZnxToMut, VecZnxToRef,
5};
6use sampling::source::Source;
7
8use crate::{
9 elem::Infos,
10 keys::{PublicKey, SecretDistribution, SecretKeyDft},
11 utils::derive_size,
12};
13
14pub struct RLWECt<C> {
15 pub data: VecZnx<C>,
16 pub log_base2k: usize,
17 pub log_k: usize,
18}
19
20impl RLWECt<Vec<u8>> {
21 pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize, cols: usize) -> Self {
22 Self {
23 data: module.new_vec_znx(cols, derive_size(log_base2k, log_k)),
24 log_base2k: log_base2k,
25 log_k: log_k,
26 }
27 }
28}
29
30impl<T> Infos for RLWECt<T> {
31 type Inner = VecZnx<T>;
32
33 fn inner(&self) -> &Self::Inner {
34 &self.data
35 }
36
37 fn log_base2k(&self) -> usize {
38 self.log_base2k
39 }
40
41 fn log_k(&self) -> usize {
42 self.log_k
43 }
44}
45
46impl<C> VecZnxToMut for RLWECt<C>
47where
48 VecZnx<C>: VecZnxToMut,
49{
50 fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
51 self.data.to_mut()
52 }
53}
54
55impl<C> VecZnxToRef for RLWECt<C>
56where
57 VecZnx<C>: VecZnxToRef,
58{
59 fn to_ref(&self) -> VecZnx<&[u8]> {
60 self.data.to_ref()
61 }
62}
63
64pub struct RLWEPt<C> {
65 pub data: VecZnx<C>,
66 pub log_base2k: usize,
67 pub log_k: usize,
68}
69
70impl<T> Infos for RLWEPt<T> {
71 type Inner = VecZnx<T>;
72
73 fn inner(&self) -> &Self::Inner {
74 &self.data
75 }
76
77 fn log_base2k(&self) -> usize {
78 self.log_base2k
79 }
80
81 fn log_k(&self) -> usize {
82 self.log_k
83 }
84}
85
86impl<C> VecZnxToMut for RLWEPt<C>
87where
88 VecZnx<C>: VecZnxToMut,
89{
90 fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
91 self.data.to_mut()
92 }
93}
94
95impl<C> VecZnxToRef for RLWEPt<C>
96where
97 VecZnx<C>: VecZnxToRef,
98{
99 fn to_ref(&self) -> VecZnx<&[u8]> {
100 self.data.to_ref()
101 }
102}
103
104impl RLWEPt<Vec<u8>> {
105 pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize) -> Self {
106 Self {
107 data: module.new_vec_znx(1, derive_size(log_base2k, log_k)),
108 log_base2k: log_base2k,
109 log_k: log_k,
110 }
111 }
112}
113
114pub struct RLWECtDft<C, B: Backend> {
115 pub data: VecZnxDft<C, B>,
116 pub log_base2k: usize,
117 pub log_k: usize,
118}
119
120impl<B: Backend> RLWECtDft<Vec<u8>, B> {
121 pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, cols: usize) -> Self {
122 Self {
123 data: module.new_vec_znx_dft(cols, derive_size(log_base2k, log_k)),
124 log_base2k: log_base2k,
125 log_k: log_k,
126 }
127 }
128}
129
130impl<T, B: Backend> Infos for RLWECtDft<T, B> {
131 type Inner = VecZnxDft<T, B>;
132
133 fn inner(&self) -> &Self::Inner {
134 &self.data
135 }
136
137 fn log_base2k(&self) -> usize {
138 self.log_base2k
139 }
140
141 fn log_k(&self) -> usize {
142 self.log_k
143 }
144}
145
146impl<C, B: Backend> VecZnxDftToMut<B> for RLWECtDft<C, B>
147where
148 VecZnxDft<C, B>: VecZnxDftToMut<B>,
149{
150 fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
151 self.data.to_mut()
152 }
153}
154
155impl<C, B: Backend> VecZnxDftToRef<B> for RLWECtDft<C, B>
156where
157 VecZnxDft<C, B>: VecZnxDftToRef<B>,
158{
159 fn to_ref(&self) -> VecZnxDft<&[u8], B> {
160 self.data.to_ref()
161 }
162}
163
164impl RLWECt<Vec<u8>> {
165 pub fn encrypt_sk_scratch_bytes<B: Backend>(module: &Module<B>, size: usize) -> usize {
166 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
167 }
168
169 pub fn encrypt_pk_scratch_bytes<B: Backend>(module: &Module<B>, pk_size: usize) -> usize {
170 ((module.bytes_of_vec_znx_dft(1, pk_size) + module.bytes_of_vec_znx_big(1, pk_size)) | module.bytes_of_scalar_znx(1))
171 + module.bytes_of_scalar_znx_dft(1)
172 + module.vec_znx_big_normalize_tmp_bytes()
173 }
174
175 pub fn decrypt_scratch_bytes<B: Backend>(module: &Module<B>, size: usize) -> usize {
176 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
177 }
178}
179
180pub fn encrypt_rlwe_sk<C, P, S>(
181 module: &Module<FFT64>,
182 ct: &mut RLWECt<C>,
183 pt: Option<(&RLWEPt<P>, usize)>,
184 sk_dft: &SecretKeyDft<S, FFT64>,
185 source_xa: &mut Source,
186 source_xe: &mut Source,
187 sigma: f64,
188 bound: f64,
189 scratch: &mut Scratch,
190) where
191 VecZnx<C>: VecZnxToMut + VecZnxToRef,
192 VecZnx<P>: VecZnxToRef,
193 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
194{
195 let log_base2k: usize = ct.log_base2k();
196 let log_k: usize = ct.log_k();
197 let size: usize = ct.size();
198
199 // c1 = a
200 ct.data.fill_uniform(log_base2k, 1, size, source_xa);
201
202 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
203
204 {
205 let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
206 module.vec_znx_dft(&mut c0_dft, 0, ct, 1);
207
208 // c0_dft = DFT(a) * DFT(s)
209 module.svp_apply_inplace(&mut c0_dft, 0, sk_dft, 0);
210
211 // c0_big = IDFT(c0_dft)
212 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
213 }
214
215 // c0_big = m - c0_big
216 if let Some((pt, col)) = pt {
217 match col {
218 0 => module.vec_znx_big_sub_small_b_inplace(&mut c0_big, 0, pt, 0),
219 1 => {
220 module.vec_znx_big_negate_inplace(&mut c0_big, 0);
221 module.vec_znx_add_inplace(ct, 1, pt, 0);
222 module.vec_znx_normalize_inplace(log_base2k, ct, 1, scratch_1);
223 }
224 _ => panic!("invalid target column: {}", col),
225 }
226 } else {
227 module.vec_znx_big_negate_inplace(&mut c0_big, 0);
228 }
229 // c0_big += e
230 c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, bound);
231
232 // c0 = norm(c0_big = -as + m + e)
233 module.vec_znx_big_normalize(log_base2k, ct, 0, &c0_big, 0, scratch_1);
234}
235
236pub fn decrypt_rlwe<P, C, S>(
237 module: &Module<FFT64>,
238 pt: &mut RLWEPt<P>,
239 ct: &RLWECt<C>,
240 sk_dft: &SecretKeyDft<S, FFT64>,
241 scratch: &mut Scratch,
242) where
243 VecZnx<P>: VecZnxToMut + VecZnxToRef,
244 VecZnx<C>: VecZnxToRef,
245 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
246{
247 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, ct.size()); // TODO optimize size when pt << ct
248
249 {
250 let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, ct.size()); // TODO optimize size when pt << ct
251 module.vec_znx_dft(&mut c0_dft, 0, ct, 1);
252
253 // c0_dft = DFT(a) * DFT(s)
254 module.svp_apply_inplace(&mut c0_dft, 0, sk_dft, 0);
255
256 // c0_big = IDFT(c0_dft)
257 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
258 }
259
260 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
261 module.vec_znx_big_add_small_inplace(&mut c0_big, 0, ct, 0);
262
263 // pt = norm(BIG(m + e))
264 module.vec_znx_big_normalize(ct.log_base2k(), pt, 0, &mut c0_big, 0, scratch_1);
265
266 pt.log_base2k = ct.log_base2k();
267 pt.log_k = pt.log_k().min(ct.log_k());
268}
269
270impl<C> RLWECt<C> {
271 pub fn encrypt_sk<P, S>(
272 &mut self,
273 module: &Module<FFT64>,
274 pt: Option<&RLWEPt<P>>,
275 sk_dft: &SecretKeyDft<S, FFT64>,
276 source_xa: &mut Source,
277 source_xe: &mut Source,
278 sigma: f64,
279 bound: f64,
280 scratch: &mut Scratch,
281 ) where
282 VecZnx<C>: VecZnxToMut + VecZnxToRef,
283 VecZnx<P>: VecZnxToRef,
284 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
285 {
286 if let Some(pt) = pt {
287 encrypt_rlwe_sk(
288 module,
289 self,
290 Some((pt, 0)),
291 sk_dft,
292 source_xa,
293 source_xe,
294 sigma,
295 bound,
296 scratch,
297 )
298 } else {
299 encrypt_rlwe_sk::<C, P, S>(
300 module, self, None, sk_dft, source_xa, source_xe, sigma, bound, scratch,
301 )
302 }
303 }
304
305 pub fn decrypt<P, S>(
306 &self,
307 module: &Module<FFT64>,
308 pt: &mut RLWEPt<P>,
309 sk_dft: &SecretKeyDft<S, FFT64>,
310 scratch: &mut Scratch,
311 ) where
312 VecZnx<P>: VecZnxToMut + VecZnxToRef,
313 VecZnx<C>: VecZnxToRef,
314 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
315 {
316 decrypt_rlwe(module, pt, self, sk_dft, scratch);
317 }
318
319 pub fn encrypt_pk<P, S>(
320 &mut self,
321 module: &Module<FFT64>,
322 pt: Option<&RLWEPt<P>>,
323 pk: &PublicKey<S, FFT64>,
324 source_xu: &mut Source,
325 source_xe: &mut Source,
326 sigma: f64,
327 bound: f64,
328 scratch: &mut Scratch,
329 ) where
330 VecZnx<C>: VecZnxToMut + VecZnxToRef,
331 VecZnx<P>: VecZnxToRef,
332 VecZnxDft<S, FFT64>: VecZnxDftToRef<FFT64>,
333 {
334 encrypt_rlwe_pk(
335 module, self, pt, pk, source_xu, source_xe, sigma, bound, scratch,
336 )
337 }
338}
339
340pub(crate) fn encrypt_zero_rlwe_dft_sk<C, S>(
341 module: &Module<FFT64>,
342 ct: &mut RLWECtDft<C, FFT64>,
343 sk: &SecretKeyDft<S, FFT64>,
344 source_xa: &mut Source,
345 source_xe: &mut Source,
346 sigma: f64,
347 bound: f64,
348 scratch: &mut Scratch,
349) where
350 VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
351 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
352{
353 let log_base2k: usize = ct.log_base2k();
354 let log_k: usize = ct.log_k();
355 let size: usize = ct.size();
356
357 #[cfg(debug_assertions)]
358 {
359 match sk.dist {
360 SecretDistribution::NONE => panic!("invalid sk.dist = SecretDistribution::NONE"),
361 _ => {}
362 }
363 assert_eq!(ct.cols(), 2);
364 }
365
366 // ct[1] = DFT(a)
367 {
368 let (mut tmp_znx, _) = scratch.tmp_vec_znx(module, 1, size);
369 tmp_znx.fill_uniform(log_base2k, 0, size, source_xa);
370 module.vec_znx_dft(ct, 1, &tmp_znx, 0);
371 }
372
373 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
374
375 {
376 let (mut tmp_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
377 // c0_dft = ct[1] * DFT(s)
378 module.svp_apply(&mut tmp_dft, 0, sk, 0, ct, 1);
379
380 // c0_big = IDFT(c0_dft)
381 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut tmp_dft, 0);
382 }
383
384 // c0_big += e
385 c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, bound);
386
387 // c0 = norm(c0_big = -as - e), NOTE: e is centered at 0.
388 let (mut tmp_znx, scratch_2) = scratch_1.tmp_vec_znx(module, 1, size);
389 module.vec_znx_big_normalize(log_base2k, &mut tmp_znx, 0, &c0_big, 0, scratch_2);
390 module.vec_znx_negate_inplace(&mut tmp_znx, 0);
391 // ct[0] = DFT(-as + e)
392 module.vec_znx_dft(ct, 0, &tmp_znx, 0);
393}
394
395impl RLWECtDft<Vec<u8>, FFT64> {
396 pub fn encrypt_zero_sk_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
397 (module.bytes_of_vec_znx(1, size) | module.bytes_of_vec_znx_dft(1, size))
398 + module.bytes_of_vec_znx_big(1, size)
399 + module.bytes_of_vec_znx(1, size)
400 + module.vec_znx_big_normalize_tmp_bytes()
401 }
402
403 pub fn decrypt_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
404 (module.vec_znx_big_normalize_tmp_bytes()
405 | module.bytes_of_vec_znx_dft(1, size)
406 | (module.bytes_of_vec_znx_big(1, size) + module.vec_znx_idft_tmp_bytes()))
407 + module.bytes_of_vec_znx_big(1, size)
408 }
409}
410
411pub fn decrypt_rlwe_dft<P, C, S>(
412 module: &Module<FFT64>,
413 pt: &mut RLWEPt<P>,
414 ct: &RLWECtDft<C, FFT64>,
415 sk: &SecretKeyDft<S, FFT64>,
416 scratch: &mut Scratch,
417) where
418 VecZnx<P>: VecZnxToMut + VecZnxToRef,
419 VecZnxDft<C, FFT64>: VecZnxDftToRef<FFT64>,
420 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
421{
422 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, ct.size()); // TODO optimize size when pt << ct
423
424 {
425 let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, ct.size()); // TODO optimize size when pt << ct
426 // c0_dft = DFT(a) * DFT(s)
427 module.svp_apply(&mut c0_dft, 0, sk, 0, ct, 1);
428 // c0_big = IDFT(c0_dft)
429 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
430 }
431
432 {
433 let (mut c1_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, ct.size());
434 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
435 module.vec_znx_idft(&mut c1_big, 0, ct, 0, scratch_2);
436 module.vec_znx_big_add_inplace(&mut c0_big, 0, &c1_big, 0);
437 }
438
439 // pt = norm(BIG(m + e))
440 module.vec_znx_big_normalize(ct.log_base2k(), pt, 0, &mut c0_big, 0, scratch_1);
441
442 pt.log_base2k = ct.log_base2k();
443 pt.log_k = pt.log_k().min(ct.log_k());
444}
445
446impl<C> RLWECtDft<C, FFT64> {
447 pub(crate) fn encrypt_zero_sk<S>(
448 &mut self,
449 module: &Module<FFT64>,
450 sk_dft: &SecretKeyDft<S, FFT64>,
451 source_xa: &mut Source,
452 source_xe: &mut Source,
453 sigma: f64,
454 bound: f64,
455 scratch: &mut Scratch,
456 ) where
457 VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
458 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
459 {
460 encrypt_zero_rlwe_dft_sk(
461 module, self, sk_dft, source_xa, source_xe, sigma, bound, scratch,
462 )
463 }
464
465 pub fn decrypt<P, S>(
466 &self,
467 module: &Module<FFT64>,
468 pt: &mut RLWEPt<P>,
469 sk_dft: &SecretKeyDft<S, FFT64>,
470 scratch: &mut Scratch,
471 ) where
472 VecZnx<P>: VecZnxToMut + VecZnxToRef,
473 VecZnxDft<C, FFT64>: VecZnxDftToRef<FFT64>,
474 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
475 {
476 decrypt_rlwe_dft(module, pt, self, sk_dft, scratch);
477 }
478}
479
480pub(crate) fn encrypt_rlwe_pk<C, P, S>(
481 module: &Module<FFT64>,
482 ct: &mut RLWECt<C>,
483 pt: Option<&RLWEPt<P>>,
484 pk: &PublicKey<S, FFT64>,
485 source_xu: &mut Source,
486 source_xe: &mut Source,
487 sigma: f64,
488 bound: f64,
489 scratch: &mut Scratch,
490) where
491 VecZnx<C>: VecZnxToMut + VecZnxToRef,
492 VecZnx<P>: VecZnxToRef,
493 VecZnxDft<S, FFT64>: VecZnxDftToRef<FFT64>,
494{
495 #[cfg(debug_assertions)]
496 {
497 assert_eq!(ct.log_base2k(), pk.log_base2k());
498 assert_eq!(ct.n(), module.n());
499 assert_eq!(pk.n(), module.n());
500 if let Some(pt) = pt {
501 assert_eq!(pt.log_base2k(), pk.log_base2k());
502 assert_eq!(pt.n(), module.n());
503 }
504 }
505
506 let log_base2k: usize = pk.log_base2k();
507 let size_pk: usize = pk.size();
508
509 // Generates u according to the underlying secret distribution.
510 let (mut u_dft, scratch_1) = scratch.tmp_scalar_znx_dft(module, 1);
511
512 {
513 let (mut u, _) = scratch_1.tmp_scalar_znx(module, 1);
514 match pk.dist {
515 SecretDistribution::NONE => panic!(
516 "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through Self::generate"
517 ),
518 SecretDistribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
519 SecretDistribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
520 }
521
522 module.svp_prepare(&mut u_dft, 0, &u, 0);
523 }
524
525 let (mut tmp_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, size_pk); // TODO optimize size (e.g. when encrypting at low homomorphic capacity)
526 let (mut tmp_dft, scratch_3) = scratch_2.tmp_vec_znx_dft(module, 1, size_pk); // TODO optimize size (e.g. when encrypting at low homomorphic capacity)
527
528 // ct[0] = pk[0] * u + m + e0
529 module.svp_apply(&mut tmp_dft, 0, &u_dft, 0, pk, 0);
530 module.vec_znx_idft_tmp_a(&mut tmp_big, 0, &mut tmp_dft, 0);
531 tmp_big.add_normal(log_base2k, 0, pk.log_k(), source_xe, sigma, bound);
532
533 if let Some(pt) = pt {
534 module.vec_znx_big_add_small_inplace(&mut tmp_big, 0, pt, 0);
535 }
536
537 module.vec_znx_big_normalize(log_base2k, ct, 0, &tmp_big, 0, scratch_3);
538
539 // ct[1] = pk[1] * u + e1
540 module.svp_apply(&mut tmp_dft, 0, &u_dft, 0, pk, 1);
541 module.vec_znx_idft_tmp_a(&mut tmp_big, 0, &mut tmp_dft, 0);
542 tmp_big.add_normal(log_base2k, 0, pk.log_k(), source_xe, sigma, bound);
543 module.vec_znx_big_normalize(log_base2k, ct, 1, &tmp_big, 0, scratch_3);
544}
545
546#[cfg(test)]
547mod tests {
548 use base2k::{Decoding, Encoding, FFT64, Module, ScratchOwned, Stats, VecZnxOps, ZnxZero};
549 use itertools::izip;
550 use sampling::source::Source;
551
552 use crate::{
553 elem_rlwe::{Infos, RLWECt, RLWECtDft, RLWEPt},
554 keys::{PublicKey, SecretKey, SecretKeyDft},
555 };
556
557 #[test]
558 fn encrypt_sk_fft64() {
559 let module: Module<FFT64> = Module::<FFT64>::new(32);
560 let log_base2k: usize = 8;
561 let log_k_ct: usize = 54;
562 let log_k_pt: usize = 30;
563
564 let sigma: f64 = 3.2;
565 let bound: f64 = sigma * 6.0;
566
567 let mut ct: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_ct, 2);
568 let mut pt: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_pt);
569
570 let mut source_xs: Source = Source::new([0u8; 32]);
571 let mut source_xe: Source = Source::new([0u8; 32]);
572 let mut source_xa: Source = Source::new([0u8; 32]);
573
574 let mut scratch: ScratchOwned = ScratchOwned::new(
575 RLWECt::encrypt_sk_scratch_bytes(&module, ct.size()) | RLWECt::decrypt_scratch_bytes(&module, ct.size()),
576 );
577
578 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
579 sk.fill_ternary_prob(0.5, &mut source_xs);
580
581 let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
582 sk_dft.dft(&module, &sk);
583
584 let mut data_want: Vec<i64> = vec![0i64; module.n()];
585
586 data_want
587 .iter_mut()
588 .for_each(|x| *x = source_xa.next_i64() & 0xFF);
589
590 pt.data
591 .encode_vec_i64(0, log_base2k, log_k_pt, &data_want, 10);
592
593 ct.encrypt_sk(
594 &module,
595 Some(&pt),
596 &sk_dft,
597 &mut source_xa,
598 &mut source_xe,
599 sigma,
600 bound,
601 scratch.borrow(),
602 );
603
604 pt.data.zero();
605
606 ct.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
607
608 let mut data_have: Vec<i64> = vec![0i64; module.n()];
609
610 pt.data
611 .decode_vec_i64(0, log_base2k, pt.size() * log_base2k, &mut data_have);
612
613 // TODO: properly assert the decryption noise through std(dec(ct) - pt)
614 let scale: f64 = (1 << (pt.size() * log_base2k - log_k_pt)) as f64;
615 izip!(data_want.iter(), data_have.iter()).for_each(|(a, b)| {
616 let b_scaled = (*b as f64) / scale;
617 assert!(
618 (*a as f64 - b_scaled).abs() < 0.1,
619 "{} {}",
620 *a as f64,
621 b_scaled
622 )
623 });
624
625 module.free();
626 }
627
628 #[test]
629 fn encrypt_zero_sk_fft64() {
630 let module: Module<FFT64> = Module::<FFT64>::new(1024);
631 let log_base2k: usize = 8;
632 let log_k_ct: usize = 55;
633
634 let sigma: f64 = 3.2;
635 let bound: f64 = sigma * 6.0;
636
637 let mut pt: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
638
639 let mut source_xs: Source = Source::new([0u8; 32]);
640 let mut source_xe: Source = Source::new([1u8; 32]);
641 let mut source_xa: Source = Source::new([0u8; 32]);
642
643 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
644 sk.fill_ternary_prob(0.5, &mut source_xs);
645 let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
646 sk_dft.dft(&module, &sk);
647
648 let mut ct_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_ct, 2);
649
650 let mut scratch: ScratchOwned = ScratchOwned::new(
651 RLWECtDft::decrypt_scratch_bytes(&module, ct_dft.size())
652 | RLWECtDft::encrypt_zero_sk_scratch_bytes(&module, ct_dft.size()),
653 );
654
655 ct_dft.encrypt_zero_sk(
656 &module,
657 &sk_dft,
658 &mut source_xa,
659 &mut source_xe,
660 sigma,
661 bound,
662 scratch.borrow(),
663 );
664 ct_dft.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
665
666 assert!((sigma - pt.data.std(0, log_base2k) * (log_k_ct as f64).exp2()) <= 0.2);
667 module.free();
668 }
669
670 #[test]
671 fn encrypt_pk_fft64() {
672 let module: Module<FFT64> = Module::<FFT64>::new(32);
673 let log_base2k: usize = 8;
674 let log_k_ct: usize = 54;
675 let log_k_pk: usize = 64;
676
677 let sigma: f64 = 3.2;
678 let bound: f64 = sigma * 6.0;
679
680 let mut ct: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_ct, 2);
681 let mut pt_want: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
682
683 let mut source_xs: Source = Source::new([0u8; 32]);
684 let mut source_xe: Source = Source::new([0u8; 32]);
685 let mut source_xa: Source = Source::new([0u8; 32]);
686 let mut source_xu: Source = Source::new([0u8; 32]);
687
688 let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
689 sk.fill_ternary_prob(0.5, &mut source_xs);
690 let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
691 sk_dft.dft(&module, &sk);
692
693 let mut pk: PublicKey<Vec<u8>, FFT64> = PublicKey::new(&module, log_base2k, log_k_pk);
694 pk.generate(
695 &module,
696 &sk_dft,
697 &mut source_xa,
698 &mut source_xe,
699 sigma,
700 bound,
701 );
702
703 let mut scratch: ScratchOwned = ScratchOwned::new(
704 RLWECt::encrypt_sk_scratch_bytes(&module, ct.size())
705 | RLWECt::decrypt_scratch_bytes(&module, ct.size())
706 | RLWECt::encrypt_pk_scratch_bytes(&module, pk.size()),
707 );
708
709 let mut data_want: Vec<i64> = vec![0i64; module.n()];
710
711 data_want
712 .iter_mut()
713 .for_each(|x| *x = source_xa.next_i64() & 0);
714
715 pt_want
716 .data
717 .encode_vec_i64(0, log_base2k, log_k_ct, &data_want, 10);
718
719 ct.encrypt_pk(
720 &module,
721 Some(&pt_want),
722 &pk,
723 &mut source_xu,
724 &mut source_xe,
725 sigma,
726 bound,
727 scratch.borrow(),
728 );
729
730 let mut pt_have: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
731
732 ct.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
733
734 module.vec_znx_sub_ab_inplace(&mut pt_want, 0, &pt_have, 0);
735
736 assert!(((1.0f64 / 12.0).sqrt() - pt_want.data.std(0, log_base2k) * (log_k_ct as f64).exp2()).abs() < 0.2);
737
738 module.free();
739 }
740}