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 let cols: usize = self.rank_out() + 1;
196
197 let mut source_xa = Source::new(seed);
198
199 let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(module, basek, k);
200 (0..rank_in).for_each(|col_i| {
201 (0..rows).for_each(|row_i| {
202 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
203 tmp_pt.data.zero(); // zeroes for next iteration
204 module.vec_znx_add_scalar_inplace(
205 &mut tmp_pt.data,
206 0,
207 (digits - 1) + row_i * digits,
208 pt,
209 col_i,
210 );
211 module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scrach_1);
212
213 let (seed, mut source_xa_tmp) = source_xa.branch();
214 self.seed[col_i * rows + row_i] = seed;
215
216 encrypt_sk_internal(
217 module,
218 self.basek(),
219 self.k(),
220 &mut self.at_mut(row_i, col_i).data,
221 cols,
222 true,
223 Some((&tmp_pt, 0)),
224 sk,
225 &mut source_xa_tmp,
226 source_xe,
227 sigma,
228 scrach_1,
229 );
230 });
231 });
232 }
233}
234
235pub trait GLWESwitchingKeyEncryptSkFamily<B: Backend> = GGLWEEncryptSkFamily<B>;
236
237impl GLWESwitchingKey<Vec<u8>> {
238 pub fn encrypt_sk_scratch_space<B: Backend>(
239 module: &Module<B>,
240 basek: usize,
241 k: usize,
242 rank_in: usize,
243 rank_out: usize,
244 ) -> usize
245 where
246 Module<B>: GLWESwitchingKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
247 {
248 (GGLWECiphertext::encrypt_sk_scratch_space(module, basek, k) | module.scalar_znx_alloc_bytes(1))
249 + module.scalar_znx_alloc_bytes(rank_in)
250 + GLWESecretExec::bytes_of(module, rank_out)
251 }
252
253 pub fn encrypt_pk_scratch_space<B: Backend>(
254 module: &Module<B>,
255 _basek: usize,
256 _k: usize,
257 _rank_in: usize,
258 _rank_out: usize,
259 ) -> usize {
260 GGLWECiphertext::encrypt_pk_scratch_space(module, _basek, _k, _rank_out)
261 }
262}
263
264impl<DataSelf: DataMut> GLWESwitchingKey<DataSelf> {
265 pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
266 &mut self,
267 module: &Module<B>,
268 sk_in: &GLWESecret<DataSkIn>,
269 sk_out: &GLWESecret<DataSkOut>,
270 source_xa: &mut Source,
271 source_xe: &mut Source,
272 sigma: f64,
273 scratch: &mut Scratch<B>,
274 ) where
275 Module<B>: GLWESwitchingKeyEncryptSkFamily<B>
276 + ScalarZnxAllocBytes
277 + VecZnxSwithcDegree
278 + VecZnxAllocBytes
279 + VecZnxAddScalarInplace,
280 Scratch<B>:
281 ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + ScratchAvailable + TakeVecZnx<B>,
282 {
283 #[cfg(debug_assertions)]
284 {
285 assert!(sk_in.n() <= module.n());
286 assert!(sk_out.n() <= module.n());
287 assert!(
288 scratch.available()
289 >= GLWESwitchingKey::encrypt_sk_scratch_space(
290 module,
291 self.basek(),
292 self.k(),
293 self.rank_in(),
294 self.rank_out()
295 ),
296 "scratch.available()={} < GLWESwitchingKey::encrypt_sk_scratch_space={}",
297 scratch.available(),
298 GLWESwitchingKey::encrypt_sk_scratch_space(
299 module,
300 self.basek(),
301 self.k(),
302 self.rank_in(),
303 self.rank_out()
304 )
305 )
306 }
307
308 let (mut sk_in_tmp, scratch1) = scratch.take_scalar_znx(module, 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(module, sk_out.rank());
319 {
320 let (mut tmp, _) = scratch2.take_scalar_znx(module, 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 basek: usize,
345 k: usize,
346 rank_in: usize,
347 rank_out: usize,
348 ) -> usize
349 where
350 Module<B>: GLWESwitchingKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
351 {
352 (GGLWECiphertext::encrypt_sk_scratch_space(module, basek, k) | module.scalar_znx_alloc_bytes(1))
353 + module.scalar_znx_alloc_bytes(rank_in)
354 + GLWESecretExec::bytes_of(module, rank_out)
355 }
356}
357
358impl<DataSelf: DataMut> GLWESwitchingKeyCompressed<DataSelf> {
359 pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
360 &mut self,
361 module: &Module<B>,
362 sk_in: &GLWESecret<DataSkIn>,
363 sk_out: &GLWESecret<DataSkOut>,
364 seed_xa: [u8; 32],
365 source_xe: &mut Source,
366 sigma: f64,
367 scratch: &mut Scratch<B>,
368 ) where
369 Module<B>: GLWESwitchingKeyEncryptSkFamily<B>
370 + ScalarZnxAllocBytes
371 + VecZnxSwithcDegree
372 + VecZnxAllocBytes
373 + VecZnxAddScalarInplace,
374 Scratch<B>:
375 ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + ScratchAvailable + TakeVecZnx<B>,
376 {
377 #[cfg(debug_assertions)]
378 {
379 assert!(sk_in.n() <= module.n());
380 assert!(sk_out.n() <= module.n());
381 assert!(
382 scratch.available()
383 >= GLWESwitchingKey::encrypt_sk_scratch_space(
384 module,
385 self.basek(),
386 self.k(),
387 self.rank_in(),
388 self.rank_out()
389 ),
390 "scratch.available()={} < GLWESwitchingKey::encrypt_sk_scratch_space={}",
391 scratch.available(),
392 GLWESwitchingKey::encrypt_sk_scratch_space(
393 module,
394 self.basek(),
395 self.k(),
396 self.rank_in(),
397 self.rank_out()
398 )
399 )
400 }
401
402 let (mut sk_in_tmp, scratch1) = scratch.take_scalar_znx(module, 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(module, sk_out.rank());
413 {
414 let (mut tmp, _) = scratch2.take_scalar_znx(module, 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>, basek: usize, k: usize, rank: usize) -> usize
439 where
440 Module<B>: AutomorphismKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
441 {
442 GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank, rank) + GLWESecret::bytes_of(module, rank)
443 }
444
445 pub fn encrypt_pk_scratch_space<B: Backend>(module: &Module<B>, _basek: usize, _k: usize, _rank: usize) -> usize {
446 GLWESwitchingKey::encrypt_pk_scratch_space(module, _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>
462 + ScalarZnxAllocBytes
463 + VecZnxAllocBytes
464 + VecZnxAutomorphism
465 + VecZnxSwithcDegree
466 + VecZnxAddScalarInplace,
467 Scratch<B>: ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx<B>,
468 {
469 #[cfg(debug_assertions)]
470 {
471 assert_eq!(self.n(), module.n());
472 assert_eq!(sk.n(), module.n());
473 assert_eq!(self.rank_out(), self.rank_in());
474 assert_eq!(sk.rank(), self.rank());
475 assert!(
476 scratch.available() >= AutomorphismKey::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank()),
477 "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
478 scratch.available(),
479 self.rank(),
480 self.size(),
481 AutomorphismKey::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank())
482 )
483 }
484
485 let (mut sk_out, scratch_1) = scratch.take_glwe_secret(module, sk.rank());
486
487 {
488 (0..self.rank()).for_each(|i| {
489 module.vec_znx_automorphism(
490 module.galois_element_inv(p),
491 &mut sk_out.data.as_vec_znx_mut(),
492 i,
493 &sk.data.as_vec_znx(),
494 i,
495 );
496 });
497 }
498
499 self.key
500 .encrypt_sk(module, &sk, &sk_out, source_xa, source_xe, sigma, scratch_1);
501
502 self.p = p;
503 }
504}
505
506impl AutomorphismKeyCompressed<Vec<u8>> {
507 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
508 where
509 Module<B>: AutomorphismKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
510 {
511 GLWESwitchingKeyCompressed::encrypt_sk_scratch_space(module, basek, k, rank, rank) + GLWESecret::bytes_of(module, rank)
512 }
513}
514
515impl<DataSelf: DataMut> AutomorphismKeyCompressed<DataSelf> {
516 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
517 &mut self,
518 module: &Module<B>,
519 p: i64,
520 sk: &GLWESecret<DataSk>,
521 seed_xa: [u8; 32],
522 source_xe: &mut Source,
523 sigma: f64,
524 scratch: &mut Scratch<B>,
525 ) where
526 Module<B>: AutomorphismKeyEncryptSkFamily<B>
527 + ScalarZnxAllocBytes
528 + VecZnxAllocBytes
529 + VecZnxSwithcDegree
530 + VecZnxAutomorphism
531 + VecZnxAddScalarInplace,
532 Scratch<B>: ScratchAvailable + TakeScalarZnx<B> + TakeVecZnxDft<B> + TakeGLWESecretExec<B> + TakeVecZnx<B>,
533 {
534 #[cfg(debug_assertions)]
535 {
536 assert_eq!(self.n(), module.n());
537 assert_eq!(sk.n(), module.n());
538 assert_eq!(self.rank_out(), self.rank_in());
539 assert_eq!(sk.rank(), self.rank());
540 assert!(
541 scratch.available()
542 >= AutomorphismKeyCompressed::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank()),
543 "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space(module, self.rank()={}, self.size()={}): {}",
544 scratch.available(),
545 self.rank(),
546 self.size(),
547 AutomorphismKeyCompressed::encrypt_sk_scratch_space(module, self.basek(), self.k(), self.rank())
548 )
549 }
550
551 let (mut sk_out, scratch_1) = scratch.take_glwe_secret(module, sk.rank());
552
553 {
554 (0..self.rank()).for_each(|i| {
555 module.vec_znx_automorphism(
556 module.galois_element_inv(p),
557 &mut sk_out.data.as_vec_znx_mut(),
558 i,
559 &sk.data.as_vec_znx(),
560 i,
561 );
562 });
563 }
564
565 self.key
566 .encrypt_sk(module, &sk, &sk_out, seed_xa, source_xe, sigma, scratch_1);
567
568 self.p = p;
569 }
570}
571
572pub trait GLWETensorKeyEncryptSkFamily<B: Backend> =
573 GGLWEEncryptSkFamily<B> + VecZnxBigAllocBytes + VecZnxDftToVecZnxBigTmpA<B> + SvpApply<B>;
574
575impl GLWETensorKey<Vec<u8>> {
576 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
577 where
578 Module<B>: GLWETensorKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
579 {
580 GLWESecretExec::bytes_of(module, rank)
581 + module.vec_znx_dft_alloc_bytes(rank, 1)
582 + module.vec_znx_big_alloc_bytes(1, 1)
583 + module.vec_znx_dft_alloc_bytes(1, 1)
584 + GLWESecret::bytes_of(module, 1)
585 + GLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k, rank, rank)
586 }
587}
588
589impl<DataSelf: DataMut> GLWETensorKey<DataSelf> {
590 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
591 &mut self,
592 module: &Module<B>,
593 sk: &GLWESecret<DataSk>,
594 source_xa: &mut Source,
595 source_xe: &mut Source,
596 sigma: f64,
597 scratch: &mut Scratch<B>,
598 ) where
599 Module<B>: GLWETensorKeyEncryptSkFamily<B>
600 + ScalarZnxAllocBytes
601 + VecZnxSwithcDegree
602 + VecZnxAllocBytes
603 + VecZnxAddScalarInplace,
604 Scratch<B>:
605 ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecretExec<B> + TakeScalarZnx<B> + TakeVecZnx<B>,
606 {
607 #[cfg(debug_assertions)]
608 {
609 assert_eq!(self.rank(), sk.rank());
610 assert_eq!(self.n(), module.n());
611 assert_eq!(sk.n(), module.n());
612 }
613
614 let rank: usize = self.rank();
615
616 let (mut sk_dft_prep, scratch1) = scratch.take_glwe_secret_exec(module, rank);
617 sk_dft_prep.prepare(module, &sk);
618
619 let (mut sk_dft, scratch2) = scratch1.take_vec_znx_dft(module, rank, 1);
620
621 (0..rank).for_each(|i| {
622 module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
623 });
624
625 let (mut sk_ij_big, scratch3) = scratch2.take_vec_znx_big(module, 1, 1);
626 let (mut sk_ij, scratch4) = scratch3.take_glwe_secret(module, 1);
627 let (mut sk_ij_dft, scratch5) = scratch4.take_vec_znx_dft(module, 1, 1);
628
629 (0..rank).for_each(|i| {
630 (i..rank).for_each(|j| {
631 module.svp_apply(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
632
633 module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
634 module.vec_znx_big_normalize(
635 self.basek(),
636 &mut sk_ij.data.as_vec_znx_mut(),
637 0,
638 &sk_ij_big,
639 0,
640 scratch5,
641 );
642
643 self.at_mut(i, j)
644 .encrypt_sk(module, &sk_ij, sk, source_xa, source_xe, sigma, scratch5);
645 });
646 })
647 }
648}
649
650impl GLWETensorKeyCompressed<Vec<u8>> {
651 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
652 where
653 Module<B>: GLWETensorKeyEncryptSkFamily<B> + ScalarZnxAllocBytes + VecZnxAllocBytes,
654 {
655 GLWETensorKey::encrypt_sk_scratch_space(module, basek, k, rank)
656 }
657}
658
659impl<DataSelf: DataMut> GLWETensorKeyCompressed<DataSelf> {
660 pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
661 &mut self,
662 module: &Module<B>,
663 sk: &GLWESecret<DataSk>,
664 seed_xa: [u8; 32],
665 source_xe: &mut Source,
666 sigma: f64,
667 scratch: &mut Scratch<B>,
668 ) where
669 Module<B>: GLWETensorKeyEncryptSkFamily<B>
670 + ScalarZnxAllocBytes
671 + VecZnxSwithcDegree
672 + VecZnxAllocBytes
673 + VecZnxAddScalarInplace,
674 Scratch<B>:
675 ScratchAvailable + TakeVecZnxDft<B> + TakeVecZnxBig<B> + TakeGLWESecretExec<B> + TakeScalarZnx<B> + TakeVecZnx<B>,
676 {
677 #[cfg(debug_assertions)]
678 {
679 assert_eq!(self.rank(), sk.rank());
680 assert_eq!(self.n(), module.n());
681 assert_eq!(sk.n(), module.n());
682 }
683
684 let rank: usize = self.rank();
685
686 let (mut sk_dft_prep, scratch1) = scratch.take_glwe_secret_exec(module, rank);
687 sk_dft_prep.prepare(module, &sk);
688
689 let (mut sk_dft, scratch2) = scratch1.take_vec_znx_dft(module, rank, 1);
690
691 (0..rank).for_each(|i| {
692 module.vec_znx_dft_from_vec_znx(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
693 });
694
695 let (mut sk_ij_big, scratch3) = scratch2.take_vec_znx_big(module, 1, 1);
696 let (mut sk_ij, scratch4) = scratch3.take_glwe_secret(module, 1);
697 let (mut sk_ij_dft, scratch5) = scratch4.take_vec_znx_dft(module, 1, 1);
698
699 let mut source_xa: Source = Source::new(seed_xa);
700
701 (0..rank).for_each(|i| {
702 (i..rank).for_each(|j| {
703 module.svp_apply(&mut sk_ij_dft, 0, &sk_dft_prep.data, j, &sk_dft, i);
704
705 module.vec_znx_dft_to_vec_znx_big_tmp_a(&mut sk_ij_big, 0, &mut sk_ij_dft, 0);
706 module.vec_znx_big_normalize(
707 self.basek(),
708 &mut sk_ij.data.as_vec_znx_mut(),
709 0,
710 &sk_ij_big,
711 0,
712 scratch5,
713 );
714
715 let (seed_xa_tmp, _) = source_xa.branch();
716
717 self.at_mut(i, j)
718 .encrypt_sk(module, &sk_ij, sk, seed_xa_tmp, source_xe, sigma, scratch5);
719 });
720 })
721 }
722}