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