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 pub fn encrypt_sk_tmp_bytes<BE: Backend, A>(module: &Module<BE>, infos: &A) -> usize
23 where
24 A: GLWEInfos,
25 Module<BE>: VecZnxNormalizeTmpBytes + VecZnxDftBytesOf,
26 {
27 let size: usize = infos.size();
28 assert_eq!(module.n() as u32, infos.n());
29 module.vec_znx_normalize_tmp_bytes() + 2 * VecZnx::bytes_of(module.n(), 1, size) + module.bytes_of_vec_znx_dft(1, size)
30 }
31 pub fn encrypt_pk_tmp_bytes<BE: Backend, A>(module: &Module<BE>, infos: &A) -> usize
32 where
33 A: GLWEInfos,
34 Module<BE>: VecZnxDftBytesOf + SvpPPolBytesOf + VecZnxBigBytesOf + VecZnxNormalizeTmpBytes,
35 {
36 let size: usize = infos.size();
37 assert_eq!(module.n() as u32, infos.n());
38 ((module.bytes_of_vec_znx_dft(1, size) + module.bytes_of_vec_znx_big(1, size)) | ScalarZnx::bytes_of(module.n(), 1))
39 + module.bytes_of_svp_ppol(1)
40 + module.vec_znx_normalize_tmp_bytes()
41 }
42}
43
44impl<D: DataMut> GLWE<D> {
45 pub fn encrypt_sk<R, P, S, BE: Backend>(
46 &mut self,
47 module: &Module<BE>,
48 pt: &P,
49 sk: &S,
50 source_xa: &mut Source,
51 source_xe: &mut Source,
52 scratch: &mut Scratch<BE>,
53 ) where
54 P: GLWEPlaintextToRef,
55 S: GLWESecretPreparedToRef<BE>,
56 Module<BE>: GLWEEncryptSk<BE>,
57 {
58 module.glwe_encrypt_sk(self, pt, sk, source_xa, source_xe, scratch);
59 }
60
61 pub fn encrypt_zero_sk<S, BE: Backend>(
62 &mut self,
63 module: &Module<BE>,
64 sk: &S,
65 source_xa: &mut Source,
66 source_xe: &mut Source,
67 scratch: &mut Scratch<BE>,
68 ) where
69 S: GLWESecretPreparedToRef<BE>,
70 Module<BE>: GLWEEncryptZeroSk<BE>,
71 {
72 module.glwe_encrypt_zero_sk(self, sk, source_xa, source_xe, scratch);
73 }
74
75 pub fn encrypt_pk<P, K, BE: Backend>(
76 &mut self,
77 module: &Module<BE>,
78 pt: &P,
79 pk: &K,
80 source_xu: &mut Source,
81 source_xe: &mut Source,
82 scratch: &mut Scratch<BE>,
83 ) where
84 P: GLWEPlaintextToRef,
85 K: GLWEPublicKeyPreparedToRef<BE>,
86 Module<BE>: GLWEEncryptPk<BE>,
87 {
88 module.glwe_encrypt_pk(self, pt, pk, source_xu, source_xe, scratch);
89 }
90
91 pub fn encrypt_zero_pk<K, BE: Backend>(
92 &mut self,
93 module: &Module<BE>,
94 pk: &K,
95 source_xu: &mut Source,
96 source_xe: &mut Source,
97 scratch: &mut Scratch<BE>,
98 ) where
99 K: GLWEPublicKeyPreparedToRef<BE>,
100 Module<BE>: GLWEEncryptZeroPk<BE>,
101 {
102 module.glwe_encrypt_zero_pk(self, pk, source_xu, source_xe, scratch);
103 }
104}
105
106pub trait GLWEEncryptSk<BE: Backend> {
107 fn glwe_encrypt_sk<R, P, S>(
108 &self,
109 res: &mut R,
110 pt: &P,
111 sk: &S,
112 source_xa: &mut Source,
113 source_xe: &mut Source,
114 scratch: &mut Scratch<BE>,
115 ) where
116 R: GLWEToMut,
117 P: GLWEPlaintextToRef,
118 S: GLWESecretPreparedToRef<BE>;
119}
120
121impl<BE: Backend> GLWEEncryptSk<BE> for Module<BE>
122where
123 Module<BE>: GLWEEncryptSkInternal<BE> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf,
124 Scratch<BE>: ScratchAvailable,
125{
126 fn glwe_encrypt_sk<R, P, S>(
127 &self,
128 res: &mut R,
129 pt: &P,
130 sk: &S,
131 source_xa: &mut Source,
132 source_xe: &mut Source,
133 scratch: &mut Scratch<BE>,
134 ) where
135 R: GLWEToMut,
136 P: GLWEPlaintextToRef,
137 S: GLWESecretPreparedToRef<BE>,
138 {
139 let mut res: GLWE<&mut [u8]> = res.to_mut();
140 let pt: GLWEPlaintext<&[u8]> = pt.to_ref();
141
142 #[cfg(debug_assertions)]
143 {
144 let sk: GLWESecretPrepared<&[u8], BE> = sk.to_ref();
145 assert_eq!(res.rank(), sk.rank());
146 assert_eq!(res.n(), self.n() as u32);
147 assert_eq!(sk.n(), self.n() as u32);
148 assert_eq!(pt.n(), self.n() as u32);
149 assert!(
150 scratch.available() >= GLWE::encrypt_sk_tmp_bytes(self, &res),
151 "scratch.available(): {} < GLWECiphertext::encrypt_sk_tmp_bytes: {}",
152 scratch.available(),
153 GLWE::encrypt_sk_tmp_bytes(self, &res)
154 )
155 }
156
157 let cols: usize = (res.rank() + 1).into();
158 self.glwe_encrypt_sk_internal(
159 res.base2k().into(),
160 res.k().into(),
161 res.data_mut(),
162 cols,
163 false,
164 Some((&pt, 0)),
165 sk,
166 source_xa,
167 source_xe,
168 SIGMA,
169 scratch,
170 );
171 }
172}
173
174pub trait GLWEEncryptZeroSk<BE: Backend> {
175 fn glwe_encrypt_zero_sk<R, S>(
176 &self,
177 res: &mut R,
178 sk: &S,
179 source_xa: &mut Source,
180 source_xe: &mut Source,
181 scratch: &mut Scratch<BE>,
182 ) where
183 R: GLWEToMut,
184 S: GLWESecretPreparedToRef<BE>;
185}
186
187impl<BE: Backend> GLWEEncryptZeroSk<BE> for Module<BE>
188where
189 Module<BE>: GLWEEncryptSkInternal<BE> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf,
190 Scratch<BE>: ScratchAvailable,
191{
192 fn glwe_encrypt_zero_sk<R, S>(
193 &self,
194 res: &mut R,
195 sk: &S,
196 source_xa: &mut Source,
197 source_xe: &mut Source,
198 scratch: &mut Scratch<BE>,
199 ) where
200 R: GLWEToMut,
201 S: GLWESecretPreparedToRef<BE>,
202 {
203 let mut res: GLWE<&mut [u8]> = res.to_mut();
204
205 #[cfg(debug_assertions)]
206 {
207 let sk: GLWESecretPrepared<&[u8], BE> = sk.to_ref();
208 assert_eq!(res.rank(), sk.rank());
209 assert_eq!(res.n(), self.n() as u32);
210 assert_eq!(sk.n(), self.n() as u32);
211 assert!(
212 scratch.available() >= GLWE::encrypt_sk_tmp_bytes(self, &res),
213 "scratch.available(): {} < GLWECiphertext::encrypt_sk_tmp_bytes: {}",
214 scratch.available(),
215 GLWE::encrypt_sk_tmp_bytes(self, &res)
216 )
217 }
218
219 let cols: usize = (res.rank() + 1).into();
220 self.glwe_encrypt_sk_internal(
221 res.base2k().into(),
222 res.k().into(),
223 res.data_mut(),
224 cols,
225 false,
226 None::<(&GLWEPlaintext<Vec<u8>>, usize)>,
227 sk,
228 source_xa,
229 source_xe,
230 SIGMA,
231 scratch,
232 );
233 }
234}
235
236pub trait GLWEEncryptPk<BE: Backend> {
237 fn glwe_encrypt_pk<R, P, K>(
238 &self,
239 res: &mut R,
240 pt: &P,
241 pk: &K,
242 source_xu: &mut Source,
243 source_xe: &mut Source,
244 scratch: &mut Scratch<BE>,
245 ) where
246 R: GLWEToMut,
247 P: GLWEPlaintextToRef,
248 K: GLWEPublicKeyPreparedToRef<BE>;
249}
250
251impl<BE: Backend> GLWEEncryptPk<BE> for Module<BE>
252where
253 Module<BE>: GLWEEncryptPkInternal<BE>,
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 self.glwe_encrypt_pk_internal(res, Some((pt, 0)), pk, source_xu, source_xe, scratch);
269 }
270}
271
272pub trait GLWEEncryptZeroPk<BE: Backend> {
273 fn glwe_encrypt_zero_pk<R, K>(
274 &self,
275 res: &mut R,
276 pk: &K,
277 source_xu: &mut Source,
278 source_xe: &mut Source,
279 scratch: &mut Scratch<BE>,
280 ) where
281 R: GLWEToMut,
282 K: GLWEPublicKeyPreparedToRef<BE>;
283}
284
285impl<BE: Backend> GLWEEncryptZeroPk<BE> for Module<BE>
286where
287 Module<BE>: GLWEEncryptPkInternal<BE>,
288{
289 fn glwe_encrypt_zero_pk<R, K>(
290 &self,
291 res: &mut R,
292 pk: &K,
293 source_xu: &mut Source,
294 source_xe: &mut Source,
295 scratch: &mut Scratch<BE>,
296 ) where
297 R: GLWEToMut,
298 K: GLWEPublicKeyPreparedToRef<BE>,
299 {
300 self.glwe_encrypt_pk_internal(
301 res,
302 None::<(&GLWEPlaintext<Vec<u8>>, usize)>,
303 pk,
304 source_xu,
305 source_xe,
306 scratch,
307 );
308 }
309}
310
311pub(crate) trait GLWEEncryptPkInternal<BE: Backend> {
312 fn glwe_encrypt_pk_internal<R, P, K>(
313 &self,
314 res: &mut R,
315 pt: Option<(&P, usize)>,
316 pk: &K,
317 source_xu: &mut Source,
318 source_xe: &mut Source,
319 scratch: &mut Scratch<BE>,
320 ) where
321 R: GLWEToMut,
322 P: GLWEPlaintextToRef,
323 K: GLWEPublicKeyPreparedToRef<BE>;
324}
325
326impl<BE: Backend> GLWEEncryptPkInternal<BE> for Module<BE>
327where
328 Module<BE>: SvpPrepare<BE>
329 + SvpApplyDftToDft<BE>
330 + VecZnxIdftApplyConsume<BE>
331 + VecZnxBigAddNormal<BE>
332 + VecZnxBigAddSmallInplace<BE>
333 + VecZnxBigNormalize<BE>
334 + SvpPPolBytesOf
335 + ModuleN
336 + VecZnxDftBytesOf,
337 Scratch<BE>: ScratchTakeBasic,
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<BE>,
347 ) where
348 R: GLWEToMut,
349 P: GLWEPlaintextToRef,
350 K: GLWEPublicKeyPreparedToRef<BE>,
351 {
352 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
353 let pk: &GLWEPublicKeyPrepared<&[u8], BE> = &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(self, 1);
372
373 {
374 let (mut u, _) = scratch_1.take_scalar_znx(self, 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(self, 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<BE: 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<BE>,
438 ) where
439 R: VecZnxToMut,
440 P: GLWEPlaintextToRef,
441 S: GLWESecretPreparedToRef<BE>;
442}
443
444impl<BE: Backend> GLWEEncryptSkInternal<BE> for Module<BE>
445where
446 Module<BE>: ModuleN
447 + VecZnxDftBytesOf
448 + VecZnxBigNormalize<BE>
449 + VecZnxDftApply<BE>
450 + SvpApplyDftToDftInplace<BE>
451 + VecZnxIdftApplyConsume<BE>
452 + VecZnxNormalizeTmpBytes
453 + VecZnxFillUniform
454 + VecZnxSubInplace
455 + VecZnxAddInplace
456 + VecZnxNormalizeInplace<BE>
457 + VecZnxAddNormal
458 + VecZnxNormalize<BE>
459 + VecZnxSub,
460 Scratch<BE>: ScratchAvailable + ScratchTakeBasic,
461{
462 fn glwe_encrypt_sk_internal<R, P, S>(
463 &self,
464 base2k: usize,
465 k: usize,
466 res: &mut R,
467 cols: usize,
468 compressed: bool,
469 pt: Option<(&P, usize)>,
470 sk: &S,
471 source_xa: &mut Source,
472 source_xe: &mut Source,
473 sigma: f64,
474 scratch: &mut Scratch<BE>,
475 ) where
476 R: VecZnxToMut,
477 P: GLWEPlaintextToRef,
478 S: GLWESecretPreparedToRef<BE>,
479 {
480 let ct: &mut VecZnx<&mut [u8]> = &mut res.to_mut();
481 let sk: GLWESecretPrepared<&[u8], BE> = sk.to_ref();
482
483 #[cfg(debug_assertions)]
484 {
485 if compressed {
486 assert_eq!(
487 ct.cols(),
488 1,
489 "invalid ciphertext: compressed tag=true but #cols={} != 1",
490 ct.cols()
491 )
492 }
493 }
494
495 let size: usize = ct.size();
496
497 let (mut c0, scratch_1) = scratch.take_vec_znx(self, 1, size);
498 c0.zero();
499
500 {
501 let (mut ci, scratch_2) = scratch_1.take_vec_znx(self, 1, size);
502
503 // ct[i] = uniform
504 // ct[0] -= c[i] * s[i],
505 (1..cols).for_each(|i| {
506 let col_ct: usize = if compressed { 0 } else { i };
507
508 // ct[i] = uniform (+ pt)
509 self.vec_znx_fill_uniform(base2k, ct, col_ct, source_xa);
510
511 let (mut ci_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, size);
512
513 // ci = ct[i] - pt
514 // i.e. we act as we sample ct[i] already as uniform + pt
515 // and if there is a pt, then we subtract it before applying DFT
516 if let Some((pt, col)) = pt {
517 if i == col {
518 self.vec_znx_sub(&mut ci, 0, ct, col_ct, &pt.to_ref().data, 0);
519 self.vec_znx_normalize_inplace(base2k, &mut ci, 0, scratch_3);
520 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, &ci, 0);
521 } else {
522 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, ct, col_ct);
523 }
524 } else {
525 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, ct, col_ct);
526 }
527
528 self.svp_apply_dft_to_dft_inplace(&mut ci_dft, 0, &sk.data, i - 1);
529 let ci_big: VecZnxBig<&mut [u8], BE> = self.vec_znx_idft_apply_consume(ci_dft);
530
531 // use c[0] as buffer, which is overwritten later by the normalization step
532 self.vec_znx_big_normalize(base2k, &mut ci, 0, base2k, &ci_big, 0, scratch_3);
533
534 // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
535 self.vec_znx_sub_inplace(&mut c0, 0, &ci, 0);
536 });
537 }
538
539 // c[0] += e
540 self.vec_znx_add_normal(base2k, &mut c0, 0, k, source_xe, sigma, SIGMA_BOUND);
541
542 // c[0] += m if col = 0
543 if let Some((pt, col)) = pt
544 && col == 0
545 {
546 self.vec_znx_add_inplace(&mut c0, 0, &pt.to_ref().data, 0);
547 }
548
549 // c[0] = norm(c[0])
550 self.vec_znx_normalize(base2k, ct, 0, base2k, &c0, 0, scratch_1);
551 }
552}