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 GLWE, GLWEInfos, GLWEPlaintext, GLWEPlaintextToRef, GLWEToMut, LWEInfos,
17 prepared::{GLWEPublicKeyPrepared, GLWEPublicKeyPreparedToRef, GLWESecretPrepared, GLWESecretPreparedToRef},
18 },
19};
20
21impl GLWE<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() + 2 * VecZnx::bytes_of(module.n(), 1, size) + module.vec_znx_dft_bytes_of(1, size)
30 }
31 pub fn encrypt_pk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
32 where
33 A: GLWEInfos,
34 Module<B>: VecZnxDftAllocBytes + SvpPPolAllocBytes + VecZnxBigAllocBytes + VecZnxNormalizeTmpBytes,
35 {
36 let size: usize = infos.size();
37 assert_eq!(module.n() as u32, infos.n());
38 ((module.vec_znx_dft_bytes_of(1, size) + module.vec_znx_big_bytes_of(1, size)) | ScalarZnx::bytes_of(module.n(), 1))
39 + module.svp_ppol_bytes_of(1)
40 + module.vec_znx_normalize_tmp_bytes()
41 }
42}
43
44impl<D: DataMut> GLWE<D> {
45 pub fn encrypt_sk<R, P, S, B: Backend>(
46 &mut self,
47 module: &Module<B>,
48 pt: &P,
49 sk: &S,
50 source_xa: &mut Source,
51 source_xe: &mut Source,
52 scratch: &mut Scratch<B>,
53 ) where
54 P: GLWEPlaintextToRef,
55 S: GLWESecretPreparedToRef<B>,
56 Module<B>: GLWEEncryptSk<B>,
57 {
58 module.glwe_encrypt_sk(self, pt, sk, source_xa, source_xe, scratch);
59 }
60
61 pub fn encrypt_zero_sk<S, B: Backend>(
62 &mut self,
63 module: &Module<B>,
64 sk: &S,
65 source_xa: &mut Source,
66 source_xe: &mut Source,
67 scratch: &mut Scratch<B>,
68 ) where
69 S: GLWESecretPreparedToRef<B>,
70 Module<B>: GLWEEncryptZeroSk<B>,
71 {
72 module.glwe_encrypt_zero_sk(self, sk, source_xa, source_xe, scratch);
73 }
74
75 pub fn encrypt_pk<P, K, B: Backend>(
76 &mut self,
77 module: &Module<B>,
78 pt: &P,
79 pk: &K,
80 source_xu: &mut Source,
81 source_xe: &mut Source,
82 scratch: &mut Scratch<B>,
83 ) where
84 P: GLWEPlaintextToRef,
85 K: GLWEPublicKeyPreparedToRef<B>,
86 Module<B>: GLWEEncryptPk<B>,
87 {
88 module.glwe_encrypt_pk(self, pt, pk, source_xu, source_xe, scratch);
89 }
90
91 pub fn encrypt_zero_pk<K, B: Backend>(
92 &mut self,
93 module: &Module<B>,
94 pk: &K,
95 source_xu: &mut Source,
96 source_xe: &mut Source,
97 scratch: &mut Scratch<B>,
98 ) where
99 K: GLWEPublicKeyPreparedToRef<B>,
100 Module<B>: GLWEEncryptZeroPk<B>,
101 {
102 module.glwe_encrypt_zero_pk(self, pk, source_xu, source_xe, scratch);
103 }
104}
105
106pub trait GLWEEncryptSk<B: 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<B>,
115 ) where
116 R: GLWEToMut,
117 P: GLWEPlaintextToRef,
118 S: GLWESecretPreparedToRef<B>;
119}
120
121impl<B: Backend> GLWEEncryptSk<B> for Module<B>
122where
123 Module<B>: GLWEEncryptSkInternal<B> + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
124 Scratch<B>: 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<B>,
134 ) where
135 R: GLWEToMut,
136 P: GLWEPlaintextToRef,
137 S: GLWESecretPreparedToRef<B>,
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], B> = 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_scratch_space(self, &res),
151 "scratch.available(): {} < GLWECiphertext::encrypt_sk_scratch_space: {}",
152 scratch.available(),
153 GLWE::encrypt_sk_scratch_space(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<B: 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<B>,
182 ) where
183 R: GLWEToMut,
184 S: GLWESecretPreparedToRef<B>;
185}
186
187impl<B: Backend> GLWEEncryptZeroSk<B> for Module<B>
188where
189 Module<B>: GLWEEncryptSkInternal<B> + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes,
190 Scratch<B>: 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<B>,
199 ) where
200 R: GLWEToMut,
201 S: GLWESecretPreparedToRef<B>,
202 {
203 let mut res: GLWE<&mut [u8]> = res.to_mut();
204
205 #[cfg(debug_assertions)]
206 {
207 let sk: GLWESecretPrepared<&[u8], B> = 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_scratch_space(self, &res),
213 "scratch.available(): {} < GLWECiphertext::encrypt_sk_scratch_space: {}",
214 scratch.available(),
215 GLWE::encrypt_sk_scratch_space(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<B: 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<B>,
245 ) where
246 R: GLWEToMut,
247 P: GLWEPlaintextToRef,
248 K: GLWEPublicKeyPreparedToRef<B>;
249}
250
251impl<B: Backend> GLWEEncryptPk<B> for Module<B>
252where
253 Module<B>: GLWEEncryptPkInternal<B>,
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<B>,
263 ) where
264 R: GLWEToMut,
265 P: GLWEPlaintextToRef,
266 K: GLWEPublicKeyPreparedToRef<B>,
267 {
268 self.glwe_encrypt_pk_internal(res, Some((pt, 0)), pk, source_xu, source_xe, scratch);
269 }
270}
271
272pub trait GLWEEncryptZeroPk<B: 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<B>,
280 ) where
281 R: GLWEToMut,
282 K: GLWEPublicKeyPreparedToRef<B>;
283}
284
285impl<B: Backend> GLWEEncryptZeroPk<B> for Module<B>
286where
287 Module<B>: GLWEEncryptPkInternal<B>,
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<B>,
296 ) where
297 R: GLWEToMut,
298 K: GLWEPublicKeyPreparedToRef<B>,
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<B: 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<B>,
320 ) where
321 R: GLWEToMut,
322 P: GLWEPlaintextToRef,
323 K: GLWEPublicKeyPreparedToRef<B>;
324}
325
326impl<B: Backend> GLWEEncryptPkInternal<B> for Module<B>
327where
328 Module<B>: SvpPrepare<B>
329 + SvpApplyDftToDft<B>
330 + VecZnxIdftApplyConsume<B>
331 + VecZnxBigAddNormal<B>
332 + VecZnxBigAddSmallInplace<B>
333 + VecZnxBigNormalize<B>,
334 Scratch<B>: TakeSvpPPol<B> + TakeScalarZnx + TakeVecZnxDft<B>,
335{
336 fn glwe_encrypt_pk_internal<R, P, K>(
337 &self,
338 res: &mut R,
339 pt: Option<(&P, usize)>,
340 pk: &K,
341 source_xu: &mut Source,
342 source_xe: &mut Source,
343 scratch: &mut Scratch<B>,
344 ) where
345 R: GLWEToMut,
346 P: GLWEPlaintextToRef,
347 K: GLWEPublicKeyPreparedToRef<B>,
348 {
349 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
350 let pk: &GLWEPublicKeyPrepared<&[u8], B> = &pk.to_ref();
351
352 #[cfg(debug_assertions)]
353 {
354 assert_eq!(res.base2k(), pk.base2k());
355 assert_eq!(res.n(), pk.n());
356 assert_eq!(res.rank(), pk.rank());
357 if let Some((pt, _)) = pt {
358 assert_eq!(pt.to_ref().base2k(), pk.base2k());
359 assert_eq!(pt.to_ref().n(), pk.n());
360 }
361 }
362
363 let base2k: usize = pk.base2k().into();
364 let size_pk: usize = pk.size();
365 let cols: usize = (res.rank() + 1).into();
366
367 // Generates u according to the underlying secret distribution.
368 let (mut u_dft, scratch_1) = scratch.take_svp_ppol(res.n().into(), 1);
369
370 {
371 let (mut u, _) = scratch_1.take_scalar_znx(res.n().into(), 1);
372 match pk.dist {
373 Distribution::NONE => panic!(
374 "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through \
375 Self::generate"
376 ),
377 Distribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
378 Distribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
379 Distribution::BinaryFixed(hw) => u.fill_binary_hw(0, hw, source_xu),
380 Distribution::BinaryProb(prob) => u.fill_binary_prob(0, prob, source_xu),
381 Distribution::BinaryBlock(block_size) => u.fill_binary_block(0, block_size, source_xu),
382 Distribution::ZERO => {}
383 }
384
385 self.svp_prepare(&mut u_dft, 0, &u, 0);
386 }
387
388 // ct[i] = pk[i] * u + ei (+ m if col = i)
389 (0..cols).for_each(|i| {
390 let (mut ci_dft, scratch_2) = scratch_1.take_vec_znx_dft(res.n().into(), 1, size_pk);
391 // ci_dft = DFT(u) * DFT(pk[i])
392 self.svp_apply_dft_to_dft(&mut ci_dft, 0, &u_dft, 0, &pk.data, i);
393
394 // ci_big = u * p[i]
395 let mut ci_big = self.vec_znx_idft_apply_consume(ci_dft);
396
397 // ci_big = u * pk[i] + e
398 self.vec_znx_big_add_normal(
399 base2k,
400 &mut ci_big,
401 0,
402 pk.k().into(),
403 source_xe,
404 SIGMA,
405 SIGMA_BOUND,
406 );
407
408 // ci_big = u * pk[i] + e + m (if col = i)
409 if let Some((pt, col)) = pt
410 && col == i
411 {
412 self.vec_znx_big_add_small_inplace(&mut ci_big, 0, &pt.to_ref().data, 0);
413 }
414
415 // ct[i] = norm(ci_big)
416 self.vec_znx_big_normalize(base2k, &mut res.data, i, base2k, &ci_big, 0, scratch_2);
417 });
418 }
419}
420
421pub(crate) trait GLWEEncryptSkInternal<B: Backend> {
422 fn glwe_encrypt_sk_internal<R, P, S>(
423 &self,
424 base2k: usize,
425 k: usize,
426 res: &mut R,
427 cols: usize,
428 compressed: bool,
429 pt: Option<(&P, usize)>,
430 sk: &S,
431 source_xa: &mut Source,
432 source_xe: &mut Source,
433 sigma: f64,
434 scratch: &mut Scratch<B>,
435 ) where
436 R: VecZnxToMut,
437 P: GLWEPlaintextToRef,
438 S: GLWESecretPreparedToRef<B>;
439}
440
441impl<B: Backend> GLWEEncryptSkInternal<B> for Module<B>
442where
443 Module<B>: VecZnxDftAllocBytes
444 + VecZnxBigNormalize<B>
445 + VecZnxDftApply<B>
446 + SvpApplyDftToDftInplace<B>
447 + VecZnxIdftApplyConsume<B>
448 + VecZnxNormalizeTmpBytes
449 + VecZnxFillUniform
450 + VecZnxSubInplace
451 + VecZnxAddInplace
452 + VecZnxNormalizeInplace<B>
453 + VecZnxAddNormal
454 + VecZnxNormalize<B>
455 + VecZnxSub,
456 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
457{
458 fn glwe_encrypt_sk_internal<R, P, S>(
459 &self,
460 base2k: usize,
461 k: usize,
462 res: &mut R,
463 cols: usize,
464 compressed: bool,
465 pt: Option<(&P, usize)>,
466 sk: &S,
467 source_xa: &mut Source,
468 source_xe: &mut Source,
469 sigma: f64,
470 scratch: &mut Scratch<B>,
471 ) where
472 R: VecZnxToMut,
473 P: GLWEPlaintextToRef,
474 S: GLWESecretPreparedToRef<B>,
475 {
476 let ct: &mut VecZnx<&mut [u8]> = &mut res.to_mut();
477 let sk: GLWESecretPrepared<&[u8], B> = sk.to_ref();
478
479 #[cfg(debug_assertions)]
480 {
481 if compressed {
482 assert_eq!(
483 ct.cols(),
484 1,
485 "invalid ciphertext: compressed tag=true but #cols={} != 1",
486 ct.cols()
487 )
488 }
489 }
490
491 let size: usize = ct.size();
492
493 let (mut c0, scratch_1) = scratch.take_vec_znx(ct.n(), 1, size);
494 c0.zero();
495
496 {
497 let (mut ci, scratch_2) = scratch_1.take_vec_znx(ct.n(), 1, size);
498
499 // ct[i] = uniform
500 // ct[0] -= c[i] * s[i],
501 (1..cols).for_each(|i| {
502 let col_ct: usize = if compressed { 0 } else { i };
503
504 // ct[i] = uniform (+ pt)
505 self.vec_znx_fill_uniform(base2k, ct, col_ct, source_xa);
506
507 let (mut ci_dft, scratch_3) = scratch_2.take_vec_znx_dft(ct.n(), 1, size);
508
509 // ci = ct[i] - pt
510 // i.e. we act as we sample ct[i] already as uniform + pt
511 // and if there is a pt, then we subtract it before applying DFT
512 if let Some((pt, col)) = pt {
513 if i == col {
514 self.vec_znx_sub(&mut ci, 0, ct, col_ct, &pt.to_ref().data, 0);
515 self.vec_znx_normalize_inplace(base2k, &mut ci, 0, scratch_3);
516 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, &ci, 0);
517 } else {
518 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, ct, col_ct);
519 }
520 } else {
521 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, ct, col_ct);
522 }
523
524 self.svp_apply_dft_to_dft_inplace(&mut ci_dft, 0, &sk.data, i - 1);
525 let ci_big: VecZnxBig<&mut [u8], B> = self.vec_znx_idft_apply_consume(ci_dft);
526
527 // use c[0] as buffer, which is overwritten later by the normalization step
528 self.vec_znx_big_normalize(base2k, &mut ci, 0, base2k, &ci_big, 0, scratch_3);
529
530 // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
531 self.vec_znx_sub_inplace(&mut c0, 0, &ci, 0);
532 });
533 }
534
535 // c[0] += e
536 self.vec_znx_add_normal(base2k, &mut c0, 0, k, source_xe, sigma, SIGMA_BOUND);
537
538 // c[0] += m if col = 0
539 if let Some((pt, col)) = pt
540 && col == 0
541 {
542 self.vec_znx_add_inplace(&mut c0, 0, &pt.to_ref().data, 0);
543 }
544
545 // c[0] = norm(c[0])
546 self.vec_znx_normalize(base2k, ct, 0, base2k, &c0, 0, scratch_1);
547 }
548}