poulpy/poulpy-core/src/encryption/glwe.rs
poulpy-core/src/encryption/glwe.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 GetDistribution, ScratchTakeCore,
14 dist::Distribution,
15 encryption::{SIGMA, SIGMA_BOUND},
16 layouts::{
17 GLWE, GLWEInfos, GLWEPlaintext, GLWEPlaintextToRef, GLWEPrepared, GLWEPreparedToRef, GLWEToMut, LWEInfos,
18 prepared::{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 + GLWEInfos,
83 K: GLWEPreparedToRef<BE> + GetDistribution + GLWEInfos,
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: GLWEPreparedToRef<BE> + GetDistribution + GLWEInfos,
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 + GLWEInfos,
249 K: GLWEPreparedToRef<BE> + GetDistribution + GLWEInfos;
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: GLWEPreparedToRef<BE> + GetDistribution + GLWEInfos;
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 + GLWEInfos,
289 K: GLWEPreparedToRef<BE> + GetDistribution + GLWEInfos,
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: GLWEPreparedToRef<BE> + GetDistribution + GLWEInfos,
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 + GLWEInfos,
328 K: GLWEPreparedToRef<BE> + GetDistribution + GLWEInfos;
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 BE: Backend<ScalarBig = i64>,
343 Scratch<BE>: ScratchTakeBasic,
344{
345 fn glwe_encrypt_pk_internal<R, P, K>(
346 &self,
347 res: &mut R,
348 pt: Option<(&P, usize)>,
349 pk: &K,
350 source_xu: &mut Source,
351 source_xe: &mut Source,
352 scratch: &mut Scratch<BE>,
353 ) where
354 R: GLWEToMut,
355 P: GLWEPlaintextToRef + GLWEInfos,
356 K: GLWEPreparedToRef<BE> + GetDistribution + GLWEInfos,
357 {
358 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
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.base2k(), pk.base2k());
365 assert_eq!(pt.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.n(), 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 {
394 let pk: &GLWEPrepared<&[u8], BE> = &pk.to_ref();
395
396 // ct[i] = pk[i] * u + ei (+ m if col = i)
397 for i in 0..cols {
398 let (mut ci_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, 1, size_pk);
399 // ci_dft = DFT(u) * DFT(pk[i])
400 self.svp_apply_dft_to_dft(&mut ci_dft, 0, &u_dft, 0, &pk.data, i);
401
402 // ci_big = u * p[i]
403 let mut ci_big = self.vec_znx_idft_apply_consume(ci_dft);
404
405 // ci_big = u * pk[i] + e
406 self.vec_znx_big_add_normal(base2k, &mut ci_big, 0, pk.k().into(), source_xe, SIGMA, SIGMA_BOUND);
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(&mut res.data, base2k, 0, i, &ci_big, base2k, 0, scratch_2);
417
418 use poulpy_hal::layouts::ZnxView;
419 for k in 0..res.size(){
420 normalize_check(i, k, ci_big.at(i, k), res.data.at(i, k));
421 }
422 }
423 }
424 }
425}
426
427// DBG
428fn normalize_check(i: usize /*col*/, k: usize, a: &[i64], b: &[i64]) {
429 assert_eq!(a.len(), b.len());
430 // let c_{i,j} be the coefficients of the normalized polynomial, then
431 // i. it is K-normalized if c_{i,j} \in [-2^{K-1}, 2^{K-1}) \forall i=0,\ldots,N-1; j \geq 1
432 // ii. it is reduced if additionally c_{i,0}=0 ~\forall~ i=0, \ldots, N-1
433 if k == 0 {
434 for j in 0..a.len() {
435 assert_eq!(b[j], a[j]);
436 }
437 }
438 // println!("a {:?}", a);
439 // println!("b {:?}", b);
440}
441
442pub(crate) trait GLWEEncryptSkInternal<BE: Backend> {
443 #[allow(clippy::too_many_arguments)]
444 fn glwe_encrypt_sk_internal<R, P, S>(
445 &self,
446 base2k: usize,
447 k: usize,
448 res: &mut R,
449 cols: usize,
450 compressed: bool,
451 pt: Option<(&P, usize)>,
452 sk: &S,
453 source_xa: &mut Source,
454 source_xe: &mut Source,
455 sigma: f64,
456 scratch: &mut Scratch<BE>,
457 ) where
458 R: VecZnxToMut,
459 P: GLWEPlaintextToRef,
460 S: GLWESecretPreparedToRef<BE>;
461}
462
463impl<BE: Backend> GLWEEncryptSkInternal<BE> for Module<BE>
464where
465 Self: ModuleN
466 + VecZnxDftBytesOf
467 + VecZnxBigNormalize<BE>
468 + VecZnxDftApply<BE>
469 + SvpApplyDftToDftInplace<BE>
470 + VecZnxIdftApplyConsume<BE>
471 + VecZnxNormalizeTmpBytes
472 + VecZnxFillUniform
473 + VecZnxSubInplace
474 + VecZnxAddInplace
475 + VecZnxNormalizeInplace<BE>
476 + VecZnxAddNormal
477 + VecZnxNormalize<BE>
478 + VecZnxSub,
479 Scratch<BE>: ScratchTakeBasic,
480{
481 fn glwe_encrypt_sk_internal<R, P, S>(
482 &self,
483 base2k: usize,
484 k: usize,
485 res: &mut R,
486 cols: usize,
487 compressed: bool,
488 pt: Option<(&P, usize)>,
489 sk: &S,
490 source_xa: &mut Source,
491 source_xe: &mut Source,
492 sigma: f64,
493 scratch: &mut Scratch<BE>,
494 ) where
495 R: VecZnxToMut,
496 P: GLWEPlaintextToRef,
497 S: GLWESecretPreparedToRef<BE>,
498 {
499 let ct: &mut VecZnx<&mut [u8]> = &mut res.to_mut();
500 let sk: GLWESecretPrepared<&[u8], BE> = sk.to_ref();
501
502 if compressed {
503 assert_eq!(ct.cols(), 1, "invalid glwe: compressed tag=true but #cols={} != 1", ct.cols())
504 }
505
506 assert!(
507 sk.dist != Distribution::NONE,
508 "glwe secret distribution is NONE (have you prepared the key?)"
509 );
510
511 let size: usize = ct.size();
512
513 let (mut c0, scratch_1) = scratch.take_vec_znx(self.n(), 1, size);
514 c0.zero();
515
516 {
517 let (mut ci, scratch_2) = scratch_1.take_vec_znx(self.n(), 1, size);
518
519 // ct[i] = uniform
520 // ct[0] -= c[i] * s[i],
521 (1..cols).for_each(|i| {
522 let col_ct: usize = if compressed { 0 } else { i };
523
524 // ct[i] = uniform (+ pt)
525 self.vec_znx_fill_uniform(base2k, ct, col_ct, source_xa);
526
527 let (mut ci_dft, scratch_3) = scratch_2.take_vec_znx_dft(self, 1, size);
528
529 // ci = ct[i] - pt
530 // i.e. we act as we sample ct[i] already as uniform + pt
531 // and if there is a pt, then we subtract it before applying DFT
532 if let Some((pt, col)) = pt {
533 if i == col {
534 self.vec_znx_sub(&mut ci, 0, ct, col_ct, &pt.to_ref().data, 0);
535 self.vec_znx_normalize_inplace(base2k, &mut ci, 0, scratch_3);
536 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, &ci, 0);
537 } else {
538 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, ct, col_ct);
539 }
540 } else {
541 self.vec_znx_dft_apply(1, 0, &mut ci_dft, 0, ct, col_ct);
542 }
543
544 self.svp_apply_dft_to_dft_inplace(&mut ci_dft, 0, &sk.data, i - 1);
545 let ci_big: VecZnxBig<&mut [u8], BE> = self.vec_znx_idft_apply_consume(ci_dft);
546
547 // use c[0] as buffer, which is overwritten later by the normalization step
548 self.vec_znx_big_normalize(&mut ci, base2k, 0, 0, &ci_big, base2k, 0, scratch_3);
549
550 // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
551 self.vec_znx_sub_inplace(&mut c0, 0, &ci, 0);
552 });
553 }
554
555 // c[0] += e
556 self.vec_znx_add_normal(base2k, &mut c0, 0, k, source_xe, sigma, SIGMA_BOUND);
557
558 // c[0] += m if col = 0
559 if let Some((pt, col)) = pt
560 && col == 0
561 {
562 self.vec_znx_add_inplace(&mut c0, 0, &pt.to_ref().data, 0);
563 }
564
565 // c[0] = norm(c[0])
566 self.vec_znx_normalize(ct, base2k, 0, 0, &c0, base2k, 0, scratch_1);
567 }
568}