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