poulpy/core/src/rlwe.rs
core/src/rlwe.rsRaw
1use base2k::{
2 AddNormal, Backend, FFT64, FillUniform, MatZnxDft, MatZnxDftToRef, Module, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc,
3 ScalarZnxDftOps, ScalarZnxDftToRef, Scratch, VecZnx, VecZnxAlloc, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft,
4 VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps, VecZnxToMut, VecZnxToRef, ZnxInfos,
5};
6use sampling::source::Source;
7
8use crate::{
9 elem::{FromProdBy, FromProdByScratchSpace, Infos, MatZnxDftProducts, ProdBy, ProdByScratchSpace},
10 grlwe::GRLWECt,
11 keys::{PublicKey, SecretDistribution, SecretKeyDft},
12 rgsw::RGSWCt,
13 utils::derive_size,
14};
15
16pub struct RLWECt<C> {
17 pub data: VecZnx<C>,
18 pub log_base2k: usize,
19 pub log_k: usize,
20}
21
22impl RLWECt<Vec<u8>> {
23 pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize) -> Self {
24 Self {
25 data: module.new_vec_znx(2, derive_size(log_base2k, log_k)),
26 log_base2k: log_base2k,
27 log_k: log_k,
28 }
29 }
30}
31
32impl<T> Infos for RLWECt<T> {
33 type Inner = VecZnx<T>;
34
35 fn inner(&self) -> &Self::Inner {
36 &self.data
37 }
38
39 fn log_base2k(&self) -> usize {
40 self.log_base2k
41 }
42
43 fn log_k(&self) -> usize {
44 self.log_k
45 }
46}
47
48impl<C> VecZnxToMut for RLWECt<C>
49where
50 VecZnx<C>: VecZnxToMut,
51{
52 fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
53 self.data.to_mut()
54 }
55}
56
57impl<C> VecZnxToRef for RLWECt<C>
58where
59 VecZnx<C>: VecZnxToRef,
60{
61 fn to_ref(&self) -> VecZnx<&[u8]> {
62 self.data.to_ref()
63 }
64}
65
66impl<C> RLWECt<C>
67where
68 VecZnx<C>: VecZnxToRef,
69{
70 #[allow(dead_code)]
71 pub(crate) fn dft<R>(&self, module: &Module<FFT64>, res: &mut RLWECtDft<R, FFT64>)
72 where
73 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + ZnxInfos,
74 {
75 #[cfg(debug_assertions)]
76 {
77 assert_eq!(self.cols(), 2);
78 assert_eq!(res.cols(), 2);
79 assert_eq!(self.log_base2k(), res.log_base2k())
80 }
81
82 module.vec_znx_dft(res, 0, self, 0);
83 module.vec_znx_dft(res, 1, self, 1);
84 }
85}
86
87impl ProdByScratchSpace for RLWECt<Vec<u8>> {
88 fn prod_by_grlwe_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize {
89 <GRLWECt<Vec<u8>, FFT64> as MatZnxDftProducts<GRLWECt<Vec<u8>, FFT64>, Vec<u8>>>::mul_rlwe_inplace_scratch_space(
90 module, lhs, rhs,
91 )
92 }
93
94 fn prod_by_rgsw_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize {
95 <RGSWCt<Vec<u8>, FFT64> as MatZnxDftProducts<RGSWCt<Vec<u8>, FFT64>, Vec<u8>>>::mul_rlwe_inplace_scratch_space(
96 module, lhs, rhs,
97 )
98 }
99}
100
101impl FromProdByScratchSpace for RLWECt<Vec<u8>> {
102 fn from_prod_by_grlwe_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
103 <GRLWECt<Vec<u8>, FFT64> as MatZnxDftProducts<GRLWECt<Vec<u8>, FFT64>, Vec<u8>>>::mul_rlwe_scratch_space(
104 module, res_size, lhs, rhs,
105 )
106 }
107
108 fn from_prod_by_rgsw_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
109 <RGSWCt<Vec<u8>, FFT64> as MatZnxDftProducts<RGSWCt<Vec<u8>, FFT64>, Vec<u8>>>::mul_rlwe_scratch_space(
110 module, res_size, lhs, rhs,
111 )
112 }
113}
114
115impl<MUT> ProdBy<RLWECt<MUT>> for RLWECt<MUT>
116where
117 VecZnx<MUT>: VecZnxToMut + VecZnxToRef,
118{
119 fn prod_by_grlwe<R>(&mut self, module: &Module<FFT64>, rhs: &GRLWECt<R, FFT64>, scratch: &mut Scratch)
120 where
121 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
122 {
123 rhs.mul_rlwe_inplace(module, self, scratch);
124 }
125
126 fn prod_by_rgsw<R>(&mut self, module: &Module<FFT64>, rhs: &RGSWCt<R, FFT64>, scratch: &mut Scratch)
127 where
128 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
129 {
130 rhs.mul_rlwe_inplace(module, self, scratch);
131 }
132}
133
134impl<MUT, REF> FromProdBy<RLWECt<MUT>, RLWECt<REF>> for RLWECt<MUT>
135where
136 VecZnx<MUT>: VecZnxToMut + VecZnxToRef,
137 VecZnx<REF>: VecZnxToRef,
138{
139 fn from_prod_by_grlwe<R>(&mut self, module: &Module<FFT64>, lhs: &RLWECt<REF>, rhs: &GRLWECt<R, FFT64>, scratch: &mut Scratch)
140 where
141 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
142 {
143 rhs.mul_rlwe(module, self, lhs, scratch);
144 }
145
146 fn from_prod_by_rgsw<R>(&mut self, module: &Module<FFT64>, lhs: &RLWECt<REF>, rhs: &RGSWCt<R, FFT64>, scratch: &mut Scratch)
147 where
148 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
149 {
150 rhs.mul_rlwe(module, self, lhs, scratch);
151 }
152}
153
154impl RLWECt<Vec<u8>> {
155 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, size: usize) -> usize {
156 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
157 }
158
159 pub fn encrypt_pk_scratch_space<B: Backend>(module: &Module<B>, pk_size: usize) -> usize {
160 ((module.bytes_of_vec_znx_dft(1, pk_size) + module.bytes_of_vec_znx_big(1, pk_size)) | module.bytes_of_scalar_znx(1))
161 + module.bytes_of_scalar_znx_dft(1)
162 + module.vec_znx_big_normalize_tmp_bytes()
163 }
164
165 pub fn decrypt_scratch_space<B: Backend>(module: &Module<B>, size: usize) -> usize {
166 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
167 }
168}
169
170pub fn encrypt_rlwe_sk<C, P, S>(
171 module: &Module<FFT64>,
172 ct: &mut RLWECt<C>,
173 pt: Option<(&RLWEPt<P>, usize)>,
174 sk_dft: &SecretKeyDft<S, FFT64>,
175 source_xa: &mut Source,
176 source_xe: &mut Source,
177 sigma: f64,
178 bound: f64,
179 scratch: &mut Scratch,
180) where
181 VecZnx<C>: VecZnxToMut + VecZnxToRef,
182 VecZnx<P>: VecZnxToRef,
183 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
184{
185 let log_base2k: usize = ct.log_base2k();
186 let log_k: usize = ct.log_k();
187 let size: usize = ct.size();
188
189 // c1 = a
190 ct.data.fill_uniform(log_base2k, 1, size, source_xa);
191
192 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
193
194 {
195 let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
196 module.vec_znx_dft(&mut c0_dft, 0, ct, 1);
197
198 // c0_dft = DFT(a) * DFT(s)
199 module.svp_apply_inplace(&mut c0_dft, 0, sk_dft, 0);
200
201 // c0_big = IDFT(c0_dft)
202 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
203 }
204
205 // c0_big = m - c0_big
206 if let Some((pt, col)) = pt {
207 match col {
208 0 => module.vec_znx_big_sub_small_b_inplace(&mut c0_big, 0, pt, 0),
209 1 => {
210 module.vec_znx_big_negate_inplace(&mut c0_big, 0);
211 module.vec_znx_add_inplace(ct, 1, pt, 0);
212 module.vec_znx_normalize_inplace(log_base2k, ct, 1, scratch_1);
213 }
214 _ => panic!("invalid target column: {}", col),
215 }
216 } else {
217 module.vec_znx_big_negate_inplace(&mut c0_big, 0);
218 }
219 // c0_big += e
220 c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, bound);
221
222 // c0 = norm(c0_big = -as + m + e)
223 module.vec_znx_big_normalize(log_base2k, ct, 0, &c0_big, 0, scratch_1);
224}
225
226pub fn decrypt_rlwe<P, C, S>(
227 module: &Module<FFT64>,
228 pt: &mut RLWEPt<P>,
229 ct: &RLWECt<C>,
230 sk_dft: &SecretKeyDft<S, FFT64>,
231 scratch: &mut Scratch,
232) where
233 VecZnx<P>: VecZnxToMut + VecZnxToRef,
234 VecZnx<C>: VecZnxToRef,
235 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
236{
237 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, ct.size()); // TODO optimize size when pt << ct
238
239 {
240 let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, ct.size()); // TODO optimize size when pt << ct
241 module.vec_znx_dft(&mut c0_dft, 0, ct, 1);
242
243 // c0_dft = DFT(a) * DFT(s)
244 module.svp_apply_inplace(&mut c0_dft, 0, sk_dft, 0);
245
246 // c0_big = IDFT(c0_dft)
247 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
248 }
249
250 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
251 module.vec_znx_big_add_small_inplace(&mut c0_big, 0, ct, 0);
252
253 // pt = norm(BIG(m + e))
254 module.vec_znx_big_normalize(ct.log_base2k(), pt, 0, &mut c0_big, 0, scratch_1);
255
256 pt.log_base2k = ct.log_base2k();
257 pt.log_k = pt.log_k().min(ct.log_k());
258}
259
260impl<C> RLWECt<C> {
261 pub fn encrypt_sk<P, S>(
262 &mut self,
263 module: &Module<FFT64>,
264 pt: Option<&RLWEPt<P>>,
265 sk_dft: &SecretKeyDft<S, FFT64>,
266 source_xa: &mut Source,
267 source_xe: &mut Source,
268 sigma: f64,
269 bound: f64,
270 scratch: &mut Scratch,
271 ) where
272 VecZnx<C>: VecZnxToMut + VecZnxToRef,
273 VecZnx<P>: VecZnxToRef,
274 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
275 {
276 if let Some(pt) = pt {
277 encrypt_rlwe_sk(
278 module,
279 self,
280 Some((pt, 0)),
281 sk_dft,
282 source_xa,
283 source_xe,
284 sigma,
285 bound,
286 scratch,
287 )
288 } else {
289 encrypt_rlwe_sk::<C, P, S>(
290 module, self, None, sk_dft, source_xa, source_xe, sigma, bound, scratch,
291 )
292 }
293 }
294
295 pub fn decrypt<P, S>(
296 &self,
297 module: &Module<FFT64>,
298 pt: &mut RLWEPt<P>,
299 sk_dft: &SecretKeyDft<S, FFT64>,
300 scratch: &mut Scratch,
301 ) where
302 VecZnx<P>: VecZnxToMut + VecZnxToRef,
303 VecZnx<C>: VecZnxToRef,
304 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
305 {
306 decrypt_rlwe(module, pt, self, sk_dft, scratch);
307 }
308
309 pub fn encrypt_pk<P, S>(
310 &mut self,
311 module: &Module<FFT64>,
312 pt: Option<&RLWEPt<P>>,
313 pk: &PublicKey<S, FFT64>,
314 source_xu: &mut Source,
315 source_xe: &mut Source,
316 sigma: f64,
317 bound: f64,
318 scratch: &mut Scratch,
319 ) where
320 VecZnx<C>: VecZnxToMut + VecZnxToRef,
321 VecZnx<P>: VecZnxToRef,
322 VecZnxDft<S, FFT64>: VecZnxDftToRef<FFT64>,
323 {
324 encrypt_rlwe_pk(
325 module, self, pt, pk, source_xu, source_xe, sigma, bound, scratch,
326 )
327 }
328}
329
330pub(crate) fn encrypt_rlwe_pk<C, P, S>(
331 module: &Module<FFT64>,
332 ct: &mut RLWECt<C>,
333 pt: Option<&RLWEPt<P>>,
334 pk: &PublicKey<S, FFT64>,
335 source_xu: &mut Source,
336 source_xe: &mut Source,
337 sigma: f64,
338 bound: f64,
339 scratch: &mut Scratch,
340) where
341 VecZnx<C>: VecZnxToMut + VecZnxToRef,
342 VecZnx<P>: VecZnxToRef,
343 VecZnxDft<S, FFT64>: VecZnxDftToRef<FFT64>,
344{
345 #[cfg(debug_assertions)]
346 {
347 assert_eq!(ct.log_base2k(), pk.log_base2k());
348 assert_eq!(ct.n(), module.n());
349 assert_eq!(pk.n(), module.n());
350 if let Some(pt) = pt {
351 assert_eq!(pt.log_base2k(), pk.log_base2k());
352 assert_eq!(pt.n(), module.n());
353 }
354 }
355
356 let log_base2k: usize = pk.log_base2k();
357 let size_pk: usize = pk.size();
358
359 // Generates u according to the underlying secret distribution.
360 let (mut u_dft, scratch_1) = scratch.tmp_scalar_znx_dft(module, 1);
361
362 {
363 let (mut u, _) = scratch_1.tmp_scalar_znx(module, 1);
364 match pk.dist {
365 SecretDistribution::NONE => panic!(
366 "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through Self::generate"
367 ),
368 SecretDistribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
369 SecretDistribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
370 SecretDistribution::ZERO => {}
371 }
372
373 module.svp_prepare(&mut u_dft, 0, &u, 0);
374 }
375
376 let (mut tmp_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, size_pk); // TODO optimize size (e.g. when encrypting at low homomorphic capacity)
377 let (mut tmp_dft, scratch_3) = scratch_2.tmp_vec_znx_dft(module, 1, size_pk); // TODO optimize size (e.g. when encrypting at low homomorphic capacity)
378
379 // ct[0] = pk[0] * u + m + e0
380 module.svp_apply(&mut tmp_dft, 0, &u_dft, 0, pk, 0);
381 module.vec_znx_idft_tmp_a(&mut tmp_big, 0, &mut tmp_dft, 0);
382 tmp_big.add_normal(log_base2k, 0, pk.log_k(), source_xe, sigma, bound);
383
384 if let Some(pt) = pt {
385 module.vec_znx_big_add_small_inplace(&mut tmp_big, 0, pt, 0);
386 }
387
388 module.vec_znx_big_normalize(log_base2k, ct, 0, &tmp_big, 0, scratch_3);
389
390 // ct[1] = pk[1] * u + e1
391 module.svp_apply(&mut tmp_dft, 0, &u_dft, 0, pk, 1);
392 module.vec_znx_idft_tmp_a(&mut tmp_big, 0, &mut tmp_dft, 0);
393 tmp_big.add_normal(log_base2k, 0, pk.log_k(), source_xe, sigma, bound);
394 module.vec_znx_big_normalize(log_base2k, ct, 1, &tmp_big, 0, scratch_3);
395}
396
397pub struct RLWEPt<C> {
398 pub data: VecZnx<C>,
399 pub log_base2k: usize,
400 pub log_k: usize,
401}
402
403impl<T> Infos for RLWEPt<T> {
404 type Inner = VecZnx<T>;
405
406 fn inner(&self) -> &Self::Inner {
407 &self.data
408 }
409
410 fn log_base2k(&self) -> usize {
411 self.log_base2k
412 }
413
414 fn log_k(&self) -> usize {
415 self.log_k
416 }
417}
418
419impl<C> VecZnxToMut for RLWEPt<C>
420where
421 VecZnx<C>: VecZnxToMut,
422{
423 fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
424 self.data.to_mut()
425 }
426}
427
428impl<C> VecZnxToRef for RLWEPt<C>
429where
430 VecZnx<C>: VecZnxToRef,
431{
432 fn to_ref(&self) -> VecZnx<&[u8]> {
433 self.data.to_ref()
434 }
435}
436
437impl RLWEPt<Vec<u8>> {
438 pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize) -> Self {
439 Self {
440 data: module.new_vec_znx(1, derive_size(log_base2k, log_k)),
441 log_base2k: log_base2k,
442 log_k: log_k,
443 }
444 }
445}
446
447pub struct RLWECtDft<C, B: Backend> {
448 pub data: VecZnxDft<C, B>,
449 pub log_base2k: usize,
450 pub log_k: usize,
451}
452
453impl<B: Backend> RLWECtDft<Vec<u8>, B> {
454 pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize) -> Self {
455 Self {
456 data: module.new_vec_znx_dft(2, derive_size(log_base2k, log_k)),
457 log_base2k: log_base2k,
458 log_k: log_k,
459 }
460 }
461}
462
463impl<T, B: Backend> Infos for RLWECtDft<T, B> {
464 type Inner = VecZnxDft<T, B>;
465
466 fn inner(&self) -> &Self::Inner {
467 &self.data
468 }
469
470 fn log_base2k(&self) -> usize {
471 self.log_base2k
472 }
473
474 fn log_k(&self) -> usize {
475 self.log_k
476 }
477}
478
479impl<C, B: Backend> VecZnxDftToMut<B> for RLWECtDft<C, B>
480where
481 VecZnxDft<C, B>: VecZnxDftToMut<B>,
482{
483 fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
484 self.data.to_mut()
485 }
486}
487
488impl<C, B: Backend> VecZnxDftToRef<B> for RLWECtDft<C, B>
489where
490 VecZnxDft<C, B>: VecZnxDftToRef<B>,
491{
492 fn to_ref(&self) -> VecZnxDft<&[u8], B> {
493 self.data.to_ref()
494 }
495}
496
497impl<C> RLWECtDft<C, FFT64>
498where
499 VecZnxDft<C, FFT64>: VecZnxDftToRef<FFT64>,
500{
501 #[allow(dead_code)]
502 pub(crate) fn idft_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
503 module.bytes_of_vec_znx(2, size) + (module.vec_znx_big_normalize_tmp_bytes() | module.vec_znx_idft_tmp_bytes())
504 }
505
506 pub(crate) fn idft<R>(&self, module: &Module<FFT64>, res: &mut RLWECt<R>, scratch: &mut Scratch)
507 where
508 VecZnx<R>: VecZnxToMut,
509 {
510 #[cfg(debug_assertions)]
511 {
512 assert_eq!(self.cols(), 2);
513 assert_eq!(res.cols(), 2);
514 assert_eq!(self.log_base2k(), res.log_base2k())
515 }
516
517 let min_size: usize = self.size().min(res.size());
518
519 let (mut res_big, scratch1) = scratch.tmp_vec_znx_big(module, 2, min_size);
520
521 module.vec_znx_idft(&mut res_big, 0, &self.data, 0, scratch1);
522 module.vec_znx_idft(&mut res_big, 1, &self.data, 1, scratch1);
523 module.vec_znx_big_normalize(self.log_base2k(), res, 0, &res_big, 0, scratch1);
524 module.vec_znx_big_normalize(self.log_base2k(), res, 1, &res_big, 1, scratch1);
525 }
526}
527
528pub(crate) fn encrypt_zero_rlwe_dft_sk<C, S>(
529 module: &Module<FFT64>,
530 ct: &mut RLWECtDft<C, FFT64>,
531 sk: &SecretKeyDft<S, FFT64>,
532 source_xa: &mut Source,
533 source_xe: &mut Source,
534 sigma: f64,
535 bound: f64,
536 scratch: &mut Scratch,
537) where
538 VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
539 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
540{
541 let log_base2k: usize = ct.log_base2k();
542 let log_k: usize = ct.log_k();
543 let size: usize = ct.size();
544
545 #[cfg(debug_assertions)]
546 {
547 match sk.dist {
548 SecretDistribution::NONE => panic!("invalid sk.dist = SecretDistribution::NONE"),
549 _ => {}
550 }
551 assert_eq!(ct.cols(), 2);
552 }
553
554 // ct[1] = DFT(a)
555 {
556 let (mut tmp_znx, _) = scratch.tmp_vec_znx(module, 1, size);
557 tmp_znx.fill_uniform(log_base2k, 0, size, source_xa);
558 module.vec_znx_dft(ct, 1, &tmp_znx, 0);
559 }
560
561 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
562
563 {
564 let (mut tmp_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
565 // c0_dft = ct[1] * DFT(s)
566 module.svp_apply(&mut tmp_dft, 0, sk, 0, ct, 1);
567
568 // c0_big = IDFT(c0_dft)
569 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut tmp_dft, 0);
570 }
571
572 // c0_big += e
573 c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, bound);
574
575 // c0 = norm(c0_big = -as - e), NOTE: e is centered at 0.
576 let (mut tmp_znx, scratch_2) = scratch_1.tmp_vec_znx(module, 1, size);
577 module.vec_znx_big_normalize(log_base2k, &mut tmp_znx, 0, &c0_big, 0, scratch_2);
578 module.vec_znx_negate_inplace(&mut tmp_znx, 0);
579 // ct[0] = DFT(-as + e)
580 module.vec_znx_dft(ct, 0, &tmp_znx, 0);
581}
582
583impl RLWECtDft<Vec<u8>, FFT64> {
584 pub fn encrypt_zero_sk_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
585 (module.bytes_of_vec_znx(1, size) | module.bytes_of_vec_znx_dft(1, size))
586 + module.bytes_of_vec_znx_big(1, size)
587 + module.bytes_of_vec_znx(1, size)
588 + module.vec_znx_big_normalize_tmp_bytes()
589 }
590
591 pub fn decrypt_scratch_space(module: &Module<FFT64>, size: usize) -> usize {
592 (module.vec_znx_big_normalize_tmp_bytes()
593 | module.bytes_of_vec_znx_dft(1, size)
594 | (module.bytes_of_vec_znx_big(1, size) + module.vec_znx_idft_tmp_bytes()))
595 + module.bytes_of_vec_znx_big(1, size)
596 }
597}
598
599pub fn decrypt_rlwe_dft<P, C, S>(
600 module: &Module<FFT64>,
601 pt: &mut RLWEPt<P>,
602 ct: &RLWECtDft<C, FFT64>,
603 sk: &SecretKeyDft<S, FFT64>,
604 scratch: &mut Scratch,
605) where
606 VecZnx<P>: VecZnxToMut + VecZnxToRef,
607 VecZnxDft<C, FFT64>: VecZnxDftToRef<FFT64>,
608 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
609{
610 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, ct.size()); // TODO optimize size when pt << ct
611
612 {
613 let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, ct.size()); // TODO optimize size when pt << ct
614 // c0_dft = DFT(a) * DFT(s)
615 module.svp_apply(&mut c0_dft, 0, sk, 0, ct, 1);
616 // c0_big = IDFT(c0_dft)
617 module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
618 }
619
620 {
621 let (mut c1_big, scratch_2) = scratch_1.tmp_vec_znx_big(module, 1, ct.size());
622 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
623 module.vec_znx_idft(&mut c1_big, 0, ct, 0, scratch_2);
624 module.vec_znx_big_add_inplace(&mut c0_big, 0, &c1_big, 0);
625 }
626
627 // pt = norm(BIG(m + e))
628 module.vec_znx_big_normalize(ct.log_base2k(), pt, 0, &mut c0_big, 0, scratch_1);
629
630 pt.log_base2k = ct.log_base2k();
631 pt.log_k = pt.log_k().min(ct.log_k());
632}
633
634impl<C> RLWECtDft<C, FFT64> {
635 pub(crate) fn encrypt_zero_sk<S>(
636 &mut self,
637 module: &Module<FFT64>,
638 sk_dft: &SecretKeyDft<S, FFT64>,
639 source_xa: &mut Source,
640 source_xe: &mut Source,
641 sigma: f64,
642 bound: f64,
643 scratch: &mut Scratch,
644 ) where
645 VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
646 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
647 {
648 encrypt_zero_rlwe_dft_sk(
649 module, self, sk_dft, source_xa, source_xe, sigma, bound, scratch,
650 )
651 }
652
653 pub fn decrypt<P, S>(
654 &self,
655 module: &Module<FFT64>,
656 pt: &mut RLWEPt<P>,
657 sk_dft: &SecretKeyDft<S, FFT64>,
658 scratch: &mut Scratch,
659 ) where
660 VecZnx<P>: VecZnxToMut + VecZnxToRef,
661 VecZnxDft<C, FFT64>: VecZnxDftToRef<FFT64>,
662 ScalarZnxDft<S, FFT64>: ScalarZnxDftToRef<FFT64>,
663 {
664 decrypt_rlwe_dft(module, pt, self, sk_dft, scratch);
665 }
666}
667
668impl ProdByScratchSpace for RLWECtDft<Vec<u8>, FFT64> {
669 fn prod_by_grlwe_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize {
670 <GRLWECt<Vec<u8>, FFT64> as MatZnxDftProducts<GRLWECt<Vec<u8>, FFT64>, Vec<u8>>>::mul_rlwe_dft_inplace_scratch_space(
671 module, lhs, rhs,
672 )
673 }
674
675 fn prod_by_rgsw_scratch_space(module: &Module<FFT64>, lhs: usize, rhs: usize) -> usize {
676 <RGSWCt<Vec<u8>, FFT64> as MatZnxDftProducts<RGSWCt<Vec<u8>, FFT64>, Vec<u8>>>::mul_rlwe_dft_inplace_scratch_space(
677 module, lhs, rhs,
678 )
679 }
680}
681
682impl FromProdByScratchSpace for RLWECtDft<Vec<u8>, FFT64> {
683 fn from_prod_by_grlwe_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
684 <GRLWECt<Vec<u8>, FFT64> as MatZnxDftProducts<GRLWECt<Vec<u8>, FFT64>, Vec<u8>>>::mul_rlwe_dft_scratch_space(
685 module, res_size, lhs, rhs,
686 )
687 }
688
689 fn from_prod_by_rgsw_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
690 <RGSWCt<Vec<u8>, FFT64> as MatZnxDftProducts<RGSWCt<Vec<u8>, FFT64>, Vec<u8>>>::mul_rlwe_dft_scratch_space(
691 module, res_size, lhs, rhs,
692 )
693 }
694}
695
696impl<MUT> ProdBy<RLWECtDft<MUT, FFT64>> for RLWECtDft<MUT, FFT64>
697where
698 VecZnxDft<MUT, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
699{
700 fn prod_by_grlwe<R>(&mut self, module: &Module<FFT64>, rhs: &GRLWECt<R, FFT64>, scratch: &mut Scratch)
701 where
702 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
703 {
704 rhs.mul_rlwe_dft_inplace(module, self, scratch);
705 }
706
707 fn prod_by_rgsw<R>(&mut self, module: &Module<FFT64>, rhs: &RGSWCt<R, FFT64>, scratch: &mut Scratch)
708 where
709 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
710 {
711 rhs.mul_rlwe_dft_inplace(module, self, scratch);
712 }
713}
714
715impl<MUT, REF> FromProdBy<RLWECtDft<MUT, FFT64>, RLWECtDft<REF, FFT64>> for RLWECtDft<MUT, FFT64>
716where
717 VecZnxDft<MUT, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
718 VecZnxDft<REF, FFT64>: VecZnxDftToRef<FFT64>,
719{
720 fn from_prod_by_grlwe<R>(
721 &mut self,
722 module: &Module<FFT64>,
723 lhs: &RLWECtDft<REF, FFT64>,
724 rhs: &GRLWECt<R, FFT64>,
725 scratch: &mut Scratch,
726 ) where
727 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
728 {
729 rhs.mul_rlwe_dft(module, self, lhs, scratch);
730 }
731
732 fn from_prod_by_rgsw<R>(
733 &mut self,
734 module: &Module<FFT64>,
735 lhs: &RLWECtDft<REF, FFT64>,
736 rhs: &RGSWCt<R, FFT64>,
737 scratch: &mut Scratch,
738 ) where
739 MatZnxDft<R, FFT64>: MatZnxDftToRef<FFT64>,
740 {
741 rhs.mul_rlwe_dft(module, self, lhs, scratch);
742 }
743}