poulpy/core/src/gglwe/encryption.rs
core/src/gglwe/encryption.rsRaw
1use backend::hal::{
2 api::{
3 ScratchAvailable, SvpApply, TakeScalarZnx, TakeVecZnx, TakeVecZnxBig, TakeVecZnxDft, VecZnxAddScalarInplace,
4 VecZnxAutomorphism, VecZnxBigAllocBytes, VecZnxDftToVecZnxBigTmpA, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes,
5 VecZnxSwithcDegree, ZnxZero,
6 },
7 layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
8};
9use sampling::source::Source;
10
11use crate::{
12 AutomorphismKey, AutomorphismKeyCompressed, GGLWECiphertext, GGLWECiphertextCompressed, GLWECiphertext, GLWEDecryptFamily,
13 GLWEEncryptSkFamily, GLWEPlaintext, GLWESecret, GLWESecretExec, GLWESecretFamily, GLWESwitchingKey,
14 GLWESwitchingKeyCompressed, GLWETensorKey, GLWETensorKeyCompressed, Infos, TakeGLWEPt, TakeGLWESecret, TakeGLWESecretExec,
15 encrypt_sk_internal,
16};
17
18pub trait GGLWEEncryptSkFamily<B: Backend> = GLWEEncryptSkFamily<B> + GLWESecretFamily<B>;
19
20impl GGLWECiphertext<Vec<u8>> {
21 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
22 where
23 Module<B>: GGLWEEncryptSkFamily<B>,
24 {
25 GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k)
26 + (GLWEPlaintext::byte_of(n, basek, k) | module.vec_znx_normalize_tmp_bytes(n))
27 }
28
29 pub fn encrypt_pk_scratch_space<B: Backend>(_module: &Module<B>, _n: usize, _basek: usize, _k: usize, _rank: usize) -> usize {
30 unimplemented!()
31 }
32}
33
34impl<DataSelf: DataMut> GGLWECiphertext<DataSelf> {
35 pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
36 &mut self,
37 module: &Module<B>,
38 pt: &ScalarZnx<DataPt>,
39 sk: &GLWESecretExec<DataSk, B>,
40 source_xa: &mut Source,
41 source_xe: &mut Source,
42 sigma: f64,
43 scratch: &mut Scratch<B>,
44 ) where
45 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAddScalarInplace,
46 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
47 {
48 #[cfg(debug_assertions)]
49 {
50 use backend::hal::api::ZnxInfos;
51
52 assert_eq!(
53 self.rank_in(),
54 pt.cols(),
55 "self.rank_in(): {} != pt.cols(): {}",
56 self.rank_in(),
57 pt.cols()
58 );
59 assert_eq!(
60 self.rank_out(),
61 sk.rank(),
62 "self.rank_out(): {} != sk.rank(): {}",
63 self.rank_out(),
64 sk.rank()
65 );
66 assert_eq!(self.n(), sk.n());
67 assert_eq!(pt.n(), sk.n());
68 assert!(
69 scratch.available() >= GGLWECiphertext::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k()),
70 "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
71 scratch.available(),
72 self.rank(),
73 self.size(),
74 GGLWECiphertext::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k())
75 );
76 assert!(
77 self.rows() * self.digits() * self.basek() <= self.k(),
78 "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
79 self.rows(),
80 self.digits(),
81 self.basek(),
82 self.rows() * self.digits() * self.basek(),
83 self.k()
84 );
85 }
86
87 let rows: usize = self.rows();
88 let digits: usize = self.digits();
89 let basek: usize = self.basek();
90 let k: usize = self.k();
91 let rank_in: usize = self.rank_in();
92
93 let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(sk.n(), basek, k);
94 // For each input column (i.e. rank) produces a GGLWE ciphertext of rank_out+1 columns
95 //
96 // Example for ksk rank 2 to rank 3:
97 //
98 // (-(a0*s0 + a1*s1 + a2*s2) + s0', a0, a1, a2)
99 // (-(b0*s0 + b1*s1 + b2*s2) + s0', b0, b1, b2)
100 //
101 // Example ksk rank 2 to rank 1
102 //
103 // (-(a*s) + s0, a)
104 // (-(b*s) + s1, b)
105 (0..rank_in).for_each(|col_i| {
106 (0..rows).for_each(|row_i| {
107 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
108 tmp_pt.data.zero(); // zeroes for next iteration
109 module.vec_znx_add_scalar_inplace(
110 &mut tmp_pt.data,
111 0,
112 (digits - 1) + row_i * digits,
113 pt,
114 col_i,
115 );
116 module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scrach_1);
117
118 // rlwe encrypt of vec_znx_pt into vec_znx_ct
119 self.at_mut(row_i, col_i)
120 .encrypt_sk(module, &tmp_pt, sk, source_xa, source_xe, sigma, scrach_1);
121 });
122 });
123 }
124}
125
126impl GGLWECiphertextCompressed<Vec<u8>> {
127 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
128 where
129 Module<B>: GLWESwitchingKeyEncryptSkFamily<B>,
130 {
131 GGLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k)
132 }
133}
134
135impl<D: DataMut> GGLWECiphertextCompressed<D> {
136 pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
137 &mut self,
138 module: &Module<B>,
139 pt: &ScalarZnx<DataPt>,
140 sk: &GLWESecretExec<DataSk, B>,
141 seed: [u8; 32],
142 source_xe: &mut Source,
143 sigma: f64,
144 scratch: &mut Scratch<B>,
145 ) where
146 Module<B>: GGLWEEncryptSkFamily<B> + VecZnxAddScalarInplace,
147 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
148 {
149 #[cfg(debug_assertions)]
150 {
151 use backend::hal::api::ZnxInfos;
152
153 assert_eq!(
154 self.rank_in(),
155 pt.cols(),
156 "self.rank_in(): {} != pt.cols(): {}",
157 self.rank_in(),
158 pt.cols()
159 );
160 assert_eq!(
161 self.rank_out(),
162 sk.rank(),
163 "self.rank_out(): {} != sk.rank(): {}",
164 self.rank_out(),
165 sk.rank()
166 );
167 assert_eq!(self.n(), sk.n());
168 assert_eq!(pt.n(), sk.n());
169 assert!(
170 scratch.available()
171 >= GGLWECiphertextCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k()),
172 "scratch.available: {} < GGLWECiphertext::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
173 scratch.available(),
174 self.rank(),
175 self.size(),
176 GGLWECiphertextCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k())
177 );
178 assert!(
179 self.rows() * self.digits() * self.basek() <= self.k(),
180 "self.rows() : {} * self.digits() : {} * self.basek() : {} = {} >= self.k() = {}",
181 self.rows(),
182 self.digits(),
183 self.basek(),
184 self.rows() * self.digits() * self.basek(),
185 self.k()
186 );
187 }
188
189 let rows: usize = self.rows();
190 let digits: usize = self.digits();
191 let basek: usize = self.basek();
192 let k: usize = self.k();
193 let rank_in: usize = self.rank_in();
194 let cols: usize = self.rank_out() + 1;
195
196 let mut source_xa = Source::new(seed);
197
198 let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(sk.n(), basek, k);
199 (0..rank_in).for_each(|col_i| {
200 (0..rows).for_each(|row_i| {
201 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
202 tmp_pt.data.zero(); // zeroes for next iteration
203 module.vec_znx_add_scalar_inplace(
204 &mut tmp_pt.data,
205 0,
206 (digits - 1) + row_i * digits,
207 pt,
208 col_i,
209 );
210 module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scrach_1);
211
212 let (seed, mut source_xa_tmp) = source_xa.branch();
213 self.seed[col_i * rows + row_i] = seed;
214
215 encrypt_sk_internal(
216 module,
217 self.basek(),
218 self.k(),
219 &mut self.at_mut(row_i, col_i).data,
220 cols,
221 true,
222 Some((&tmp_pt, 0)),
223 sk,
224 &mut source_xa_tmp,
225 source_xe,
226 sigma,
227 scrach_1,
228 );
229 });
230 });
231 }
232}
233
234pub trait GLWESwitchingKeyEncryptSkFamily<B: Backend> = GGLWEEncryptSkFamily<B>;
235
236impl GLWESwitchingKey<Vec<u8>> {
237 pub fn encrypt_sk_scratch_space<B: Backend>(
238 module: &Module<B>,
239 n: usize,
240 basek: usize,
241 k: usize,
242 rank_in: usize,
243 rank_out: usize,
244 ) -> usize
245 where
246 Module<B>: GLWESwitchingKeyEncryptSkFamily<B>,
247 {
248 (GGLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k) | ScalarZnx::alloc_bytes(n, 1))
249 + ScalarZnx::alloc_bytes(n, rank_in)
250 + GLWESecretExec::bytes_of(module, n, rank_out)
251 }
252
253 pub fn encrypt_pk_scratch_space<B: Backend>(
254 module: &Module<B>,
255 _n: usize,
256 _basek: usize,
257 _k: usize,
258 _rank_in: usize,
259 _rank_out: usize,
260 ) -> usize {
261 GGLWECiphertext::encrypt_pk_scratch_space(module, _n, _basek, _k, _rank_out)
262 }
263}
264
265impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
266 pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
267 &mut self,
268 module: &Module<B>,
269 sk_in: &GLWESecret<DataSkIn>,
270 sk_out: &GLWESecret<DataSkOut>,
271 source_xa: &mut Source,
272 source_xe: &mut Source,
273 sigma: f64,
274 scratch: &mut Scratch<B>,
275 ) where
276 Module<B>: GLWESwitchingKeyEncryptSkFamily<B> + VecZnxSwithcDegree + VecZnxAddScalarInplace,
277 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + ScratchAvailable + TakeVecZnx,
278 {
279 #[cfg(debug_assertions)]
280 {
281 assert!(sk_in.n() <= module.n());
282 assert!(sk_out.n() <= module.n());
283 assert!(
284 scratch.available()
285 >= GLWESwitchingKey::encrypt_sk_scratch_space(
286 module,
287 sk_out.n(),
288 self.basek(),
289 self.k(),
290 self.rank_in(),
291 self.rank_out()
292 ),
293 "scratch.available()={} < GLWESwitchingKey::encrypt_sk_scratch_space={}",
294 scratch.available(),
295 GLWESwitchingKey::encrypt_sk_scratch_space(
296 module,
297 sk_out.n(),
298 self.basek(),
299 self.k(),
300 self.rank_in(),
301 self.rank_out()
302 )
303 )
304 }
305
306 let n: usize = sk_in.n().max(sk_out.n());
307
308 let (mut sk_in_tmp, scratch1) = scratch.take_scalar_znx(n, sk_in.rank());
309 (0..sk_in.rank()).for_each(|i| {
310 module.vec_znx_switch_degree(
311 &mut sk_in_tmp.as_vec_znx_mut(),
312 i,
313 &sk_in.data.as_vec_znx(),
314 i,
315 );
316 });
317
318 let (mut sk_out_tmp, scratch2) = scratch1.take_glwe_secret_exec(n, sk_out.rank());
319 {
320 let (mut tmp, _) = scratch2.take_scalar_znx(n, 1);
321 (0..sk_out.rank()).for_each(|i| {
322 module.vec_znx_switch_degree(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
323 module.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
324 });
325 }
326
327 self.key.encrypt_sk(
328 module,
329 &sk_in_tmp,
330 &sk_out_tmp,
331 source_xa,
332 source_xe,
333 sigma,
334 scratch2,
335 );
336 self.sk_in_n = sk_in.n();
337 self.sk_out_n = sk_out.n();
338 }
339}
340
341impl GLWESwitchingKeyCompressed<Vec<u8>> {
342 pub fn encrypt_sk_scratch_space<B: Backend>(
343 module: &Module<B>,
344 n: usize,
345 basek: usize,
346 k: usize,
347 rank_in: usize,
348 rank_out: usize,
349 ) -> usize
350 where
351 Module<B>: GLWESwitchingKeyEncryptSkFamily<B>,
352 {
353 (GGLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k) | ScalarZnx::alloc_bytes(n, 1))
354 + ScalarZnx::alloc_bytes(n, rank_in)
355 + GLWESecretExec::bytes_of(module, n, rank_out)
356 }
357}
358
359impl<DataSelf: DataMut> GLWESwitchingKeyCompressed<DataSelf> {
360 pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
361 &mut self,
362 module: &Module<B>,
363 sk_in: &GLWESecret<DataSkIn>,
364 sk_out: &GLWESecret<DataSkOut>,
365 seed_xa: [u8; 32],
366 source_xe: &mut Source,
367 sigma: f64,
368 scratch: &mut Scratch<B>,
369 ) where
370 Module<B>: GLWESwitchingKeyEncryptSkFamily<B> + VecZnxSwithcDegree + VecZnxAddScalarInplace,
371 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + ScratchAvailable + TakeVecZnx,
372 {
373 #[cfg(debug_assertions)]
374 {
375 assert!(sk_in.n() <= module.n());
376 assert!(sk_out.n() <= module.n());
377 assert!(
378 scratch.available()
379 >= GLWESwitchingKey::encrypt_sk_scratch_space(
380 module,
381 sk_out.n(),
382 self.basek(),
383 self.k(),
384 self.rank_in(),
385 self.rank_out()
386 ),
387 "scratch.available()={} < GLWESwitchingKey::encrypt_sk_scratch_space={}",
388 scratch.available(),
389 GLWESwitchingKey::encrypt_sk_scratch_space(
390 module,
391 sk_out.n(),
392 self.basek(),
393 self.k(),
394 self.rank_in(),
395 self.rank_out()
396 )
397 )
398 }
399
400 let n: usize = sk_in.n().max(sk_out.n());
401
402 let (mut sk_in_tmp, scratch1) = scratch.take_scalar_znx(n, sk_in.rank());
403 (0..sk_in.rank()).for_each(|i| {
404 module.vec_znx_switch_degree(
405 &mut sk_in_tmp.as_vec_znx_mut(),
406 i,
407 &sk_in.data.as_vec_znx(),
408 i,
409 );
410 });
411
412 let (mut sk_out_tmp, scratch2) = scratch1.take_glwe_secret_exec(n, sk_out.rank());
413 {
414 let (mut tmp, _) = scratch2.take_scalar_znx(n, 1);
415 (0..sk_out.rank()).for_each(|i| {
416 module.vec_znx_switch_degree(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
417 module.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
418 });
419 }
420
421 self.key.encrypt_sk(
422 module,
423 &sk_in_tmp,
424 &sk_out_tmp,
425 seed_xa,
426 source_xe,
427 sigma,
428 scratch2,
429 );
430 self.sk_in_n = sk_in.n();
431 self.sk_out_n = sk_out.n();
432 }
433}
434
435pub trait AutomorphismKeyEncryptSkFamily<B: Backend> = GGLWEEncryptSkFamily<B>;
436
437impl AutomorphismKey<Vec<u8>> {
438 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
439 where
440 Module<B>: AutomorphismKeyEncryptSkFamily<B>,
441 {
442 GLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, rank, rank) + GLWESecret::bytes_of(n, rank)
443 }
444
445 pub fn encrypt_pk_scratch_space<B: Backend>(module: &Module<B>, _n: usize, _basek: usize, _k: usize, _rank: usize) -> usize {
446 GLWESwitchingKey::encrypt_pk_scratch_space(module, _n, _basek, _k, _rank, _rank)
447 }
448}
449
450impl<DataSelf: DataMut> AutomorphismKey<DataSelf> {
451 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
452 &mut self,
453 module: &Module<B>,
454 p: i64,
455 sk: &GLWESecret<DataSk>,
456 source_xa: &mut Source,
457 source_xe: &mut Source,
458 sigma: f64,
459 scratch: &mut Scratch<B>,
460 ) where
461 Module<B>: AutomorphismKeyEncryptSkFamily<B> + VecZnxAutomorphism + VecZnxSwithcDegree + VecZnxAddScalarInplace,
462 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
463 {
464 #[cfg(debug_assertions)]
465 {
466 assert_eq!(self.n(), sk.n());
467 assert_eq!(self.rank_out(), self.rank_in());
468 assert_eq!(sk.rank(), self.rank());
469 assert!(
470 scratch.available()
471 >= AutomorphismKey::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank()),
472 "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
473 scratch.available(),
474 self.rank(),
475 self.size(),
476 AutomorphismKey::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank())
477 )
478 }
479
480 let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
481
482 {
483 (0..self.rank()).for_each(|i| {
484 module.vec_znx_automorphism(
485 module.galois_element_inv(p),
486 &mut sk_out.data.as_vec_znx_mut(),
487 i,
488 &sk.data.as_vec_znx(),
489 i,
490 );
491 });
492 }
493
494 self.key
495 .encrypt_sk(module, &sk, &sk_out, source_xa, source_xe, sigma, scratch_1);
496
497 self.p = p;
498 }
499}
500
501impl AutomorphismKeyCompressed<Vec<u8>> {
502 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
503 where
504 Module<B>: AutomorphismKeyEncryptSkFamily<B>,
505 {
506 GLWESwitchingKeyCompressed::encrypt_sk_scratch_space(module, n, basek, k, rank, rank) + GLWESecret::bytes_of(n, rank)
507 }
508}
509
510impl<DataSelf: DataMut> AutomorphismKeyCompressed<DataSelf> {
511 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
512 &mut self,
513 module: &Module<B>,
514 p: i64,
515 sk: &GLWESecret<DataSk>,
516 seed_xa: [u8; 32],
517 source_xe: &mut Source,
518 sigma: f64,
519 scratch: &mut Scratch<B>,
520 ) where
521 Module<B>: AutomorphismKeyEncryptSkFamily<B> + VecZnxSwithcDegree + VecZnxAutomorphism + VecZnxAddScalarInplace,
522 Scratch<B>: ScratchAvailable + TakeScalarZnx + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx,
523 {
524 #[cfg(debug_assertions)]
525 {
526 assert_eq!(self.n(), sk.n());
527 assert_eq!(self.rank_out(), self.rank_in());
528 assert_eq!(sk.rank(), self.rank());
529 assert!(
530 scratch.available()
531 >= AutomorphismKeyCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank()),
532 "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
533 scratch.available(),
534 self.rank(),
535 self.size(),
536 AutomorphismKeyCompressed::encrypt_sk_scratch_space(module, sk.n(), self.basek(), self.k(), self.rank())
537 )
538 }
539
540 let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
541
542 {
543 (0..self.rank()).for_each(|i| {
544 module.vec_znx_automorphism(
545 module.galois_element_inv(p),
546 &mut sk_out.data.as_vec_znx_mut(),
547 i,
548 &sk.data.as_vec_znx(),
549 i,
550 );
551 });
552 }
553
554 self.key
555 .encrypt_sk(module, &sk, &sk_out, seed_xa, source_xe, sigma, scratch_1);
556
557 self.p = p;
558 }
559}
560
561pub trait GLWETensorKeyEncryptSkFamily<B: Backend> =
562 GGLWEEncryptSkFamily<B> + VecZnxBigAllocBytes + VecZnxDftToVecZnxBigTmpA<B> + SvpApply<B>;
563
564impl GLWETensorKey<Vec<u8>> {
565 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
566 where
567 Module<B>: GLWETensorKeyEncryptSkFamily<B>,
568 {
569 GLWESecretExec::bytes_of(module, n, rank)
570 + module.vec_znx_dft_alloc_bytes(n, rank, 1)
571 + module.vec_znx_big_alloc_bytes(n, 1, 1)
572 + module.vec_znx_dft_alloc_bytes(n, 1, 1)
573 + GLWESecret::bytes_of(n, 1)
574 + GLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, rank, rank)
575 }
576}
577
578impl<DataSelf: DataMut> GLWETensorKey<DataSelf> {
579 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
580 &mut self,
581 module: &Module<B>,
582 sk: &GLWESecret<DataSk>,
583 source_xa: &mut Source,
584 source_xe: &mut Source,
585 sigma: f64,
586 scratch: &mut Scratch<B>,
587 ) where
588 Module<B>: GLWETensorKeyEncryptSkFamily<B> + VecZnxSwithcDegree + VecZnxAddScalarInplace,
589 Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecretExec<B> + TakeScalarZnx + TakeVecZnx,
590 {
591 #[cfg(debug_assertions)]
592 {
593 assert_eq!(self.rank(), sk.rank());
594 assert_eq!(self.n(), sk.n());
595 }
596
597 let n: usize = sk.n();
598
599 let rank: usize = self.rank();
600
601 let (mut sk_dft_prep, scratch1) = scratch.take_glwe_secret_exec(n, rank);
602 sk_dft_prep.prepare(module, &sk);
603
604 let (mut sk_dft, scratch2) = scratch1.take_vec_znx_dft(n, rank, 1);
605
606 (0..rank).for_each(|i| {
607 module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
608 });
609
610 let (mut sk_ij_big, scratch3) = scratch2.take_vec_znx_big(n, 1, 1);
611 let (mut sk_ij, scratch4) = scratch3.take_glwe_secret(n, 1);
612 let (mut sk_ij_dft, scratch5) = scratch4.take_vec_znx_dft(n, 1, 1);
613
614 (0..rank).for_each(|i| {
615 (i..rank).for_each(|j| {
616 module.svp_apply(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
617
618 module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
619 module.vec_znx_big_normalize(
620 self.basek(),
621 &mut sk_ij.data.as_vec_znx_mut(),
622 0,
623 &sk_ij_big,
624 0,
625 scratch5,
626 );
627
628 self.at_mut(i, j)
629 .encrypt_sk(module, &sk_ij, sk, source_xa, source_xe, sigma, scratch5);
630 });
631 })
632 }
633}
634
635impl GLWETensorKeyCompressed<Vec<u8>> {
636 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank: usize) -> usize
637 where
638 Module<B>: GLWETensorKeyEncryptSkFamily<B>,
639 {
640 GLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k, rank)
641 }
642}
643
644impl<DataSelf: DataMut> GLWETensorKeyCompressed<DataSelf> {
645 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
646 &mut self,
647 module: &Module<B>,
648 sk: &GLWESecret<DataSk>,
649 seed_xa: [u8; 32],
650 source_xe: &mut Source,
651 sigma: f64,
652 scratch: &mut Scratch<B>,
653 ) where
654 Module<B>: GLWETensorKeyEncryptSkFamily<B> + VecZnxSwithcDegree + VecZnxAddScalarInplace,
655 Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecretExec<B> + TakeScalarZnx + TakeVecZnx,
656 {
657 #[cfg(debug_assertions)]
658 {
659 assert_eq!(self.rank(), sk.rank());
660 assert_eq!(self.n(), sk.n());
661 }
662
663 let n: usize = sk.n();
664 let rank: usize = self.rank();
665
666 let (mut sk_dft_prep, scratch1) = scratch.take_glwe_secret_exec(n, rank);
667 sk_dft_prep.prepare(module, &sk);
668
669 let (mut sk_dft, scratch2) = scratch1.take_vec_znx_dft(n, rank, 1);
670
671 (0..rank).for_each(|i| {
672 module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
673 });
674
675 let (mut sk_ij_big, scratch3) = scratch2.take_vec_znx_big(n, 1, 1);
676 let (mut sk_ij, scratch4) = scratch3.take_glwe_secret(n, 1);
677 let (mut sk_ij_dft, scratch5) = scratch4.take_vec_znx_dft(n, 1, 1);
678
679 let mut source_xa: Source = Source::new(seed_xa);
680
681 (0..rank).for_each(|i| {
682 (i..rank).for_each(|j| {
683 module.svp_apply(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
684
685 module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
686 module.vec_znx_big_normalize(
687 self.basek(),
688 &mut sk_ij.data.as_vec_znx_mut(),
689 0,
690 &sk_ij_big,
691 0,
692 scratch5,
693 );
694
695 let (seed_xa_tmp, _) = source_xa.branch();
696
697 self.at_mut(i, j)
698 .encrypt_sk(module, &sk_ij, sk, seed_xa_tmp, source_xe, sigma, scratch5);
699 });
700 })
701 }
702}