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