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