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 Degree, GGLWEAutomorphismKey, GGLWECiphertext, GGLWEInfos, GGLWESwitchingKey, GGLWETensorKey, GGSWCiphertext, GGSWInfos,
10 GLWECiphertext, GLWEInfos, GLWEPlaintext, GLWEPublicKey, GLWESecret, Rank,
11 prepared::{
12 GGLWEAutomorphismKeyPrepared, GGLWECiphertextPrepared, GGLWESwitchingKeyPrepared, GGLWETensorKeyPrepared,
13 GGSWCiphertextPrepared, GLWEPublicKeyPrepared, GLWESecretPrepared,
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) -> (GGLWECiphertext<&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) -> (GGLWECiphertextPrepared<&mut [u8], B>, &mut Self)
44 where
45 A: GGLWEInfos;
46}
47
48pub trait TakeGGSW {
49 fn take_ggsw<A>(&mut self, infos: &A) -> (GGSWCiphertext<&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) -> (GGSWCiphertextPrepared<&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<GGSWCiphertextPrepared<&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) -> (GGLWESwitchingKey<&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) -> (GGLWESwitchingKeyPrepared<&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) -> (GGLWETensorKey<&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) -> (GGLWETensorKeyPrepared<&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) -> (GGLWEAutomorphismKey<&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) -> (GGLWEAutomorphismKeyPrepared<&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::builder()
133 .base2k(infos.base2k())
134 .k(infos.k())
135 .data(data)
136 .build()
137 .unwrap(),
138 scratch,
139 )
140 }
141}
142
143impl<B: Backend> TakeGLWECtSlice for Scratch<B>
144where
145 Scratch<B>: TakeVecZnx,
146{
147 fn take_glwe_ct_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GLWECiphertext<&mut [u8]>>, &mut Self)
148 where
149 A: GLWEInfos,
150 {
151 let mut scratch: &mut Scratch<B> = self;
152 let mut cts: Vec<GLWECiphertext<&mut [u8]>> = Vec::with_capacity(size);
153 for _ in 0..size {
154 let (ct, new_scratch) = scratch.take_glwe_ct(infos);
155 scratch = new_scratch;
156 cts.push(ct);
157 }
158 (cts, scratch)
159 }
160}
161
162impl<B: Backend> TakeGLWEPt<B> for Scratch<B>
163where
164 Scratch<B>: TakeVecZnx,
165{
166 fn take_glwe_pt<A>(&mut self, infos: &A) -> (GLWEPlaintext<&mut [u8]>, &mut Self)
167 where
168 A: GLWEInfos,
169 {
170 let (data, scratch) = self.take_vec_znx(infos.n().into(), 1, infos.size());
171 (
172 GLWEPlaintext::builder()
173 .base2k(infos.base2k())
174 .k(infos.k())
175 .data(data)
176 .build()
177 .unwrap(),
178 scratch,
179 )
180 }
181}
182
183impl<B: Backend> TakeGGLWE for Scratch<B>
184where
185 Scratch<B>: TakeMatZnx,
186{
187 fn take_gglwe<A>(&mut self, infos: &A) -> (GGLWECiphertext<&mut [u8]>, &mut Self)
188 where
189 A: GGLWEInfos,
190 {
191 let (data, scratch) = self.take_mat_znx(
192 infos.n().into(),
193 infos.dnum().0.div_ceil(infos.dsize().0) as usize,
194 infos.rank_in().into(),
195 (infos.rank_out() + 1).into(),
196 infos.size(),
197 );
198 (
199 GGLWECiphertext::builder()
200 .base2k(infos.base2k())
201 .k(infos.k())
202 .dsize(infos.dsize())
203 .data(data)
204 .build()
205 .unwrap(),
206 scratch,
207 )
208 }
209}
210
211impl<B: Backend> TakeGGLWEPrepared<B> for Scratch<B>
212where
213 Scratch<B>: TakeVmpPMat<B>,
214{
215 fn take_gglwe_prepared<A>(&mut self, infos: &A) -> (GGLWECiphertextPrepared<&mut [u8], B>, &mut Self)
216 where
217 A: GGLWEInfos,
218 {
219 let (data, scratch) = self.take_vmp_pmat(
220 infos.n().into(),
221 infos.dnum().into(),
222 infos.rank_in().into(),
223 (infos.rank_out() + 1).into(),
224 infos.size(),
225 );
226 (
227 GGLWECiphertextPrepared::builder()
228 .base2k(infos.base2k())
229 .dsize(infos.dsize())
230 .k(infos.k())
231 .data(data)
232 .build()
233 .unwrap(),
234 scratch,
235 )
236 }
237}
238
239impl<B: Backend> TakeGGSW for Scratch<B>
240where
241 Scratch<B>: TakeMatZnx,
242{
243 fn take_ggsw<A>(&mut self, infos: &A) -> (GGSWCiphertext<&mut [u8]>, &mut Self)
244 where
245 A: GGSWInfos,
246 {
247 let (data, scratch) = self.take_mat_znx(
248 infos.n().into(),
249 infos.dnum().into(),
250 (infos.rank() + 1).into(),
251 (infos.rank() + 1).into(),
252 infos.size(),
253 );
254 (
255 GGSWCiphertext::builder()
256 .base2k(infos.base2k())
257 .dsize(infos.dsize())
258 .k(infos.k())
259 .data(data)
260 .build()
261 .unwrap(),
262 scratch,
263 )
264 }
265}
266
267impl<B: Backend> TakeGGSWPrepared<B> for Scratch<B>
268where
269 Scratch<B>: TakeVmpPMat<B>,
270{
271 fn take_ggsw_prepared<A>(&mut self, infos: &A) -> (GGSWCiphertextPrepared<&mut [u8], B>, &mut Self)
272 where
273 A: GGSWInfos,
274 {
275 let (data, scratch) = self.take_vmp_pmat(
276 infos.n().into(),
277 infos.dnum().into(),
278 (infos.rank() + 1).into(),
279 (infos.rank() + 1).into(),
280 infos.size(),
281 );
282 (
283 GGSWCiphertextPrepared::builder()
284 .base2k(infos.base2k())
285 .dsize(infos.dsize())
286 .k(infos.k())
287 .data(data)
288 .build()
289 .unwrap(),
290 scratch,
291 )
292 }
293}
294
295impl<B: Backend> TakeGGSWPreparedSlice<B> for Scratch<B>
296where
297 Scratch<B>: TakeGGSWPrepared<B>,
298{
299 fn take_ggsw_prepared_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GGSWCiphertextPrepared<&mut [u8], B>>, &mut Self)
300 where
301 A: GGSWInfos,
302 {
303 let mut scratch: &mut Scratch<B> = self;
304 let mut cts: Vec<GGSWCiphertextPrepared<&mut [u8], B>> = Vec::with_capacity(size);
305 for _ in 0..size {
306 let (ct, new_scratch) = scratch.take_ggsw_prepared(infos);
307 scratch = new_scratch;
308 cts.push(ct)
309 }
310 (cts, scratch)
311 }
312}
313
314impl<B: Backend> TakeGLWEPk for Scratch<B>
315where
316 Scratch<B>: TakeVecZnx,
317{
318 fn take_glwe_pk<A>(&mut self, infos: &A) -> (GLWEPublicKey<&mut [u8]>, &mut Self)
319 where
320 A: GLWEInfos,
321 {
322 let (data, scratch) = self.take_vec_znx(infos.n().into(), (infos.rank() + 1).into(), infos.size());
323 (
324 GLWEPublicKey::builder()
325 .base2k(infos.base2k())
326 .k(infos.k())
327 .base2k(infos.base2k())
328 .data(data)
329 .build()
330 .unwrap(),
331 scratch,
332 )
333 }
334}
335
336impl<B: Backend> TakeGLWEPkPrepared<B> for Scratch<B>
337where
338 Scratch<B>: TakeVecZnxDft<B>,
339{
340 fn take_glwe_pk_prepared<A>(&mut self, infos: &A) -> (GLWEPublicKeyPrepared<&mut [u8], B>, &mut Self)
341 where
342 A: GLWEInfos,
343 {
344 let (data, scratch) = self.take_vec_znx_dft(infos.n().into(), (infos.rank() + 1).into(), infos.size());
345 (
346 GLWEPublicKeyPrepared::builder()
347 .base2k(infos.base2k())
348 .k(infos.k())
349 .data(data)
350 .build()
351 .unwrap(),
352 scratch,
353 )
354 }
355}
356
357impl<B: Backend> TakeGLWESecret for Scratch<B>
358where
359 Scratch<B>: TakeScalarZnx,
360{
361 fn take_glwe_secret(&mut self, n: Degree, rank: Rank) -> (GLWESecret<&mut [u8]>, &mut Self) {
362 let (data, scratch) = self.take_scalar_znx(n.into(), rank.into());
363 (
364 GLWESecret {
365 data,
366 dist: Distribution::NONE,
367 },
368 scratch,
369 )
370 }
371}
372
373impl<B: Backend> TakeGLWESecretPrepared<B> for Scratch<B>
374where
375 Scratch<B>: TakeSvpPPol<B>,
376{
377 fn take_glwe_secret_prepared(&mut self, n: Degree, rank: Rank) -> (GLWESecretPrepared<&mut [u8], B>, &mut Self) {
378 let (data, scratch) = self.take_svp_ppol(n.into(), rank.into());
379 (
380 GLWESecretPrepared {
381 data,
382 dist: Distribution::NONE,
383 },
384 scratch,
385 )
386 }
387}
388
389impl<B: Backend> TakeGLWESwitchingKey for Scratch<B>
390where
391 Scratch<B>: TakeMatZnx,
392{
393 fn take_glwe_switching_key<A>(&mut self, infos: &A) -> (GGLWESwitchingKey<&mut [u8]>, &mut Self)
394 where
395 A: GGLWEInfos,
396 {
397 let (data, scratch) = self.take_gglwe(infos);
398 (
399 GGLWESwitchingKey {
400 key: data,
401 sk_in_n: 0,
402 sk_out_n: 0,
403 },
404 scratch,
405 )
406 }
407}
408
409impl<B: Backend> TakeGGLWESwitchingKeyPrepared<B> for Scratch<B>
410where
411 Scratch<B>: TakeGGLWEPrepared<B>,
412{
413 fn take_gglwe_switching_key_prepared<A>(&mut self, infos: &A) -> (GGLWESwitchingKeyPrepared<&mut [u8], B>, &mut Self)
414 where
415 A: GGLWEInfos,
416 {
417 let (data, scratch) = self.take_gglwe_prepared(infos);
418 (
419 GGLWESwitchingKeyPrepared {
420 key: data,
421 sk_in_n: 0,
422 sk_out_n: 0,
423 },
424 scratch,
425 )
426 }
427}
428
429impl<B: Backend> TakeGGLWEAutomorphismKey for Scratch<B>
430where
431 Scratch<B>: TakeMatZnx,
432{
433 fn take_gglwe_automorphism_key<A>(&mut self, infos: &A) -> (GGLWEAutomorphismKey<&mut [u8]>, &mut Self)
434 where
435 A: GGLWEInfos,
436 {
437 let (data, scratch) = self.take_glwe_switching_key(infos);
438 (GGLWEAutomorphismKey { key: data, p: 0 }, scratch)
439 }
440}
441
442impl<B: Backend> TakeGGLWEAutomorphismKeyPrepared<B> for Scratch<B>
443where
444 Scratch<B>: TakeGGLWESwitchingKeyPrepared<B>,
445{
446 fn take_gglwe_automorphism_key_prepared<A>(&mut self, infos: &A) -> (GGLWEAutomorphismKeyPrepared<&mut [u8], B>, &mut Self)
447 where
448 A: GGLWEInfos,
449 {
450 let (data, scratch) = self.take_gglwe_switching_key_prepared(infos);
451 (GGLWEAutomorphismKeyPrepared { key: data, p: 0 }, scratch)
452 }
453}
454
455impl<B: Backend> TakeTensorKey for Scratch<B>
456where
457 Scratch<B>: TakeMatZnx,
458{
459 fn take_tensor_key<A>(&mut self, infos: &A) -> (GGLWETensorKey<&mut [u8]>, &mut Self)
460 where
461 A: GGLWEInfos,
462 {
463 assert_eq!(
464 infos.rank_in(),
465 infos.rank_out(),
466 "rank_in != rank_out is not supported for GGLWETensorKey"
467 );
468 let mut keys: Vec<GGLWESwitchingKey<&mut [u8]>> = Vec::new();
469 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
470
471 let mut scratch: &mut Scratch<B> = self;
472
473 let mut ksk_infos: crate::layouts::GGLWECiphertextLayout = infos.layout();
474 ksk_infos.rank_in = Rank(1);
475
476 if pairs != 0 {
477 let (gglwe, s) = scratch.take_glwe_switching_key(&ksk_infos);
478 scratch = s;
479 keys.push(gglwe);
480 }
481 for _ in 1..pairs {
482 let (gglwe, s) = scratch.take_glwe_switching_key(&ksk_infos);
483 scratch = s;
484 keys.push(gglwe);
485 }
486 (GGLWETensorKey { keys }, scratch)
487 }
488}
489
490impl<B: Backend> TakeGGLWETensorKeyPrepared<B> for Scratch<B>
491where
492 Scratch<B>: TakeVmpPMat<B>,
493{
494 fn take_gglwe_tensor_key_prepared<A>(&mut self, infos: &A) -> (GGLWETensorKeyPrepared<&mut [u8], B>, &mut Self)
495 where
496 A: GGLWEInfos,
497 {
498 assert_eq!(
499 infos.rank_in(),
500 infos.rank_out(),
501 "rank_in != rank_out is not supported for GGLWETensorKeyPrepared"
502 );
503
504 let mut keys: Vec<GGLWESwitchingKeyPrepared<&mut [u8], B>> = Vec::new();
505 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
506
507 let mut scratch: &mut Scratch<B> = self;
508
509 let mut ksk_infos: crate::layouts::GGLWECiphertextLayout = infos.layout();
510 ksk_infos.rank_in = Rank(1);
511
512 if pairs != 0 {
513 let (gglwe, s) = scratch.take_gglwe_switching_key_prepared(&ksk_infos);
514 scratch = s;
515 keys.push(gglwe);
516 }
517 for _ in 1..pairs {
518 let (gglwe, s) = scratch.take_gglwe_switching_key_prepared(&ksk_infos);
519 scratch = s;
520 keys.push(gglwe);
521 }
522 (GGLWETensorKeyPrepared { keys }, scratch)
523 }
524}