poulpy/poulpy-core/src/scratch.rs
poulpy-core/src/scratch.rsRaw
1use poulpy_hal::{
2 api::{TakeMatZnx, TakeScalarZnx, TakeSvpPPol, TakeVecZnx, TakeVecZnxDft, TakeVmpPMat},
3 layouts::{Backend, Scratch},
4};
5
6use crate::{
7 dist::Distribution,
8 layouts::{
9 AutomorphismKey, Degree, GGLWE, GGLWEInfos, GGSW, GGSWInfos, GLWECiphertext, GLWEInfos, GLWEPlaintext, GLWEPublicKey,
10 GLWESecret, GLWESwitchingKey, Rank, TensorKey,
11 prepared::{
12 AutomorphismKeyPrepared, GGLWEPrepared, GGSWPrepared, GLWEPublicKeyPrepared, GLWESecretPrepared,
13 GLWESwitchingKeyPrepared, TensorKeyPrepared,
14 },
15 },
16};
17
18pub trait TakeGLWECt {
19 fn take_glwe_ct<A>(&mut self, infos: &A) -> (GLWECiphertext<&mut [u8]>, &mut Self)
20 where
21 A: GLWEInfos;
22}
23
24pub trait TakeGLWECtSlice {
25 fn take_glwe_ct_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GLWECiphertext<&mut [u8]>>, &mut Self)
26 where
27 A: GLWEInfos;
28}
29
30pub trait TakeGLWEPt<B: Backend> {
31 fn take_glwe_pt<A>(&mut self, infos: &A) -> (GLWEPlaintext<&mut [u8]>, &mut Self)
32 where
33 A: GLWEInfos;
34}
35
36pub trait TakeGGLWE {
37 fn take_gglwe<A>(&mut self, infos: &A) -> (GGLWE<&mut [u8]>, &mut Self)
38 where
39 A: GGLWEInfos;
40}
41
42pub trait TakeGGLWEPrepared<B: Backend> {
43 fn take_gglwe_prepared<A>(&mut self, infos: &A) -> (GGLWEPrepared<&mut [u8], B>, &mut Self)
44 where
45 A: GGLWEInfos;
46}
47
48pub trait TakeGGSW {
49 fn take_ggsw<A>(&mut self, infos: &A) -> (GGSW<&mut [u8]>, &mut Self)
50 where
51 A: GGSWInfos;
52}
53
54pub trait TakeGGSWPrepared<B: Backend> {
55 fn take_ggsw_prepared<A>(&mut self, infos: &A) -> (GGSWPrepared<&mut [u8], B>, &mut Self)
56 where
57 A: GGSWInfos;
58}
59
60pub trait TakeGGSWPreparedSlice<B: Backend> {
61 fn take_ggsw_prepared_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GGSWPrepared<&mut [u8], B>>, &mut Self)
62 where
63 A: GGSWInfos;
64}
65
66pub trait TakeGLWESecret {
67 fn take_glwe_secret(&mut self, n: Degree, rank: Rank) -> (GLWESecret<&mut [u8]>, &mut Self);
68}
69
70pub trait TakeGLWESecretPrepared<B: Backend> {
71 fn take_glwe_secret_prepared(&mut self, n: Degree, rank: Rank) -> (GLWESecretPrepared<&mut [u8], B>, &mut Self);
72}
73
74pub trait TakeGLWEPk {
75 fn take_glwe_pk<A>(&mut self, infos: &A) -> (GLWEPublicKey<&mut [u8]>, &mut Self)
76 where
77 A: GLWEInfos;
78}
79
80pub trait TakeGLWEPkPrepared<B: Backend> {
81 fn take_glwe_pk_prepared<A>(&mut self, infos: &A) -> (GLWEPublicKeyPrepared<&mut [u8], B>, &mut Self)
82 where
83 A: GLWEInfos;
84}
85
86pub trait TakeGLWESwitchingKey {
87 fn take_glwe_switching_key<A>(&mut self, infos: &A) -> (GLWESwitchingKey<&mut [u8]>, &mut Self)
88 where
89 A: GGLWEInfos;
90}
91
92pub trait TakeGGLWESwitchingKeyPrepared<B: Backend> {
93 fn take_gglwe_switching_key_prepared<A>(&mut self, infos: &A) -> (GLWESwitchingKeyPrepared<&mut [u8], B>, &mut Self)
94 where
95 A: GGLWEInfos;
96}
97
98pub trait TakeTensorKey {
99 fn take_tensor_key<A>(&mut self, infos: &A) -> (TensorKey<&mut [u8]>, &mut Self)
100 where
101 A: GGLWEInfos;
102}
103
104pub trait TakeGGLWETensorKeyPrepared<B: Backend> {
105 fn take_gglwe_tensor_key_prepared<A>(&mut self, infos: &A) -> (TensorKeyPrepared<&mut [u8], B>, &mut Self)
106 where
107 A: GGLWEInfos;
108}
109
110pub trait TakeGGLWEAutomorphismKey {
111 fn take_gglwe_automorphism_key<A>(&mut self, infos: &A) -> (AutomorphismKey<&mut [u8]>, &mut Self)
112 where
113 A: GGLWEInfos;
114}
115
116pub trait TakeGGLWEAutomorphismKeyPrepared<B: Backend> {
117 fn take_gglwe_automorphism_key_prepared<A>(&mut self, infos: &A) -> (AutomorphismKeyPrepared<&mut [u8], B>, &mut Self)
118 where
119 A: GGLWEInfos;
120}
121
122impl<B: Backend> TakeGLWECt for Scratch<B>
123where
124 Scratch<B>: TakeVecZnx,
125{
126 fn take_glwe_ct<A>(&mut self, infos: &A) -> (GLWECiphertext<&mut [u8]>, &mut Self)
127 where
128 A: GLWEInfos,
129 {
130 let (data, scratch) = self.take_vec_znx(infos.n().into(), (infos.rank() + 1).into(), infos.size());
131 (
132 GLWECiphertext {
133 k: infos.k(),
134 base2k: infos.base2k(),
135 data,
136 },
137 scratch,
138 )
139 }
140}
141
142impl<B: Backend> TakeGLWECtSlice for Scratch<B>
143where
144 Scratch<B>: TakeVecZnx,
145{
146 fn take_glwe_ct_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GLWECiphertext<&mut [u8]>>, &mut Self)
147 where
148 A: GLWEInfos,
149 {
150 let mut scratch: &mut Scratch<B> = self;
151 let mut cts: Vec<GLWECiphertext<&mut [u8]>> = Vec::with_capacity(size);
152 for _ in 0..size {
153 let (ct, new_scratch) = scratch.take_glwe_ct(infos);
154 scratch = new_scratch;
155 cts.push(ct);
156 }
157 (cts, scratch)
158 }
159}
160
161impl<B: Backend> TakeGLWEPt<B> for Scratch<B>
162where
163 Scratch<B>: TakeVecZnx,
164{
165 fn take_glwe_pt<A>(&mut self, infos: &A) -> (GLWEPlaintext<&mut [u8]>, &mut Self)
166 where
167 A: GLWEInfos,
168 {
169 let (data, scratch) = self.take_vec_znx(infos.n().into(), 1, infos.size());
170 (
171 GLWEPlaintext {
172 k: infos.k(),
173 base2k: infos.base2k(),
174 data,
175 },
176 scratch,
177 )
178 }
179}
180
181impl<B: Backend> TakeGGLWE for Scratch<B>
182where
183 Scratch<B>: TakeMatZnx,
184{
185 fn take_gglwe<A>(&mut self, infos: &A) -> (GGLWE<&mut [u8]>, &mut Self)
186 where
187 A: GGLWEInfos,
188 {
189 let (data, scratch) = self.take_mat_znx(
190 infos.n().into(),
191 infos.dnum().0.div_ceil(infos.dsize().0) as usize,
192 infos.rank_in().into(),
193 (infos.rank_out() + 1).into(),
194 infos.size(),
195 );
196 (
197 GGLWE {
198 k: infos.k(),
199 base2k: infos.base2k(),
200 dsize: infos.dsize(),
201 data,
202 },
203 scratch,
204 )
205 }
206}
207
208impl<B: Backend> TakeGGLWEPrepared<B> for Scratch<B>
209where
210 Scratch<B>: TakeVmpPMat<B>,
211{
212 fn take_gglwe_prepared<A>(&mut self, infos: &A) -> (GGLWEPrepared<&mut [u8], B>, &mut Self)
213 where
214 A: GGLWEInfos,
215 {
216 let (data, scratch) = self.take_vmp_pmat(
217 infos.n().into(),
218 infos.dnum().into(),
219 infos.rank_in().into(),
220 (infos.rank_out() + 1).into(),
221 infos.size(),
222 );
223 (
224 GGLWEPrepared {
225 k: infos.k(),
226 base2k: infos.base2k(),
227 dsize: infos.dsize(),
228 data,
229 },
230 scratch,
231 )
232 }
233}
234
235impl<B: Backend> TakeGGSW for Scratch<B>
236where
237 Scratch<B>: TakeMatZnx,
238{
239 fn take_ggsw<A>(&mut self, infos: &A) -> (GGSW<&mut [u8]>, &mut Self)
240 where
241 A: GGSWInfos,
242 {
243 let (data, scratch) = self.take_mat_znx(
244 infos.n().into(),
245 infos.dnum().into(),
246 (infos.rank() + 1).into(),
247 (infos.rank() + 1).into(),
248 infos.size(),
249 );
250 (
251 GGSW {
252 k: infos.k(),
253 base2k: infos.base2k(),
254 dsize: infos.dsize(),
255 data,
256 },
257 scratch,
258 )
259 }
260}
261
262impl<B: Backend> TakeGGSWPrepared<B> for Scratch<B>
263where
264 Scratch<B>: TakeVmpPMat<B>,
265{
266 fn take_ggsw_prepared<A>(&mut self, infos: &A) -> (GGSWPrepared<&mut [u8], B>, &mut Self)
267 where
268 A: GGSWInfos,
269 {
270 let (data, scratch) = self.take_vmp_pmat(
271 infos.n().into(),
272 infos.dnum().into(),
273 (infos.rank() + 1).into(),
274 (infos.rank() + 1).into(),
275 infos.size(),
276 );
277 (
278 GGSWPrepared {
279 k: infos.k(),
280 base2k: infos.base2k(),
281 dsize: infos.dsize(),
282 data,
283 },
284 scratch,
285 )
286 }
287}
288
289impl<B: Backend> TakeGGSWPreparedSlice<B> for Scratch<B>
290where
291 Scratch<B>: TakeGGSWPrepared<B>,
292{
293 fn take_ggsw_prepared_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GGSWPrepared<&mut [u8], B>>, &mut Self)
294 where
295 A: GGSWInfos,
296 {
297 let mut scratch: &mut Scratch<B> = self;
298 let mut cts: Vec<GGSWPrepared<&mut [u8], B>> = Vec::with_capacity(size);
299 for _ in 0..size {
300 let (ct, new_scratch) = scratch.take_ggsw_prepared(infos);
301 scratch = new_scratch;
302 cts.push(ct)
303 }
304 (cts, scratch)
305 }
306}
307
308impl<B: Backend> TakeGLWEPk for Scratch<B>
309where
310 Scratch<B>: TakeVecZnx,
311{
312 fn take_glwe_pk<A>(&mut self, infos: &A) -> (GLWEPublicKey<&mut [u8]>, &mut Self)
313 where
314 A: GLWEInfos,
315 {
316 let (data, scratch) = self.take_vec_znx(infos.n().into(), (infos.rank() + 1).into(), infos.size());
317 (
318 GLWEPublicKey {
319 k: infos.k(),
320 dist: Distribution::NONE,
321 base2k: infos.base2k(),
322 data,
323 },
324 scratch,
325 )
326 }
327}
328
329impl<B: Backend> TakeGLWEPkPrepared<B> for Scratch<B>
330where
331 Scratch<B>: TakeVecZnxDft<B>,
332{
333 fn take_glwe_pk_prepared<A>(&mut self, infos: &A) -> (GLWEPublicKeyPrepared<&mut [u8], B>, &mut Self)
334 where
335 A: GLWEInfos,
336 {
337 let (data, scratch) = self.take_vec_znx_dft(infos.n().into(), (infos.rank() + 1).into(), infos.size());
338 (
339 GLWEPublicKeyPrepared {
340 k: infos.k(),
341 dist: Distribution::NONE,
342 base2k: infos.base2k(),
343 data,
344 },
345 scratch,
346 )
347 }
348}
349
350impl<B: Backend> TakeGLWESecret for Scratch<B>
351where
352 Scratch<B>: TakeScalarZnx,
353{
354 fn take_glwe_secret(&mut self, n: Degree, rank: Rank) -> (GLWESecret<&mut [u8]>, &mut Self) {
355 let (data, scratch) = self.take_scalar_znx(n.into(), rank.into());
356 (
357 GLWESecret {
358 data,
359 dist: Distribution::NONE,
360 },
361 scratch,
362 )
363 }
364}
365
366impl<B: Backend> TakeGLWESecretPrepared<B> for Scratch<B>
367where
368 Scratch<B>: TakeSvpPPol<B>,
369{
370 fn take_glwe_secret_prepared(&mut self, n: Degree, rank: Rank) -> (GLWESecretPrepared<&mut [u8], B>, &mut Self) {
371 let (data, scratch) = self.take_svp_ppol(n.into(), rank.into());
372 (
373 GLWESecretPrepared {
374 data,
375 dist: Distribution::NONE,
376 },
377 scratch,
378 )
379 }
380}
381
382impl<B: Backend> TakeGLWESwitchingKey for Scratch<B>
383where
384 Scratch<B>: TakeMatZnx,
385{
386 fn take_glwe_switching_key<A>(&mut self, infos: &A) -> (GLWESwitchingKey<&mut [u8]>, &mut Self)
387 where
388 A: GGLWEInfos,
389 {
390 let (data, scratch) = self.take_gglwe(infos);
391 (
392 GLWESwitchingKey {
393 key: data,
394 sk_in_n: 0,
395 sk_out_n: 0,
396 },
397 scratch,
398 )
399 }
400}
401
402impl<B: Backend> TakeGGLWESwitchingKeyPrepared<B> for Scratch<B>
403where
404 Scratch<B>: TakeGGLWEPrepared<B>,
405{
406 fn take_gglwe_switching_key_prepared<A>(&mut self, infos: &A) -> (GLWESwitchingKeyPrepared<&mut [u8], B>, &mut Self)
407 where
408 A: GGLWEInfos,
409 {
410 let (data, scratch) = self.take_gglwe_prepared(infos);
411 (
412 GLWESwitchingKeyPrepared {
413 key: data,
414 sk_in_n: 0,
415 sk_out_n: 0,
416 },
417 scratch,
418 )
419 }
420}
421
422impl<B: Backend> TakeGGLWEAutomorphismKey for Scratch<B>
423where
424 Scratch<B>: TakeMatZnx,
425{
426 fn take_gglwe_automorphism_key<A>(&mut self, infos: &A) -> (AutomorphismKey<&mut [u8]>, &mut Self)
427 where
428 A: GGLWEInfos,
429 {
430 let (data, scratch) = self.take_glwe_switching_key(infos);
431 (AutomorphismKey { key: data, p: 0 }, scratch)
432 }
433}
434
435impl<B: Backend> TakeGGLWEAutomorphismKeyPrepared<B> for Scratch<B>
436where
437 Scratch<B>: TakeGGLWESwitchingKeyPrepared<B>,
438{
439 fn take_gglwe_automorphism_key_prepared<A>(&mut self, infos: &A) -> (AutomorphismKeyPrepared<&mut [u8], B>, &mut Self)
440 where
441 A: GGLWEInfos,
442 {
443 let (data, scratch) = self.take_gglwe_switching_key_prepared(infos);
444 (AutomorphismKeyPrepared { key: data, p: 0 }, scratch)
445 }
446}
447
448impl<B: Backend> TakeTensorKey for Scratch<B>
449where
450 Scratch<B>: TakeMatZnx,
451{
452 fn take_tensor_key<A>(&mut self, infos: &A) -> (TensorKey<&mut [u8]>, &mut Self)
453 where
454 A: GGLWEInfos,
455 {
456 assert_eq!(
457 infos.rank_in(),
458 infos.rank_out(),
459 "rank_in != rank_out is not supported for GGLWETensorKey"
460 );
461 let mut keys: Vec<GLWESwitchingKey<&mut [u8]>> = Vec::new();
462 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
463
464 let mut scratch: &mut Scratch<B> = self;
465
466 let mut ksk_infos: crate::layouts::GGLWECiphertextLayout = infos.layout();
467 ksk_infos.rank_in = Rank(1);
468
469 if pairs != 0 {
470 let (gglwe, s) = scratch.take_glwe_switching_key(&ksk_infos);
471 scratch = s;
472 keys.push(gglwe);
473 }
474 for _ in 1..pairs {
475 let (gglwe, s) = scratch.take_glwe_switching_key(&ksk_infos);
476 scratch = s;
477 keys.push(gglwe);
478 }
479 (TensorKey { keys }, scratch)
480 }
481}
482
483impl<B: Backend> TakeGGLWETensorKeyPrepared<B> for Scratch<B>
484where
485 Scratch<B>: TakeVmpPMat<B>,
486{
487 fn take_gglwe_tensor_key_prepared<A>(&mut self, infos: &A) -> (TensorKeyPrepared<&mut [u8], B>, &mut Self)
488 where
489 A: GGLWEInfos,
490 {
491 assert_eq!(
492 infos.rank_in(),
493 infos.rank_out(),
494 "rank_in != rank_out is not supported for GGLWETensorKeyPrepared"
495 );
496
497 let mut keys: Vec<GLWESwitchingKeyPrepared<&mut [u8], B>> = Vec::new();
498 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
499
500 let mut scratch: &mut Scratch<B> = self;
501
502 let mut ksk_infos: crate::layouts::GGLWECiphertextLayout = infos.layout();
503 ksk_infos.rank_in = Rank(1);
504
505 if pairs != 0 {
506 let (gglwe, s) = scratch.take_gglwe_switching_key_prepared(&ksk_infos);
507 scratch = s;
508 keys.push(gglwe);
509 }
510 for _ in 1..pairs {
511 let (gglwe, s) = scratch.take_gglwe_switching_key_prepared(&ksk_infos);
512 scratch = s;
513 keys.push(gglwe);
514 }
515 (TensorKeyPrepared { keys }, scratch)
516 }
517}