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