poulpy/core/src/glwe_ciphertext.rs
core/src/glwe_ciphertext.rsRaw
1use base2k::{
2 AddNormal, Backend, FFT64, FillUniform, MatZnxDft, MatZnxDftOps, MatZnxDftScratch, MatZnxDftToRef, Module, ScalarZnxAlloc,
3 ScalarZnxDft, ScalarZnxDftAlloc, ScalarZnxDftOps, ScalarZnxDftToRef, Scratch, VecZnx, VecZnxAlloc, VecZnxBig, VecZnxBigAlloc,
4 VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps,
5 VecZnxToMut, VecZnxToRef, ZnxInfos, ZnxZero,
6};
7use sampling::source::Source;
8
9use crate::{
10 SIX_SIGMA,
11 elem::Infos,
12 ggsw_ciphertext::GGSWCiphertext,
13 glwe_ciphertext_fourier::GLWECiphertextFourier,
14 glwe_plaintext::GLWEPlaintext,
15 keys::{GLWEPublicKey, SecretDistribution, SecretKeyFourier},
16 keyswitch_key::GLWESwitchingKey,
17 utils::derive_size,
18};
19
20pub struct GLWECiphertext<C> {
21 pub data: VecZnx<C>,
22 pub basek: usize,
23 pub k: usize,
24}
25
26impl GLWECiphertext<Vec<u8>> {
27 pub fn new<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
28 Self {
29 data: module.new_vec_znx(rank + 1, derive_size(basek, k)),
30 basek,
31 k,
32 }
33 }
34}
35
36impl<T> Infos for GLWECiphertext<T> {
37 type Inner = VecZnx<T>;
38
39 fn inner(&self) -> &Self::Inner {
40 &self.data
41 }
42
43 fn basek(&self) -> usize {
44 self.basek
45 }
46
47 fn k(&self) -> usize {
48 self.k
49 }
50}
51
52impl<T> GLWECiphertext<T> {
53 pub fn rank(&self) -> usize {
54 self.cols() - 1
55 }
56}
57
58impl<C> VecZnxToMut for GLWECiphertext<C>
59where
60 VecZnx<C>: VecZnxToMut,
61{
62 fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
63 self.data.to_mut()
64 }
65}
66
67impl<C> VecZnxToRef for GLWECiphertext<C>
68where
69 VecZnx<C>: VecZnxToRef,
70{
71 fn to_ref(&self) -> VecZnx<&[u8]> {
72 self.data.to_ref()
73 }
74}
75
76impl<C> GLWECiphertext<C>
77where
78 VecZnx<C>: VecZnxToRef,
79{
80 #[allow(dead_code)]
81 pub(crate) fn dft<R>(&self, module: &Module<FFT64>, res: &mut GLWECiphertextFourier<R, FFT64>)
82 where
83 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + ZnxInfos,
84 {
85 #[cfg(debug_assertions)]
86 {
87 assert_eq!(self.rank(), res.rank());
88 assert_eq!(self.basek(), res.basek())
89 }
90
91 (0..self.rank() + 1).for_each(|i| {
92 module.vec_znx_dft(res, i, self, i);
93 })
94 }
95}
96
97impl GLWECiphertext<Vec<u8>> {
98 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, ct_size: usize) -> usize {
99 module.vec_znx_big_normalize_tmp_bytes()
100 + module.bytes_of_vec_znx_dft(1, ct_size)
101 + module.bytes_of_vec_znx_big(1, ct_size)
102 }
103 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, pk_size: usize) -> usize {
104 ((module.bytes_of_vec_znx_dft(1, pk_size) + module.bytes_of_vec_znx_big(1, pk_size)) | module.bytes_of_scalar_znx(1))
105 + module.bytes_of_scalar_znx_dft(1)
106 + module.vec_znx_big_normalize_tmp_bytes()
107 }
108
109 pub fn decrypt_scratch_space(module: &Module<FFT64>, ct_size: usize) -> usize {
110 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, ct_size))
111 + module.bytes_of_vec_znx_big(1, ct_size)
112 }
113
114 pub fn keyswitch_scratch_space(
115 module: &Module<FFT64>,
116 out_size: usize,
117 out_rank: usize,
118 in_size: usize,
119 in_rank: usize,
120 ksk_size: usize,
121 ) -> usize {
122 let res_dft: usize = module.bytes_of_vec_znx_dft(out_rank + 1, ksk_size);
123 let vmp: usize = module.vmp_apply_tmp_bytes(
124 out_size,
125 in_size,
126 in_size,
127 in_rank + 1,
128 out_rank + 1,
129 ksk_size,
130 ) + module.bytes_of_vec_znx_dft(in_rank, in_size);
131 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
132
133 return res_dft + (vmp | normalize);
134 }
135
136 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
137 GLWECiphertext::keyswitch_scratch_space(module, out_size, out_rank, out_size, out_rank, ksk_size)
138 }
139
140 pub fn external_product_scratch_space(
141 module: &Module<FFT64>,
142 out_size: usize,
143 in_size: usize,
144 ggsw_size: usize,
145 rank: usize,
146 ) -> usize {
147 let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, ggsw_size);
148 let vmp: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size)
149 + module.vmp_apply_tmp_bytes(
150 out_size,
151 in_size,
152 in_size, // rows
153 rank + 1, // cols in
154 rank + 1, // cols out
155 ggsw_size,
156 );
157 let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
158
159 res_dft + (vmp | normalize)
160 }
161
162 pub fn external_product_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize, rank: usize) -> usize {
163 GLWECiphertext::external_product_scratch_space(module, res_size, res_size, rhs, rank)
164 }
165}
166
167impl<DataSelf> GLWECiphertext<DataSelf>
168where
169 VecZnx<DataSelf>: VecZnxToMut + VecZnxToRef,
170{
171 pub fn encrypt_sk<DataPt, DataSk>(
172 &mut self,
173 module: &Module<FFT64>,
174 pt: &GLWEPlaintext<DataPt>,
175 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
176 source_xa: &mut Source,
177 source_xe: &mut Source,
178 sigma: f64,
179 scratch: &mut Scratch,
180 ) where
181 VecZnx<DataPt>: VecZnxToRef,
182 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
183 {
184 self.encrypt_sk_private(
185 module,
186 Some((pt, 0)),
187 sk_dft,
188 source_xa,
189 source_xe,
190 sigma,
191 scratch,
192 );
193 }
194
195 pub fn encrypt_zero_sk<DataSk>(
196 &mut self,
197 module: &Module<FFT64>,
198 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
199 source_xa: &mut Source,
200 source_xe: &mut Source,
201 sigma: f64,
202 scratch: &mut Scratch,
203 ) where
204 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
205 {
206 self.encrypt_sk_private(module, None, sk_dft, source_xa, source_xe, sigma, scratch);
207 }
208
209 pub fn encrypt_pk<DataPt, DataPk>(
210 &mut self,
211 module: &Module<FFT64>,
212 pt: &GLWEPlaintext<DataPt>,
213 pk: &GLWEPublicKey<DataPk, FFT64>,
214 source_xu: &mut Source,
215 source_xe: &mut Source,
216 sigma: f64,
217 scratch: &mut Scratch,
218 ) where
219 VecZnx<DataPt>: VecZnxToRef,
220 VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
221 {
222 self.encrypt_pk_private(
223 module,
224 Some((pt, 0)),
225 pk,
226 source_xu,
227 source_xe,
228 sigma,
229 scratch,
230 );
231 }
232
233 pub fn encrypt_zero_pk<DataPk>(
234 &mut self,
235 module: &Module<FFT64>,
236 pk: &GLWEPublicKey<DataPk, FFT64>,
237 source_xu: &mut Source,
238 source_xe: &mut Source,
239 sigma: f64,
240 scratch: &mut Scratch,
241 ) where
242 VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
243 {
244 self.encrypt_pk_private(module, None, pk, source_xu, source_xe, sigma, scratch);
245 }
246
247 pub fn keyswitch<DataLhs, DataRhs>(
248 &mut self,
249 module: &Module<FFT64>,
250 lhs: &GLWECiphertext<DataLhs>,
251 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
252 scratch: &mut Scratch,
253 ) where
254 VecZnx<DataLhs>: VecZnxToRef,
255 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
256 {
257 let basek: usize = self.basek();
258
259 #[cfg(debug_assertions)]
260 {
261 assert_eq!(lhs.rank(), rhs.rank_in());
262 assert_eq!(self.rank(), rhs.rank_out());
263 assert_eq!(self.basek(), basek);
264 assert_eq!(lhs.basek(), basek);
265 assert_eq!(rhs.n(), module.n());
266 assert_eq!(self.n(), module.n());
267 assert_eq!(lhs.n(), module.n());
268 assert!(
269 scratch.available()
270 >= GLWECiphertext::keyswitch_scratch_space(
271 module,
272 self.size(),
273 self.rank(),
274 lhs.size(),
275 lhs.rank(),
276 rhs.size(),
277 )
278 );
279 }
280
281 let cols_in: usize = rhs.rank_in();
282 let cols_out: usize = rhs.rank_out() + 1;
283
284 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
285
286 {
287 let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, lhs.size());
288 (0..cols_in).for_each(|col_i| {
289 module.vec_znx_dft(&mut ai_dft, col_i, lhs, col_i + 1);
290 });
291 module.vmp_apply(&mut res_dft, &ai_dft, rhs, scratch2);
292 }
293
294 let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
295
296 module.vec_znx_big_add_small_inplace(&mut res_big, 0, lhs, 0);
297
298 (0..cols_out).for_each(|i| {
299 module.vec_znx_big_normalize(basek, self, i, &res_big, i, scratch1);
300 });
301 }
302
303 pub fn keyswitch_inplace<DataRhs>(
304 &mut self,
305 module: &Module<FFT64>,
306 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
307 scratch: &mut Scratch,
308 ) where
309 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
310 {
311 unsafe {
312 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
313 self.keyswitch(&module, &*self_ptr, rhs, scratch);
314 }
315 }
316
317 pub fn external_product<DataLhs, DataRhs>(
318 &mut self,
319 module: &Module<FFT64>,
320 lhs: &GLWECiphertext<DataLhs>,
321 rhs: &GGSWCiphertext<DataRhs, FFT64>,
322 scratch: &mut Scratch,
323 ) where
324 VecZnx<DataLhs>: VecZnxToRef,
325 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
326 {
327 let basek: usize = self.basek();
328
329 #[cfg(debug_assertions)]
330 {
331 assert_eq!(rhs.rank(), lhs.rank());
332 assert_eq!(rhs.rank(), self.rank());
333 assert_eq!(self.basek(), basek);
334 assert_eq!(lhs.basek(), basek);
335 assert_eq!(rhs.n(), module.n());
336 assert_eq!(self.n(), module.n());
337 assert_eq!(lhs.n(), module.n());
338 }
339
340 let cols: usize = rhs.rank() + 1;
341
342 let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size()); // Todo optimise
343
344 {
345 let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, lhs.size());
346 (0..cols).for_each(|col_i| {
347 module.vec_znx_dft(&mut a_dft, col_i, lhs, col_i);
348 });
349 module.vmp_apply(&mut res_dft, &a_dft, rhs, scratch2);
350 }
351
352 let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
353
354 (0..cols).for_each(|i| {
355 module.vec_znx_big_normalize(basek, self, i, &res_big, i, scratch1);
356 });
357 }
358
359 pub fn external_product_inplace<DataRhs>(
360 &mut self,
361 module: &Module<FFT64>,
362 rhs: &GGSWCiphertext<DataRhs, FFT64>,
363 scratch: &mut Scratch,
364 ) where
365 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
366 {
367 unsafe {
368 let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
369 self.external_product(&module, &*self_ptr, rhs, scratch);
370 }
371 }
372
373 pub(crate) fn encrypt_sk_private<DataPt, DataSk>(
374 &mut self,
375 module: &Module<FFT64>,
376 pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
377 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
378 source_xa: &mut Source,
379 source_xe: &mut Source,
380 sigma: f64,
381 scratch: &mut Scratch,
382 ) where
383 VecZnx<DataPt>: VecZnxToRef,
384 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
385 {
386 #[cfg(debug_assertions)]
387 {
388 assert_eq!(self.rank(), sk_dft.rank());
389 assert_eq!(sk_dft.n(), module.n());
390 assert_eq!(self.n(), module.n());
391 if let Some((pt, col)) = pt {
392 assert_eq!(pt.n(), module.n());
393 assert!(col < self.rank() + 1);
394 }
395 }
396
397 let log_base2k: usize = self.basek();
398 let log_k: usize = self.k();
399 let size: usize = self.size();
400 let cols: usize = self.rank() + 1;
401
402 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
403 c0_big.zero();
404
405 {
406 // c[i] = uniform
407 // c[0] -= c[i] * s[i],
408 (1..cols).for_each(|i| {
409 let (mut ci_dft, scratch_2) = scratch_1.tmp_vec_znx_dft(module, 1, size);
410
411 // c[i] = uniform
412 self.data.fill_uniform(log_base2k, i, size, source_xa);
413
414 // c[i] = norm(IDFT(DFT(c[i]) * DFT(s[i])))
415 module.vec_znx_dft(&mut ci_dft, 0, self, i);
416 module.svp_apply_inplace(&mut ci_dft, 0, sk_dft, i - 1);
417 let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
418
419 // use c[0] as buffer, which is overwritten later by the normalization step
420 module.vec_znx_big_normalize(log_base2k, self, 0, &ci_big, 0, scratch_2);
421
422 // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
423 module.vec_znx_sub_ab_inplace(&mut c0_big, 0, self, 0);
424
425 // c[i] += m if col = i
426 if let Some((pt, col)) = pt {
427 if i == col {
428 module.vec_znx_add_inplace(self, i, pt, 0);
429 module.vec_znx_normalize_inplace(log_base2k, self, i, scratch_2);
430 }
431 }
432 });
433 }
434
435 // c[0] += e
436 c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, sigma * SIX_SIGMA);
437
438 // c[0] += m if col = 0
439 if let Some((pt, col)) = pt {
440 if col == 0 {
441 module.vec_znx_add_inplace(&mut c0_big, 0, pt, 0);
442 }
443 }
444
445 // c[0] = norm(c[0])
446 module.vec_znx_normalize(log_base2k, self, 0, &c0_big, 0, scratch_1);
447 }
448
449 pub(crate) fn encrypt_pk_private<DataPt, DataPk>(
450 &mut self,
451 module: &Module<FFT64>,
452 pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
453 pk: &GLWEPublicKey<DataPk, FFT64>,
454 source_xu: &mut Source,
455 source_xe: &mut Source,
456 sigma: f64,
457 scratch: &mut Scratch,
458 ) where
459 VecZnx<DataPt>: VecZnxToRef,
460 VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
461 {
462 #[cfg(debug_assertions)]
463 {
464 assert_eq!(self.basek(), pk.basek());
465 assert_eq!(self.n(), module.n());
466 assert_eq!(pk.n(), module.n());
467 assert_eq!(self.rank(), pk.rank());
468 if let Some((pt, _)) = pt {
469 assert_eq!(pt.basek(), pk.basek());
470 assert_eq!(pt.n(), module.n());
471 }
472 }
473
474 let log_base2k: usize = pk.basek();
475 let size_pk: usize = pk.size();
476 let cols: usize = self.rank() + 1;
477
478 // Generates u according to the underlying secret distribution.
479 let (mut u_dft, scratch_1) = scratch.tmp_scalar_znx_dft(module, 1);
480
481 {
482 let (mut u, _) = scratch_1.tmp_scalar_znx(module, 1);
483 match pk.dist {
484 SecretDistribution::NONE => panic!(
485 "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through \
486 Self::generate"
487 ),
488 SecretDistribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
489 SecretDistribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
490 SecretDistribution::ZERO => {}
491 }
492
493 module.svp_prepare(&mut u_dft, 0, &u, 0);
494 }
495
496 // ct[i] = pk[i] * u + ei (+ m if col = i)
497 (0..cols).for_each(|i| {
498 let (mut ci_dft, scratch_2) = scratch_1.tmp_vec_znx_dft(module, 1, size_pk);
499 // ci_dft = DFT(u) * DFT(pk[i])
500 module.svp_apply(&mut ci_dft, 0, &u_dft, 0, pk, i);
501
502 // ci_big = u * p[i]
503 let mut ci_big = module.vec_znx_idft_consume(ci_dft);
504
505 // ci_big = u * pk[i] + e
506 ci_big.add_normal(log_base2k, 0, pk.k(), source_xe, sigma, sigma * SIX_SIGMA);
507
508 // ci_big = u * pk[i] + e + m (if col = i)
509 if let Some((pt, col)) = pt {
510 if col == i {
511 module.vec_znx_big_add_small_inplace(&mut ci_big, 0, pt, 0);
512 }
513 }
514
515 // ct[i] = norm(ci_big)
516 module.vec_znx_big_normalize(log_base2k, self, i, &ci_big, 0, scratch_2);
517 });
518 }
519}
520
521impl<DataSelf> GLWECiphertext<DataSelf>
522where
523 VecZnx<DataSelf>: VecZnxToRef,
524{
525 pub fn decrypt<DataPt, DataSk>(
526 &self,
527 module: &Module<FFT64>,
528 pt: &mut GLWEPlaintext<DataPt>,
529 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
530 scratch: &mut Scratch,
531 ) where
532 VecZnx<DataPt>: VecZnxToMut,
533 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
534 {
535 #[cfg(debug_assertions)]
536 {
537 assert_eq!(self.rank(), sk_dft.rank());
538 assert_eq!(self.n(), module.n());
539 assert_eq!(pt.n(), module.n());
540 assert_eq!(sk_dft.n(), module.n());
541 }
542
543 let cols: usize = self.rank() + 1;
544
545 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
546 c0_big.zero();
547
548 {
549 (1..cols).for_each(|i| {
550 // ci_dft = DFT(a[i]) * DFT(s[i])
551 let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
552 module.vec_znx_dft(&mut ci_dft, 0, self, i);
553 module.svp_apply_inplace(&mut ci_dft, 0, sk_dft, i - 1);
554 let ci_big = module.vec_znx_idft_consume(ci_dft);
555
556 // c0_big += a[i] * s[i]
557 module.vec_znx_big_add_inplace(&mut c0_big, 0, &ci_big, 0);
558 });
559 }
560
561 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
562 module.vec_znx_big_add_small_inplace(&mut c0_big, 0, self, 0);
563
564 // pt = norm(BIG(m + e))
565 module.vec_znx_big_normalize(self.basek(), pt, 0, &mut c0_big, 0, scratch_1);
566
567 pt.basek = self.basek();
568 pt.k = pt.k().min(self.k());
569 }
570}