poulpy/core/src/scratch.rs
core/src/scratch.rsRaw
1use backend::hal::{
2 api::{TakeMatZnx, TakeScalarZnx, TakeSvpPPol, TakeVecZnx, TakeVecZnxDft, TakeVmpPMat},
3 layouts::{Backend, DataRef, Module, Scratch},
4 oep::{TakeMatZnxImpl, TakeScalarZnxImpl, TakeSvpPPolImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, TakeVmpPMatImpl},
5};
6
7use crate::{
8 AutomorphismKey, AutomorphismKeyExec, GGLWECiphertext, GGLWECiphertextExec, GGSWCiphertext, GGSWCiphertextExec,
9 GLWECiphertext, GLWEPlaintext, GLWEPublicKey, GLWEPublicKeyExec, GLWESecret, GLWESecretExec, GLWESwitchingKey,
10 GLWESwitchingKeyExec, GLWETensorKey, GLWETensorKeyExec, Infos, dist::Distribution,
11};
12
13pub trait TakeLike<'a, B: Backend, T> {
14 type Output;
15 fn take_like(&'a mut self, template: &T) -> (Self::Output, &'a mut Self);
16}
17
18pub trait TakeGLWECt<B: Backend> {
19 fn take_glwe_ct(&mut self, module: &Module<B>, basek: usize, k: usize, rank: usize)
20 -> (GLWECiphertext<&mut [u8]>, &mut Self);
21}
22
23pub trait TakeGLWECtSlice<B: Backend> {
24 fn take_glwe_ct_slice(
25 &mut self,
26 size: usize,
27 module: &Module<B>,
28 basek: usize,
29 k: usize,
30 rank: usize,
31 ) -> (Vec<GLWECiphertext<&mut [u8]>>, &mut Self);
32}
33
34pub trait TakeGLWEPt<B: Backend> {
35 fn take_glwe_pt(&mut self, module: &Module<B>, basek: usize, k: usize) -> (GLWEPlaintext<&mut [u8]>, &mut Self);
36}
37
38pub trait TakeGGLWE<B: Backend> {
39 fn take_gglwe(
40 &mut self,
41 module: &Module<B>,
42 basek: usize,
43 k: usize,
44 rows: usize,
45 digits: usize,
46 rank_in: usize,
47 rank_out: usize,
48 ) -> (GGLWECiphertext<&mut [u8]>, &mut Self);
49}
50
51pub trait TakeGGLWEExec<B: Backend> {
52 fn take_gglwe_exec(
53 &mut self,
54 module: &Module<B>,
55 basek: usize,
56 k: usize,
57 rows: usize,
58 digits: usize,
59 rank_in: usize,
60 rank_out: usize,
61 ) -> (GGLWECiphertextExec<&mut [u8], B>, &mut Self);
62}
63
64pub trait TakeGGSW<B: Backend> {
65 fn take_ggsw(
66 &mut self,
67 module: &Module<B>,
68 basek: usize,
69 k: usize,
70 rows: usize,
71 digits: usize,
72 rank: usize,
73 ) -> (GGSWCiphertext<&mut [u8]>, &mut Self);
74}
75
76pub trait TakeGGSWExec<B: Backend> {
77 fn take_ggsw_exec(
78 &mut self,
79 module: &Module<B>,
80 basek: usize,
81 k: usize,
82 rows: usize,
83 digits: usize,
84 rank: usize,
85 ) -> (GGSWCiphertextExec<&mut [u8], B>, &mut Self);
86}
87
88pub trait TakeGLWESecret<B: Backend> {
89 fn take_glwe_secret(&mut self, module: &Module<B>, rank: usize) -> (GLWESecret<&mut [u8]>, &mut Self);
90}
91
92pub trait TakeGLWESecretExec<B: Backend> {
93 fn take_glwe_secret_exec(&mut self, module: &Module<B>, rank: usize) -> (GLWESecretExec<&mut [u8], B>, &mut Self);
94}
95
96pub trait TakeGLWEPk<B: Backend> {
97 fn take_glwe_pk(&mut self, module: &Module<B>, basek: usize, k: usize, rank: usize) -> (GLWEPublicKey<&mut [u8]>, &mut Self);
98}
99
100pub trait TakeGLWEPkExec<B: Backend> {
101 fn take_glwe_pk_exec(
102 &mut self,
103 module: &Module<B>,
104 basek: usize,
105 k: usize,
106 rank: usize,
107 ) -> (GLWEPublicKeyExec<&mut [u8], B>, &mut Self);
108}
109
110pub trait TakeGLWESwitchingKey<B: Backend> {
111 fn take_glwe_switching_key(
112 &mut self,
113 module: &Module<B>,
114 basek: usize,
115 k: usize,
116 rows: usize,
117 digits: usize,
118 rank_in: usize,
119 rank_out: usize,
120 ) -> (GLWESwitchingKey<&mut [u8]>, &mut Self);
121}
122
123pub trait TakeGLWESwitchingKeyExec<B: Backend> {
124 fn take_glwe_switching_key_exec(
125 &mut self,
126 module: &Module<B>,
127 basek: usize,
128 k: usize,
129 rows: usize,
130 digits: usize,
131 rank_in: usize,
132 rank_out: usize,
133 ) -> (GLWESwitchingKeyExec<&mut [u8], B>, &mut Self);
134}
135
136pub trait TakeTensorKey<B: Backend> {
137 fn take_tensor_key(
138 &mut self,
139 module: &Module<B>,
140 basek: usize,
141 k: usize,
142 rows: usize,
143 digits: usize,
144 rank: usize,
145 ) -> (GLWETensorKey<&mut [u8]>, &mut Self);
146}
147
148pub trait TakeTensorKeyExec<B: Backend> {
149 fn take_tensor_key_exec(
150 &mut self,
151 module: &Module<B>,
152 basek: usize,
153 k: usize,
154 rows: usize,
155 digits: usize,
156 rank: usize,
157 ) -> (GLWETensorKeyExec<&mut [u8], B>, &mut Self);
158}
159
160pub trait TakeAutomorphismKey<B: Backend> {
161 fn take_automorphism_key(
162 &mut self,
163 module: &Module<B>,
164 basek: usize,
165 k: usize,
166 rows: usize,
167 digits: usize,
168 rank: usize,
169 ) -> (AutomorphismKey<&mut [u8]>, &mut Self);
170}
171
172pub trait TakeAutomorphismKeyExec<B: Backend> {
173 fn take_automorphism_key_exec(
174 &mut self,
175 module: &Module<B>,
176 basek: usize,
177 k: usize,
178 rows: usize,
179 digits: usize,
180 rank: usize,
181 ) -> (AutomorphismKeyExec<&mut [u8], B>, &mut Self);
182}
183
184impl<B: Backend> TakeGLWECt<B> for Scratch<B>
185where
186 Scratch<B>: TakeVecZnx<B>,
187{
188 fn take_glwe_ct(
189 &mut self,
190 module: &Module<B>,
191 basek: usize,
192 k: usize,
193 rank: usize,
194 ) -> (GLWECiphertext<&mut [u8]>, &mut Self) {
195 let (data, scratch) = self.take_vec_znx(module, rank + 1, k.div_ceil(basek));
196 (GLWECiphertext { data, basek, k }, scratch)
197 }
198}
199
200impl<'a, B, D> TakeLike<'a, B, GLWECiphertext<D>> for Scratch<B>
201where
202 B: Backend + TakeVecZnxImpl<B>,
203 D: DataRef,
204{
205 type Output = GLWECiphertext<&'a mut [u8]>;
206
207 fn take_like(&'a mut self, template: &GLWECiphertext<D>) -> (Self::Output, &'a mut Self) {
208 let (data, scratch) = B::take_vec_znx_impl(self, template.n(), template.cols(), template.size());
209 (
210 GLWECiphertext {
211 data,
212 basek: template.basek(),
213 k: template.k(),
214 },
215 scratch,
216 )
217 }
218}
219
220impl<B: Backend> TakeGLWECtSlice<B> for Scratch<B>
221where
222 Scratch<B>: TakeVecZnx<B>,
223{
224 fn take_glwe_ct_slice(
225 &mut self,
226 size: usize,
227 module: &Module<B>,
228 basek: usize,
229 k: usize,
230 rank: usize,
231 ) -> (Vec<GLWECiphertext<&mut [u8]>>, &mut Self) {
232 let mut scratch: &mut Scratch<B> = self;
233 let mut cts: Vec<GLWECiphertext<&mut [u8]>> = Vec::with_capacity(size);
234 for _ in 0..size {
235 let (ct, new_scratch) = scratch.take_glwe_ct(module, basek, k, rank);
236 scratch = new_scratch;
237 cts.push(ct);
238 }
239 (cts, scratch)
240 }
241}
242
243impl<B: Backend> TakeGLWEPt<B> for Scratch<B>
244where
245 Scratch<B>: TakeVecZnx<B>,
246{
247 fn take_glwe_pt(&mut self, module: &Module<B>, basek: usize, k: usize) -> (GLWEPlaintext<&mut [u8]>, &mut Self) {
248 let (data, scratch) = self.take_vec_znx(module, 1, k.div_ceil(basek));
249 (GLWEPlaintext { data, basek, k }, scratch)
250 }
251}
252
253impl<'a, B, D> TakeLike<'a, B, GLWEPlaintext<D>> for Scratch<B>
254where
255 B: Backend + TakeVecZnxImpl<B>,
256 D: DataRef,
257{
258 type Output = GLWEPlaintext<&'a mut [u8]>;
259
260 fn take_like(&'a mut self, template: &GLWEPlaintext<D>) -> (Self::Output, &'a mut Self) {
261 let (data, scratch) = B::take_vec_znx_impl(self, template.n(), template.cols(), template.size());
262 (
263 GLWEPlaintext {
264 data,
265 basek: template.basek(),
266 k: template.k(),
267 },
268 scratch,
269 )
270 }
271}
272
273impl<B: Backend> TakeGGLWE<B> for Scratch<B>
274where
275 Scratch<B>: TakeMatZnx<B>,
276{
277 fn take_gglwe(
278 &mut self,
279 module: &Module<B>,
280 basek: usize,
281 k: usize,
282 rows: usize,
283 digits: usize,
284 rank_in: usize,
285 rank_out: usize,
286 ) -> (GGLWECiphertext<&mut [u8]>, &mut Self) {
287 let (data, scratch) = self.take_mat_znx(
288 module,
289 rows.div_ceil(digits),
290 rank_in,
291 rank_out + 1,
292 k.div_ceil(basek),
293 );
294 (
295 GGLWECiphertext {
296 data: data,
297 basek: basek,
298 k,
299 digits,
300 },
301 scratch,
302 )
303 }
304}
305
306impl<'a, B, D> TakeLike<'a, B, GGLWECiphertext<D>> for Scratch<B>
307where
308 B: Backend + TakeMatZnxImpl<B>,
309 D: DataRef,
310{
311 type Output = GGLWECiphertext<&'a mut [u8]>;
312
313 fn take_like(&'a mut self, template: &GGLWECiphertext<D>) -> (Self::Output, &'a mut Self) {
314 let (data, scratch) = B::take_mat_znx_impl(
315 self,
316 template.n(),
317 template.rows(),
318 template.data.cols_in(),
319 template.data.cols_out(),
320 template.size(),
321 );
322 (
323 GGLWECiphertext {
324 data,
325 basek: template.basek(),
326 k: template.k(),
327 digits: template.digits(),
328 },
329 scratch,
330 )
331 }
332}
333
334impl<B: Backend> TakeGGLWEExec<B> for Scratch<B>
335where
336 Scratch<B>: TakeVmpPMat<B>,
337{
338 fn take_gglwe_exec(
339 &mut self,
340 module: &Module<B>,
341 basek: usize,
342 k: usize,
343 rows: usize,
344 digits: usize,
345 rank_in: usize,
346 rank_out: usize,
347 ) -> (GGLWECiphertextExec<&mut [u8], B>, &mut Self) {
348 let (data, scratch) = self.take_vmp_pmat(
349 module,
350 rows.div_ceil(digits),
351 rank_in,
352 rank_out + 1,
353 k.div_ceil(basek),
354 );
355 (
356 GGLWECiphertextExec {
357 data: data,
358 basek: basek,
359 k,
360 digits,
361 },
362 scratch,
363 )
364 }
365}
366
367impl<'a, B, D> TakeLike<'a, B, GGLWECiphertextExec<D, B>> for Scratch<B>
368where
369 B: Backend + TakeVmpPMatImpl<B>,
370 D: DataRef,
371{
372 type Output = GGLWECiphertextExec<&'a mut [u8], B>;
373
374 fn take_like(&'a mut self, template: &GGLWECiphertextExec<D, B>) -> (Self::Output, &'a mut Self) {
375 let (data, scratch) = B::take_vmp_pmat_impl(
376 self,
377 template.n(),
378 template.rows(),
379 template.data.cols_in(),
380 template.data.cols_out(),
381 template.size(),
382 );
383 (
384 GGLWECiphertextExec {
385 data,
386 basek: template.basek(),
387 k: template.k(),
388 digits: template.digits(),
389 },
390 scratch,
391 )
392 }
393}
394
395impl<B: Backend> TakeGGSW<B> for Scratch<B>
396where
397 Scratch<B>: TakeMatZnx<B>,
398{
399 fn take_ggsw(
400 &mut self,
401 module: &Module<B>,
402 basek: usize,
403 k: usize,
404 rows: usize,
405 digits: usize,
406 rank: usize,
407 ) -> (GGSWCiphertext<&mut [u8]>, &mut Self) {
408 let (data, scratch) = self.take_mat_znx(
409 module,
410 rows.div_ceil(digits),
411 rank + 1,
412 rank + 1,
413 k.div_ceil(basek),
414 );
415 (
416 GGSWCiphertext {
417 data,
418 basek,
419 k,
420 digits,
421 },
422 scratch,
423 )
424 }
425}
426
427impl<'a, B, D> TakeLike<'a, B, GGSWCiphertext<D>> for Scratch<B>
428where
429 B: Backend + TakeMatZnxImpl<B>,
430 D: DataRef,
431{
432 type Output = GGSWCiphertext<&'a mut [u8]>;
433
434 fn take_like(&'a mut self, template: &GGSWCiphertext<D>) -> (Self::Output, &'a mut Self) {
435 let (data, scratch) = B::take_mat_znx_impl(
436 self,
437 template.n(),
438 template.rows(),
439 template.data.cols_in(),
440 template.data.cols_out(),
441 template.size(),
442 );
443 (
444 GGSWCiphertext {
445 data,
446 basek: template.basek(),
447 k: template.k(),
448 digits: template.digits(),
449 },
450 scratch,
451 )
452 }
453}
454
455impl<B: Backend> TakeGGSWExec<B> for Scratch<B>
456where
457 Scratch<B>: TakeVmpPMat<B>,
458{
459 fn take_ggsw_exec(
460 &mut self,
461 module: &Module<B>,
462 basek: usize,
463 k: usize,
464 rows: usize,
465 digits: usize,
466 rank: usize,
467 ) -> (GGSWCiphertextExec<&mut [u8], B>, &mut Self) {
468 let (data, scratch) = self.take_vmp_pmat(
469 module,
470 rows.div_ceil(digits),
471 rank + 1,
472 rank + 1,
473 k.div_ceil(basek),
474 );
475 (
476 GGSWCiphertextExec {
477 data,
478 basek,
479 k,
480 digits,
481 },
482 scratch,
483 )
484 }
485}
486
487impl<'a, B, D> TakeLike<'a, B, GGSWCiphertextExec<D, B>> for Scratch<B>
488where
489 B: Backend + TakeVmpPMatImpl<B>,
490 D: DataRef,
491{
492 type Output = GGSWCiphertextExec<&'a mut [u8], B>;
493
494 fn take_like(&'a mut self, template: &GGSWCiphertextExec<D, B>) -> (Self::Output, &'a mut Self) {
495 let (data, scratch) = B::take_vmp_pmat_impl(
496 self,
497 template.n(),
498 template.rows(),
499 template.data.cols_in(),
500 template.data.cols_out(),
501 template.size(),
502 );
503 (
504 GGSWCiphertextExec {
505 data,
506 basek: template.basek(),
507 k: template.k(),
508 digits: template.digits(),
509 },
510 scratch,
511 )
512 }
513}
514
515impl<B: Backend> TakeGLWEPk<B> for Scratch<B>
516where
517 Scratch<B>: TakeVecZnx<B>,
518{
519 fn take_glwe_pk(&mut self, module: &Module<B>, basek: usize, k: usize, rank: usize) -> (GLWEPublicKey<&mut [u8]>, &mut Self) {
520 let (data, scratch) = self.take_vec_znx(module, rank + 1, k.div_ceil(basek));
521 (
522 GLWEPublicKey {
523 data,
524 k,
525 basek,
526 dist: Distribution::NONE,
527 },
528 scratch,
529 )
530 }
531}
532
533impl<'a, B, D> TakeLike<'a, B, GLWEPublicKey<D>> for Scratch<B>
534where
535 B: Backend + TakeVecZnxImpl<B>,
536 D: DataRef,
537{
538 type Output = GLWEPublicKey<&'a mut [u8]>;
539
540 fn take_like(&'a mut self, template: &GLWEPublicKey<D>) -> (Self::Output, &'a mut Self) {
541 let (data, scratch) = B::take_vec_znx_impl(self, template.n(), template.cols(), template.size());
542 (
543 GLWEPublicKey {
544 data,
545 basek: template.basek(),
546 k: template.k(),
547 dist: template.dist,
548 },
549 scratch,
550 )
551 }
552}
553
554impl<B: Backend> TakeGLWEPkExec<B> for Scratch<B>
555where
556 Scratch<B>: TakeVecZnxDft<B>,
557{
558 fn take_glwe_pk_exec(
559 &mut self,
560 module: &Module<B>,
561 basek: usize,
562 k: usize,
563 rank: usize,
564 ) -> (GLWEPublicKeyExec<&mut [u8], B>, &mut Self) {
565 let (data, scratch) = self.take_vec_znx_dft(module, rank + 1, k.div_ceil(basek));
566 (
567 GLWEPublicKeyExec {
568 data,
569 k,
570 basek,
571 dist: Distribution::NONE,
572 },
573 scratch,
574 )
575 }
576}
577
578impl<'a, B, D> TakeLike<'a, B, GLWEPublicKeyExec<D, B>> for Scratch<B>
579where
580 B: Backend + TakeVecZnxDftImpl<B>,
581 D: DataRef,
582{
583 type Output = GLWEPublicKeyExec<&'a mut [u8], B>;
584
585 fn take_like(&'a mut self, template: &GLWEPublicKeyExec<D, B>) -> (Self::Output, &'a mut Self) {
586 let (data, scratch) = B::take_vec_znx_dft_impl(self, template.n(), template.cols(), template.size());
587 (
588 GLWEPublicKeyExec {
589 data,
590 basek: template.basek(),
591 k: template.k(),
592 dist: template.dist,
593 },
594 scratch,
595 )
596 }
597}
598
599impl<B: Backend> TakeGLWESecret<B> for Scratch<B>
600where
601 Scratch<B>: TakeScalarZnx<B>,
602{
603 fn take_glwe_secret(&mut self, module: &Module<B>, rank: usize) -> (GLWESecret<&mut [u8]>, &mut Self) {
604 let (data, scratch) = self.take_scalar_znx(module, rank);
605 (
606 GLWESecret {
607 data,
608 dist: Distribution::NONE,
609 },
610 scratch,
611 )
612 }
613}
614
615impl<'a, B, D> TakeLike<'a, B, GLWESecret<D>> for Scratch<B>
616where
617 B: Backend + TakeScalarZnxImpl<B>,
618 D: DataRef,
619{
620 type Output = GLWESecret<&'a mut [u8]>;
621
622 fn take_like(&'a mut self, template: &GLWESecret<D>) -> (Self::Output, &'a mut Self) {
623 let (data, scratch) = B::take_scalar_znx_impl(self, template.n(), template.rank());
624 (
625 GLWESecret {
626 data,
627 dist: template.dist,
628 },
629 scratch,
630 )
631 }
632}
633
634impl<B: Backend> TakeGLWESecretExec<B> for Scratch<B>
635where
636 Scratch<B>: TakeSvpPPol<B>,
637{
638 fn take_glwe_secret_exec(&mut self, module: &Module<B>, rank: usize) -> (GLWESecretExec<&mut [u8], B>, &mut Self) {
639 let (data, scratch) = self.take_svp_ppol(module, rank);
640 (
641 GLWESecretExec {
642 data,
643 dist: Distribution::NONE,
644 },
645 scratch,
646 )
647 }
648}
649
650impl<'a, B, D> TakeLike<'a, B, GLWESecretExec<D, B>> for Scratch<B>
651where
652 B: Backend + TakeSvpPPolImpl<B>,
653 D: DataRef,
654{
655 type Output = GLWESecretExec<&'a mut [u8], B>;
656
657 fn take_like(&'a mut self, template: &GLWESecretExec<D, B>) -> (Self::Output, &'a mut Self) {
658 let (data, scratch) = B::take_svp_ppol_impl(self, template.n(), template.rank());
659 (
660 GLWESecretExec {
661 data,
662 dist: template.dist,
663 },
664 scratch,
665 )
666 }
667}
668
669impl<B: Backend> TakeGLWESwitchingKey<B> for Scratch<B>
670where
671 Scratch<B>: TakeMatZnx<B>,
672{
673 fn take_glwe_switching_key(
674 &mut self,
675 module: &Module<B>,
676 basek: usize,
677 k: usize,
678 rows: usize,
679 digits: usize,
680 rank_in: usize,
681 rank_out: usize,
682 ) -> (GLWESwitchingKey<&mut [u8]>, &mut Self) {
683 let (data, scratch) = self.take_gglwe(module, basek, k, rows, digits, rank_in, rank_out);
684 (
685 GLWESwitchingKey {
686 key: data,
687 sk_in_n: 0,
688 sk_out_n: 0,
689 },
690 scratch,
691 )
692 }
693}
694
695impl<'a, B, D> TakeLike<'a, B, GLWESwitchingKey<D>> for Scratch<B>
696where
697 Scratch<B>: TakeLike<'a, B, GGLWECiphertext<D>, Output = GGLWECiphertext<&'a mut [u8]>>,
698 B: Backend + TakeMatZnxImpl<B>,
699 D: DataRef,
700{
701 type Output = GLWESwitchingKey<&'a mut [u8]>;
702
703 fn take_like(&'a mut self, template: &GLWESwitchingKey<D>) -> (Self::Output, &'a mut Self) {
704 let (key, scratch) = self.take_like(&template.key);
705 (
706 GLWESwitchingKey {
707 key,
708 sk_in_n: template.sk_in_n,
709 sk_out_n: template.sk_out_n,
710 },
711 scratch,
712 )
713 }
714}
715
716impl<B: Backend> TakeGLWESwitchingKeyExec<B> for Scratch<B>
717where
718 Scratch<B>: TakeGGLWEExec<B>,
719{
720 fn take_glwe_switching_key_exec(
721 &mut self,
722 module: &Module<B>,
723 basek: usize,
724 k: usize,
725 rows: usize,
726 digits: usize,
727 rank_in: usize,
728 rank_out: usize,
729 ) -> (GLWESwitchingKeyExec<&mut [u8], B>, &mut Self) {
730 let (data, scratch) = self.take_gglwe_exec(module, basek, k, rows, digits, rank_in, rank_out);
731 (
732 GLWESwitchingKeyExec {
733 key: data,
734 sk_in_n: 0,
735 sk_out_n: 0,
736 },
737 scratch,
738 )
739 }
740}
741
742impl<'a, B, D> TakeLike<'a, B, GLWESwitchingKeyExec<D, B>> for Scratch<B>
743where
744 Scratch<B>: TakeLike<'a, B, GGLWECiphertextExec<D, B>, Output = GGLWECiphertextExec<&'a mut [u8], B>>,
745 B: Backend + TakeMatZnxImpl<B>,
746 D: DataRef,
747{
748 type Output = GLWESwitchingKeyExec<&'a mut [u8], B>;
749
750 fn take_like(&'a mut self, template: &GLWESwitchingKeyExec<D, B>) -> (Self::Output, &'a mut Self) {
751 let (key, scratch) = self.take_like(&template.key);
752 (
753 GLWESwitchingKeyExec {
754 key,
755 sk_in_n: template.sk_in_n,
756 sk_out_n: template.sk_out_n,
757 },
758 scratch,
759 )
760 }
761}
762
763impl<B: Backend> TakeAutomorphismKey<B> for Scratch<B>
764where
765 Scratch<B>: TakeMatZnx<B>,
766{
767 fn take_automorphism_key(
768 &mut self,
769 module: &Module<B>,
770 basek: usize,
771 k: usize,
772 rows: usize,
773 digits: usize,
774 rank: usize,
775 ) -> (AutomorphismKey<&mut [u8]>, &mut Self) {
776 let (data, scratch) = self.take_glwe_switching_key(module, basek, k, rows, digits, rank, rank);
777 (AutomorphismKey { key: data, p: 0 }, scratch)
778 }
779}
780
781impl<'a, B, D> TakeLike<'a, B, AutomorphismKey<D>> for Scratch<B>
782where
783 Scratch<B>: TakeLike<'a, B, GLWESwitchingKey<D>, Output = GLWESwitchingKey<&'a mut [u8]>>,
784 B: Backend + TakeMatZnxImpl<B>,
785 D: DataRef,
786{
787 type Output = AutomorphismKey<&'a mut [u8]>;
788
789 fn take_like(&'a mut self, template: &AutomorphismKey<D>) -> (Self::Output, &'a mut Self) {
790 let (key, scratch) = self.take_like(&template.key);
791 (AutomorphismKey { key, p: template.p }, scratch)
792 }
793}
794
795impl<B: Backend> TakeAutomorphismKeyExec<B> for Scratch<B>
796where
797 Scratch<B>: TakeGLWESwitchingKeyExec<B>,
798{
799 fn take_automorphism_key_exec(
800 &mut self,
801 module: &Module<B>,
802 basek: usize,
803 k: usize,
804 rows: usize,
805 digits: usize,
806 rank: usize,
807 ) -> (AutomorphismKeyExec<&mut [u8], B>, &mut Self) {
808 let (data, scratch) = self.take_glwe_switching_key_exec(module, basek, k, rows, digits, rank, rank);
809 (AutomorphismKeyExec { key: data, p: 0 }, scratch)
810 }
811}
812
813impl<'a, B, D> TakeLike<'a, B, AutomorphismKeyExec<D, B>> for Scratch<B>
814where
815 Scratch<B>: TakeLike<'a, B, GLWESwitchingKeyExec<D, B>, Output = GLWESwitchingKeyExec<&'a mut [u8], B>>,
816 B: Backend + TakeMatZnxImpl<B>,
817 D: DataRef,
818{
819 type Output = AutomorphismKeyExec<&'a mut [u8], B>;
820
821 fn take_like(&'a mut self, template: &AutomorphismKeyExec<D, B>) -> (Self::Output, &'a mut Self) {
822 let (key, scratch) = self.take_like(&template.key);
823 (AutomorphismKeyExec { key, p: template.p }, scratch)
824 }
825}
826
827impl<B: Backend> TakeTensorKey<B> for Scratch<B>
828where
829 Scratch<B>: TakeMatZnx<B>,
830{
831 fn take_tensor_key(
832 &mut self,
833 module: &Module<B>,
834 basek: usize,
835 k: usize,
836 rows: usize,
837 digits: usize,
838 rank: usize,
839 ) -> (GLWETensorKey<&mut [u8]>, &mut Self) {
840 let mut keys: Vec<GLWESwitchingKey<&mut [u8]>> = Vec::new();
841 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
842
843 let mut scratch: &mut Scratch<B> = self;
844
845 if pairs != 0 {
846 let (gglwe, s) = scratch.take_glwe_switching_key(module, basek, k, rows, digits, 1, rank);
847 scratch = s;
848 keys.push(gglwe);
849 }
850 for _ in 1..pairs {
851 let (gglwe, s) = scratch.take_glwe_switching_key(module, basek, k, rows, digits, 1, rank);
852 scratch = s;
853 keys.push(gglwe);
854 }
855 (GLWETensorKey { keys }, scratch)
856 }
857}
858
859impl<'a, B, D> TakeLike<'a, B, GLWETensorKey<D>> for Scratch<B>
860where
861 Scratch<B>: TakeLike<'a, B, GLWESwitchingKey<D>, Output = GLWESwitchingKey<&'a mut [u8]>>,
862 B: Backend + TakeMatZnxImpl<B>,
863 D: DataRef,
864{
865 type Output = GLWETensorKey<&'a mut [u8]>;
866
867 fn take_like(&'a mut self, template: &GLWETensorKey<D>) -> (Self::Output, &'a mut Self) {
868 let mut keys: Vec<GLWESwitchingKey<&mut [u8]>> = Vec::new();
869 let pairs: usize = template.keys.len();
870
871 let mut scratch: &mut Scratch<B> = self;
872
873 if pairs != 0 {
874 let (gglwe, s) = scratch.take_like(template.at(0, 0));
875 scratch = s;
876 keys.push(gglwe);
877 }
878 for _ in 1..pairs {
879 let (gglwe, s) = scratch.take_like(template.at(0, 0));
880 scratch = s;
881 keys.push(gglwe);
882 }
883
884 (GLWETensorKey { keys }, scratch)
885 }
886}
887
888impl<B: Backend> TakeTensorKeyExec<B> for Scratch<B>
889where
890 Scratch<B>: TakeVmpPMat<B>,
891{
892 fn take_tensor_key_exec(
893 &mut self,
894 module: &Module<B>,
895 basek: usize,
896 k: usize,
897 rows: usize,
898 digits: usize,
899 rank: usize,
900 ) -> (GLWETensorKeyExec<&mut [u8], B>, &mut Self) {
901 let mut keys: Vec<GLWESwitchingKeyExec<&mut [u8], B>> = Vec::new();
902 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
903
904 let mut scratch: &mut Scratch<B> = self;
905
906 if pairs != 0 {
907 let (gglwe, s) = scratch.take_glwe_switching_key_exec(module, basek, k, rows, digits, 1, rank);
908 scratch = s;
909 keys.push(gglwe);
910 }
911 for _ in 1..pairs {
912 let (gglwe, s) = scratch.take_glwe_switching_key_exec(module, basek, k, rows, digits, 1, rank);
913 scratch = s;
914 keys.push(gglwe);
915 }
916 (GLWETensorKeyExec { keys }, scratch)
917 }
918}
919
920impl<'a, B, D> TakeLike<'a, B, GLWETensorKeyExec<D, B>> for Scratch<B>
921where
922 Scratch<B>: TakeLike<'a, B, GLWESwitchingKeyExec<D, B>, Output = GLWESwitchingKeyExec<&'a mut [u8], B>>,
923 B: Backend + TakeMatZnxImpl<B>,
924 D: DataRef,
925{
926 type Output = GLWETensorKeyExec<&'a mut [u8], B>;
927
928 fn take_like(&'a mut self, template: &GLWETensorKeyExec<D, B>) -> (Self::Output, &'a mut Self) {
929 let mut keys: Vec<GLWESwitchingKeyExec<&mut [u8], B>> = Vec::new();
930 let pairs: usize = template.keys.len();
931
932 let mut scratch: &mut Scratch<B> = self;
933
934 if pairs != 0 {
935 let (gglwe, s) = scratch.take_like(template.at(0, 0));
936 scratch = s;
937 keys.push(gglwe);
938 }
939 for _ in 1..pairs {
940 let (gglwe, s) = scratch.take_like(template.at(0, 0));
941 scratch = s;
942 keys.push(gglwe);
943 }
944
945 (GLWETensorKeyExec { keys }, scratch)
946 }
947}