poulpy/poulpy-core/src/encryption/glwe_ct.rs
poulpy-core/src/encryption/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchAvailable, SvpApplyDftToDft, SvpApplyDftToDftInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeSvpPPol,
4 TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxBigAddNormal, VecZnxBigAddSmallInplace,
5 VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume,
6 VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
7 },
8 layouts::{Backend, DataMut, Module, ScalarZnx, Scratch, VecZnx, VecZnxBig, VecZnxToMut, ZnxInfos, ZnxZero},
9 source::Source,
10};
11
12use crate::{
13 dist::Distribution,
14 encryption::{SIGMA, SIGMA_BOUND},
15 layouts::{
16 GLWECiphertext, GLWECiphertextToMut, GLWEInfos, GLWEPlaintext, GLWEPlaintextToRef, LWEInfos,
17 prepared::{GLWEPublicKeyPrepared, GLWEPublicKeyPreparedToRef, GLWESecretPrepared, GLWESecretPreparedToRef},
18 },
19};
20
21impl GLWECiphertext<Vec<u8>> {
22 pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
23 where
24 A: GLWEInfos,
25 Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
26 {
27 let size: usize = infos.size();
28 assert_eq!(module.n() as u32, infos.n());
29 module.vec_znx_normalize_tmp_bytes()
30 + 2 * VecZnx::alloc_bytes(module.n(), 1, size)
31 + module.vec_znx_dft_alloc_bytes(1, size)
32 }
33 pub fn encrypt_pk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
34 where
35 A: GLWEInfos,
36 Module<B>: VecZnxDftAllocBytes + SvpPPolAllocBytes + VecZnxBigAllocBytes + VecZnxNormalizeTmpBytes,
37 {
38 let size: usize = infos.size();
39 assert_eq!(module.n() as u32, infos.n());
40 ((module.vec_znx_dft_alloc_bytes(1, size) + module.vec_znx_big_alloc_bytes(1, size))
41 | ScalarZnx::alloc_bytes(module.n(), 1))
42 + module.svp_ppol_alloc_bytes(1)
43 + module.vec_znx_normalize_tmp_bytes()
44 }
45}
46
47impl<D: DataMut> GLWECiphertext<D> {
48 pub fn encrypt_sk<R, P, S, B: Backend>(
49 &mut self,
50 module: &Module<B>,
51 pt: &P,
52 sk: &S,
53 source_xa: &mut Source,
54 source_xe: &mut Source,
55 scratch: &mut Scratch<B>,
56 ) where
57 P: GLWEPlaintextToRef,
58 S: GLWESecretPreparedToRef<B>,
59 Module<B>: GLWEEncryptSk<B>,
60 {
61 module.glwe_encrypt_sk(self, pt, sk, source_xa, source_xe, scratch);
62 }
63
64 pub fn encrypt_zero_sk<S, B: Backend>(
65 &mut self,
66 module: &Module<B>,
67 sk: &S,
68 source_xa: &mut Source,
69 source_xe: &mut Source,
70 scratch: &mut Scratch<B>,
71 ) where
72 S: GLWESecretPreparedToRef<B>,
73 Module<B>: GLWEEncryptZeroSk<B>,
74 {
75 module.glwe_encrypt_zero_sk(self, sk, source_xa, source_xe, scratch);
76 }
77
78 pub fn encrypt_pk<P, K, B: Backend>(
79 &mut self,
80 module: &Module<B>,
81 pt: &P,
82 pk: &K,
83 source_xu: &mut Source,
84 source_xe: &mut Source,
85 scratch: &mut Scratch<B>,
86 ) where
87 P: GLWEPlaintextToRef,
88 K: GLWEPublicKeyPreparedToRef<B>,
89 Module<B>: GLWEEncryptPk<B>,
90 {
91 module.glwe_encrypt_pk(self, pt, pk, source_xu, source_xe, scratch);
92 }
93
94 pub fn encrypt_zero_pk<K, B: Backend>(
95 &mut self,
96 module: &Module<B>,
97 pk: &K,
98 source_xu: &mut Source,
99 source_xe: &mut Source,
100 scratch: &mut Scratch<B>,
101 ) where
102 K: GLWEPublicKeyPreparedToRef<B>,
103 Module<B>: GLWEEncryptZeroPk<B>,
104 {
105 module.glwe_encrypt_zero_pk(self, pk, source_xu, source_xe, scratch);
106 }
107}
108
109pub trait GLWEEncryptSk<B: Backend> {
110 fn glwe_encrypt_sk<R, P, S>(
111 &self,
112 res: &mut R,
113 pt: &P,
114 sk: &S,
115 source_xa: &mut Source,
116 source_xe: &mut Source,
117 scratch: &mut Scratch<B>,
118 ) where
119 R: GLWECiphertextToMut,
120 P: GLWEPlaintextToRef,
121 S: GLWESecretPreparedToRef<B>;
122}
123
124impl<B: Backend> GLWEEncryptSk<B> for Module<B>
125where
126 Module<B>: GLWEEncryptSkInternal<B> + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
127 Scratch<B>: ScratchAvailable,
128{
129 fn glwe_encrypt_sk<R, P, S>(
130 &self,
131 res: &mut R,
132 pt: &P,
133 sk: &S,
134 source_xa: &mut Source,
135 source_xe: &mut Source,
136 scratch: &mut Scratch<B>,
137 ) where
138 R: GLWECiphertextToMut,
139 P: GLWEPlaintextToRef,
140 S: GLWESecretPreparedToRef<B>,
141 {
142 let mut res: GLWECiphertext<&mut [u8]> = res.to_mut();
143 let pt: GLWEPlaintext<&[u8]> = pt.to_ref();
144
145 #[cfg(debug_assertions)]
146 {
147 let sk: GLWESecretPrepared<&[u8], B> = sk.to_ref();
148 assert_eq!(res.rank(), sk.rank());
149 assert_eq!(res.n(), self.n() as u32);
150 assert_eq!(sk.n(), self.n() as u32);
151 assert_eq!(pt.n(), self.n() as u32);
152 assert!(
153 scratch.available() >= GLWECiphertext::encrypt_sk_scratch_space(self, &res),
154 "scratch.available(): {} < GLWECiphertext::encrypt_sk_scratch_space: {}",
155 scratch.available(),
156 GLWECiphertext::encrypt_sk_scratch_space(self, &res)
157 )
158 }
159
160 let cols: usize = (res.rank() + 1).into();
161 self.glwe_encrypt_sk_internal(
162 res.base2k().into(),
163 res.k().into(),
164 res.data_mut(),
165 cols,
166 false,
167 Some((&pt, 0)),
168 sk,
169 source_xa,
170 source_xe,
171 SIGMA,
172 scratch,
173 );
174 }
175}
176
177pub trait GLWEEncryptZeroSk<B: Backend> {
178 fn glwe_encrypt_zero_sk<R, S>(
179 &self,
180 res: &mut R,
181 sk: &S,
182 source_xa: &mut Source,
183 source_xe: &mut Source,
184 scratch: &mut Scratch<B>,
185 ) where
186 R: GLWECiphertextToMut,
187 S: GLWESecretPreparedToRef<B>;
188}
189
190impl<B: Backend> GLWEEncryptZeroSk<B> for Module<B>
191where
192 Module<B>: GLWEEncryptSkInternal<B> + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
193 Scratch<B>: ScratchAvailable,
194{
195 fn glwe_encrypt_zero_sk<R, S>(
196 &self,
197 res: &mut R,
198 sk: &S,
199 source_xa: &mut Source,
200 source_xe: &mut Source,
201 scratch: &mut Scratch<B>,
202 ) where
203 R: GLWECiphertextToMut,
204 S: GLWESecretPreparedToRef<B>,
205 {
206 let mut res: GLWECiphertext<&mut [u8]> = res.to_mut();
207
208 #[cfg(debug_assertions)]
209 {
210 let sk: GLWESecretPrepared<&[u8], B> = sk.to_ref();
211 assert_eq!(res.rank(), sk.rank());
212 assert_eq!(res.n(), self.n() as u32);
213 assert_eq!(sk.n(), self.n() as u32);
214 assert!(
215 scratch.available() >= GLWECiphertext::encrypt_sk_scratch_space(self, &res),
216 "scratch.available(): {} < GLWECiphertext::encrypt_sk_scratch_space: {}",
217 scratch.available(),
218 GLWECiphertext::encrypt_sk_scratch_space(self, &res)
219 )
220 }
221
222 let cols: usize = (res.rank() + 1).into();
223 self.glwe_encrypt_sk_internal(
224 res.base2k().into(),
225 res.k().into(),
226 res.data_mut(),
227 cols,
228 false,
229 None::<(&GLWEPlaintext<Vec<u8>>, usize)>,
230 sk,
231 source_xa,
232 source_xe,
233 SIGMA,
234 scratch,
235 );
236 }
237}
238
239pub trait GLWEEncryptPk<B: Backend> {
240 fn glwe_encrypt_pk<R, P, K>(
241 &self,
242 res: &mut R,
243 pt: &P,
244 pk: &K,
245 source_xu: &mut Source,
246 source_xe: &mut Source,
247 scratch: &mut Scratch<B>,
248 ) where
249 R: GLWECiphertextToMut,
250 P: GLWEPlaintextToRef,
251 K: GLWEPublicKeyPreparedToRef<B>;
252}
253
254impl<B: Backend> GLWEEncryptPk<B> for Module<B>
255where
256 Module<B>: GLWEEncryptPkInternal<B>,
257{
258 fn glwe_encrypt_pk<R, P, K>(
259 &self,
260 res: &mut R,
261 pt: &P,
262 pk: &K,
263 source_xu: &mut Source,
264 source_xe: &mut Source,
265 scratch: &mut Scratch<B>,
266 ) where
267 R: GLWECiphertextToMut,
268 P: GLWEPlaintextToRef,
269 K: GLWEPublicKeyPreparedToRef<B>,
270 {
271 self.glwe_encrypt_pk_internal(res, Some((pt, 0)), pk, source_xu, source_xe, scratch);
272 }
273}
274
275pub trait GLWEEncryptZeroPk<B: Backend> {
276 fn glwe_encrypt_zero_pk<R, K>(
277 &self,
278 res: &mut R,
279 pk: &K,
280 source_xu: &mut Source,
281 source_xe: &mut Source,
282 scratch: &mut Scratch<B>,
283 ) where
284 R: GLWECiphertextToMut,
285 K: GLWEPublicKeyPreparedToRef<B>;
286}
287
288impl<B: Backend> GLWEEncryptZeroPk<B> for Module<B>
289where
290 Module<B>: GLWEEncryptPkInternal<B>,
291{
292 fn glwe_encrypt_zero_pk<R, K>(
293 &self,
294 res: &mut R,
295 pk: &K,
296 source_xu: &mut Source,
297 source_xe: &mut Source,
298 scratch: &mut Scratch<B>,
299 ) where
300 R: GLWECiphertextToMut,
301 K: GLWEPublicKeyPreparedToRef<B>,
302 {
303 self.glwe_encrypt_pk_internal(
304 res,
305 None::<(&GLWEPlaintext<Vec<u8>>, usize)>,
306 pk,
307 source_xu,
308 source_xe,
309 scratch,
310 );
311 }
312}
313
314pub(crate) trait GLWEEncryptPkInternal<B: Backend> {
315 fn glwe_encrypt_pk_internal<R, P, K>(
316 &self,
317 res: &mut R,
318 pt: Option<(&P, usize)>,
319 pk: &K,
320 source_xu: &mut Source,
321 source_xe: &mut Source,
322 scratch: &mut Scratch<B>,
323 ) where
324 R: GLWECiphertextToMut,
325 P: GLWEPlaintextToRef,
326 K: GLWEPublicKeyPreparedToRef<B>;
327}
328
329impl<B: Backend> GLWEEncryptPkInternal<B> for Module<B>
330where
331 Module<B>: SvpPrepare<B>
332 + SvpApplyDftToDft<B>
333 + VecZnxIdftApplyConsume<B>
334 + VecZnxBigAddNormal<B>
335 + VecZnxBigAddSmallInplace<B>
336 + VecZnxBigNormalize<B>,
337 Scratch<B>: TakeSvpPPol<B> + TakeScalarZnx + TakeVecZnxDft<B>,
338{
339 fn glwe_encrypt_pk_internal<R, P, K>(
340 &self,
341 res: &mut R,
342 pt: Option<(&P, usize)>,
343 pk: &K,
344 source_xu: &mut Source,
345 source_xe: &mut Source,
346 scratch: &mut Scratch<B>,
347 ) where
348 R: GLWECiphertextToMut,
349 P: GLWEPlaintextToRef,
350 K: GLWEPublicKeyPreparedToRef<B>,
351 {
352 let res: &mut GLWECiphertext<&mut [u8]> = &mut res.to_mut();
353 let pk: &GLWEPublicKeyPrepared<&[u8], B> = &pk.to_ref();
354
355 #[cfg(debug_assertions)]
356 {
357 assert_eq!(res.base2k(), pk.base2k());
358 assert_eq!(res.n(), pk.n());
359 assert_eq!(res.rank(), pk.rank());
360 if let Some((pt, _)) = pt {
361 assert_eq!(pt.to_ref().base2k(), pk.base2k());
362 assert_eq!(pt.to_ref().n(), pk.n());
363 }
364 }
365
366 let base2k: usize = pk.base2k().into();
367 let size_pk: usize = pk.size();
368 let cols: usize = (res.rank() + 1).into();
369
370 // Generates u according to the underlying secret distribution.
371 let (mut u_dft, scratch_1) = scratch.take_svp_ppol(res.n().into(), 1);
372
373 {
374 let (mut u, _) = scratch_1.take_scalar_znx(res.n().into(), 1);
375 match pk.dist {
376 Distribution::NONE => panic!(
377 "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through \
378 Self::generate"
379 ),
380 Distribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
381 Distribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
382 Distribution::BinaryFixed(hw) => u.fill_binary_hw(0, hw, source_xu),
383 Distribution::BinaryProb(prob) => u.fill_binary_prob(0, prob, source_xu),
384 Distribution::BinaryBlock(block_size) => u.fill_binary_block(0, block_size, source_xu),
385 Distribution::ZERO => {}
386 }
387
388 self.svp_prepare(&mut u_dft, 0, &u, 0);
389 }
390
391 // ct[i] = pk[i] * u + ei (+ m if col = i)
392 (0..cols).for_each(|i| {
393 let (mut ci_dft, scratch_2) = scratch_1.take_vec_znx_dft(res.n().into(), 1, size_pk);
394 // ci_dft = DFT(u) * DFT(pk[i])
395 self.svp_apply_dft_to_dft(&mut ci_dft, 0, &u_dft, 0, &pk.data, i);
396
397 // ci_big = u * p[i]
398 let mut ci_big = self.vec_znx_idft_apply_consume(ci_dft);
399
400 // ci_big = u * pk[i] + e
401 self.vec_znx_big_add_normal(
402 base2k,
403 &mut ci_big,
404 0,
405 pk.k().into(),
406 source_xe,
407 SIGMA,
408 SIGMA_BOUND,
409 );
410
411 // ci_big = u * pk[i] + e + m (if col = i)
412 if let Some((pt, col)) = pt
413 && col == i
414 {
415 self.vec_znx_big_add_small_inplace(&mut ci_big, 0, &pt.to_ref().data, 0);
416 }
417
418 // ct[i] = norm(ci_big)
419 self.vec_znx_big_normalize(base2k, &mut res.data, i, base2k, &ci_big, 0, scratch_2);
420 });
421 }
422}
423
424pub(crate) trait GLWEEncryptSkInternal<B: Backend> {
425 fn glwe_encrypt_sk_internal<R, P, S>(
426 &self,
427 base2k: usize,
428 k: usize,
429 res: &mut R,
430 cols: usize,
431 compressed: bool,
432 pt: Option<(&P, usize)>,
433 sk: &S,
434 source_xa: &mut Source,
435 source_xe: &mut Source,
436 sigma: f64,
437 scratch: &mut Scratch<B>,
438 ) where
439 R: VecZnxToMut,
440 P: GLWEPlaintextToRef,
441 S: GLWESecretPreparedToRef<B>;
442}
443
444impl<B: Backend> GLWEEncryptSkInternal<B> for Module<B>
445where
446 Module<B>: VecZnxDftAllocBytes
447 + VecZnxBigNormalize<B>
448 + VecZnxDftApply<B>
449 + SvpApplyDftToDftInplace<B>
450 + VecZnxIdftApplyConsume<B>
451 + VecZnxNormalizeTmpBytes
452 + VecZnxFillUniform
453 + VecZnxSubInplace
454 + VecZnxAddInplace
455 + VecZnxNormalizeInplace<B>
456 + VecZnxAddNormal
457 + VecZnxNormalize<B>
458 + VecZnxSub,
459 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
460{
461 fn glwe_encrypt_sk_internal<R, P, S>(
462 &self,
463 base2k: usize,
464 k: usize,
465 res: &mut R,
466 cols: usize,
467 compressed: bool,
468 pt: Option<(&P, usize)>,
469 sk: &S,
470 source_xa: &mut Source,
471 source_xe: &mut Source,
472 sigma: f64,
473 scratch: &mut Scratch<B>,
474 ) where
475 R: VecZnxToMut,
476 P: GLWEPlaintextToRef,
477 S: GLWESecretPreparedToRef<B>,
478 {
479 let ct: &mut VecZnx<&mut [u8]> = &mut res.to_mut();
480 let sk: GLWESecretPrepared<&[u8], B> = sk.to_ref();
481
482 #[cfg(debug_assertions)]
483 {
484 if compressed {
485 assert_eq!(
486 ct.cols(),
487 1,
488 "invalid ciphertext: compressed tag=true but #cols={} != 1",
489 ct.cols()
490 )
491 }
492 }
493
494 let size: usize = ct.size();
495
496 let (mut c0, scratch_1) = scratch.take_vec_znx(ct.n(), 1, size);
497 c0.zero();
498
499 {
500 let (mut ci, scratch_2) = scratch_1.take_vec_znx(ct.n(), 1, size);
501
502 // ct[i] = uniform
503 // ct[0] -= c[i] * s[i],
504 (1..cols).for_each(|i| {
505 let col_ct: usize = if compressed { 0 } else { i };
506
507 // ct[i] = uniform (+ pt)
508 self.vec_znx_fill_uniform(base2k, ct, col_ct, source_xa);
509
510 let (mut ci_dft, scratch_3) = scratch_2.take_vec_znx_dft(ct.n(), 1, size);
511
512 // ci = ct[i] - pt
513 // i.e. we act as we sample ct[i] already as uniform + pt
514 // and if there is a pt, then we subtract it before applying DFT
515 if let Some((pt, col)) = pt {
516 if i == col {
517 self.vec_znx_sub(&mut ci, 0, ct, col_ct, &pt.to_ref().data, 0);
518 self.vec_znx_normalize_inplace(base2k, &mut ci, 0, scratch_3);
519 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, &ci, 0);
520 } else {
521 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, ct, col_ct);
522 }
523 } else {
524 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, ct, col_ct);
525 }
526
527 self.svp_apply_dft_to_dft_inplace(&mut ci_dft, 0, &sk.data, i - 1);
528 let ci_big: VecZnxBig<&mut [u8], B> = self.vec_znx_idft_apply_consume(ci_dft);
529
530 // use c[0] as buffer, which is overwritten later by the normalization step
531 self.vec_znx_big_normalize(base2k, &mut ci, 0, base2k, &ci_big, 0, scratch_3);
532
533 // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
534 self.vec_znx_sub_inplace(&mut c0, 0, &ci, 0);
535 });
536 }
537
538 // c[0] += e
539 self.vec_znx_add_normal(base2k, &mut c0, 0, k, source_xe, sigma, SIGMA_BOUND);
540
541 // c[0] += m if col = 0
542 if let Some((pt, col)) = pt
543 && col == 0
544 {
545 self.vec_znx_add_inplace(&mut c0, 0, &pt.to_ref().data, 0);
546 }
547
548 // c[0] = norm(c[0])
549 self.vec_znx_normalize(base2k, ct, 0, base2k, &c0, 0, scratch_1);
550 }
551}