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