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