poulpy/poulpy-core/src/scratch.rs
poulpy-core/src/scratch.rsRaw
1use poulpy_hal::{
2 api::{ModuleN, ScratchAvailable, ScratchFromBytes, ScratchTakeBasic, SvpPPolBytesOf, VecZnxDftBytesOf, VmpPMatBytesOf},
3 layouts::{Backend, Scratch},
4};
5
6use crate::{
7 dist::Distribution,
8 layouts::{
9 Degree, GGLWE, GGLWEInfos, GGLWELayout, GGSW, GGSWInfos, GLWE, GLWEAutomorphismKey, GLWEInfos, GLWEPlaintext,
10 GLWEPrepared, GLWEPublicKey, GLWESecret, GLWESecretTensor, GLWESwitchingKey, GLWETensorKey, LWE, LWEInfos, LWEPlaintext,
11 Rank,
12 prepared::{
13 GGLWEPrepared, GGSWPrepared, GLWEAutomorphismKeyPrepared, GLWEPublicKeyPrepared, GLWESecretPrepared,
14 GLWESwitchingKeyPrepared, GLWETensorKeyPrepared,
15 },
16 },
17};
18
19pub trait ScratchTakeCore<B: Backend>
20where
21 Self: ScratchTakeBasic + ScratchAvailable + ScratchFromBytes<B>,
22{
23 fn take_lwe<A>(&mut self, infos: &A) -> (LWE<&mut [u8]>, &mut Self)
24 where
25 A: LWEInfos,
26 {
27 let (data, scratch) = self.take_vec_znx(infos.n().into(), 1, infos.size());
28 (
29 LWE {
30 k: infos.k(),
31 base2k: infos.base2k(),
32 data,
33 },
34 scratch,
35 )
36 }
37
38 fn take_lwe_plaintext<A>(&mut self, infos: &A) -> (LWEPlaintext<&mut [u8]>, &mut Self)
39 where
40 A: LWEInfos,
41 {
42 let (data, scratch) = self.take_vec_znx(1, 1, infos.size());
43 (
44 LWEPlaintext {
45 k: infos.k(),
46 base2k: infos.base2k(),
47 data,
48 },
49 scratch,
50 )
51 }
52
53 fn take_glwe<A>(&mut self, infos: &A) -> (GLWE<&mut [u8]>, &mut Self)
54 where
55 A: GLWEInfos,
56 {
57 let (data, scratch) = self.take_vec_znx(infos.n().into(), (infos.rank() + 1).into(), infos.size());
58 (
59 GLWE {
60 k: infos.k(),
61 base2k: infos.base2k(),
62 data,
63 },
64 scratch,
65 )
66 }
67
68 fn take_glwe_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GLWE<&mut [u8]>>, &mut Self)
69 where
70 A: GLWEInfos,
71 {
72 let mut scratch: &mut Self = self;
73 let mut cts: Vec<GLWE<&mut [u8]>> = Vec::with_capacity(size);
74 for _ in 0..size {
75 let (ct, new_scratch) = scratch.take_glwe(infos);
76 scratch = new_scratch;
77 cts.push(ct);
78 }
79 (cts, scratch)
80 }
81
82 fn take_glwe_plaintext<A>(&mut self, infos: &A) -> (GLWEPlaintext<&mut [u8]>, &mut Self)
83 where
84 A: GLWEInfos,
85 {
86 let (data, scratch) = self.take_vec_znx(infos.n().into(), 1, infos.size());
87 (
88 GLWEPlaintext {
89 k: infos.k(),
90 base2k: infos.base2k(),
91 data,
92 },
93 scratch,
94 )
95 }
96
97 fn take_gglwe<A>(&mut self, infos: &A) -> (GGLWE<&mut [u8]>, &mut Self)
98 where
99 A: GGLWEInfos,
100 {
101 let (data, scratch) = self.take_mat_znx(
102 infos.n().into(),
103 infos.dnum().0.div_ceil(infos.dsize().0) as usize,
104 infos.rank_in().into(),
105 (infos.rank_out() + 1).into(),
106 infos.size(),
107 );
108 (
109 GGLWE {
110 k: infos.k(),
111 base2k: infos.base2k(),
112 dsize: infos.dsize(),
113 data,
114 },
115 scratch,
116 )
117 }
118
119 fn take_gglwe_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GGLWEPrepared<&mut [u8], B>, &mut Self)
120 where
121 A: GGLWEInfos,
122 M: ModuleN + VmpPMatBytesOf,
123 {
124 assert_eq!(module.n() as u32, infos.n());
125 let (data, scratch) = self.take_vmp_pmat(
126 module,
127 infos.dnum().into(),
128 infos.rank_in().into(),
129 (infos.rank_out() + 1).into(),
130 infos.size(),
131 );
132 (
133 GGLWEPrepared {
134 k: infos.k(),
135 base2k: infos.base2k(),
136 dsize: infos.dsize(),
137 data,
138 },
139 scratch,
140 )
141 }
142
143 fn take_ggsw<A>(&mut self, infos: &A) -> (GGSW<&mut [u8]>, &mut Self)
144 where
145 A: GGSWInfos,
146 {
147 let (data, scratch) = self.take_mat_znx(
148 infos.n().into(),
149 infos.dnum().into(),
150 (infos.rank() + 1).into(),
151 (infos.rank() + 1).into(),
152 infos.size(),
153 );
154 (
155 GGSW {
156 k: infos.k(),
157 base2k: infos.base2k(),
158 dsize: infos.dsize(),
159 data,
160 },
161 scratch,
162 )
163 }
164
165 fn take_ggsw_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GGSWPrepared<&mut [u8], B>, &mut Self)
166 where
167 A: GGSWInfos,
168 M: ModuleN + VmpPMatBytesOf,
169 {
170 assert_eq!(module.n() as u32, infos.n());
171 let (data, scratch) = self.take_vmp_pmat(
172 module,
173 infos.dnum().into(),
174 (infos.rank() + 1).into(),
175 (infos.rank() + 1).into(),
176 infos.size(),
177 );
178 (
179 GGSWPrepared {
180 k: infos.k(),
181 base2k: infos.base2k(),
182 dsize: infos.dsize(),
183 data,
184 },
185 scratch,
186 )
187 }
188
189 fn take_ggsw_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GGSW<&mut [u8]>>, &mut Self)
190 where
191 A: GGSWInfos,
192 {
193 let mut scratch: &mut Self = self;
194 let mut cts: Vec<GGSW<&mut [u8]>> = Vec::with_capacity(size);
195 for _ in 0..size {
196 let (ct, new_scratch) = scratch.take_ggsw(infos);
197 scratch = new_scratch;
198 cts.push(ct)
199 }
200 (cts, scratch)
201 }
202
203 fn take_ggsw_prepared_slice<A, M>(
204 &mut self,
205 module: &M,
206 size: usize,
207 infos: &A,
208 ) -> (Vec<GGSWPrepared<&mut [u8], B>>, &mut Self)
209 where
210 A: GGSWInfos,
211 M: ModuleN + VmpPMatBytesOf,
212 {
213 let mut scratch: &mut Self = self;
214 let mut cts: Vec<GGSWPrepared<&mut [u8], B>> = Vec::with_capacity(size);
215 for _ in 0..size {
216 let (ct, new_scratch) = scratch.take_ggsw_prepared(module, infos);
217 scratch = new_scratch;
218 cts.push(ct)
219 }
220 (cts, scratch)
221 }
222
223 fn take_glwe_public_key<A>(&mut self, infos: &A) -> (GLWEPublicKey<&mut [u8]>, &mut Self)
224 where
225 A: GLWEInfos,
226 {
227 let (data, scratch) = self.take_glwe(infos);
228 (
229 GLWEPublicKey {
230 dist: Distribution::NONE,
231 key: data,
232 },
233 scratch,
234 )
235 }
236
237 fn take_glwe_public_key_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPublicKeyPrepared<&mut [u8], B>, &mut Self)
238 where
239 A: GLWEInfos,
240 M: ModuleN + VecZnxDftBytesOf,
241 {
242 let (data, scratch) = self.take_glwe_prepared(module, infos);
243 (
244 GLWEPublicKeyPrepared {
245 dist: Distribution::NONE,
246 key: data,
247 },
248 scratch,
249 )
250 }
251
252 fn take_glwe_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPrepared<&mut [u8], B>, &mut Self)
253 where
254 A: GLWEInfos,
255 M: ModuleN + VecZnxDftBytesOf,
256 {
257 assert_eq!(module.n() as u32, infos.n());
258 let (data, scratch) = self.take_vec_znx_dft(module, (infos.rank() + 1).into(), infos.size());
259 (
260 GLWEPrepared {
261 k: infos.k(),
262 base2k: infos.base2k(),
263 data,
264 },
265 scratch,
266 )
267 }
268
269 fn take_glwe_secret(&mut self, n: Degree, rank: Rank) -> (GLWESecret<&mut [u8]>, &mut Self) {
270 let (data, scratch) = self.take_scalar_znx(n.into(), rank.into());
271 (
272 GLWESecret {
273 data,
274 dist: Distribution::NONE,
275 },
276 scratch,
277 )
278 }
279
280 fn take_glwe_secret_tensor(&mut self, n: Degree, rank: Rank) -> (GLWESecretTensor<&mut [u8]>, &mut Self) {
281 let (data, scratch) = self.take_scalar_znx(n.into(), GLWESecretTensor::pairs(rank.into()));
282 (
283 GLWESecretTensor {
284 data,
285 rank,
286 dist: Distribution::NONE,
287 },
288 scratch,
289 )
290 }
291
292 fn take_glwe_secret_prepared<M>(&mut self, module: &M, rank: Rank) -> (GLWESecretPrepared<&mut [u8], B>, &mut Self)
293 where
294 M: ModuleN + SvpPPolBytesOf,
295 {
296 let (data, scratch) = self.take_svp_ppol(module, rank.into());
297 (
298 GLWESecretPrepared {
299 data,
300 dist: Distribution::NONE,
301 },
302 scratch,
303 )
304 }
305
306 fn take_glwe_switching_key<A>(&mut self, infos: &A) -> (GLWESwitchingKey<&mut [u8]>, &mut Self)
307 where
308 A: GGLWEInfos,
309 {
310 let (data, scratch) = self.take_gglwe(infos);
311 (
312 GLWESwitchingKey {
313 key: data,
314 input_degree: Degree(0),
315 output_degree: Degree(0),
316 },
317 scratch,
318 )
319 }
320
321 fn take_glwe_switching_key_prepared<A, M>(
322 &mut self,
323 module: &M,
324 infos: &A,
325 ) -> (GLWESwitchingKeyPrepared<&mut [u8], B>, &mut Self)
326 where
327 A: GGLWEInfos,
328 M: ModuleN + VmpPMatBytesOf,
329 {
330 assert_eq!(module.n() as u32, infos.n());
331 let (data, scratch) = self.take_gglwe_prepared(module, infos);
332 (
333 GLWESwitchingKeyPrepared {
334 key: data,
335 input_degree: Degree(0),
336 output_degree: Degree(0),
337 },
338 scratch,
339 )
340 }
341
342 fn take_glwe_automorphism_key<A>(&mut self, infos: &A) -> (GLWEAutomorphismKey<&mut [u8]>, &mut Self)
343 where
344 A: GGLWEInfos,
345 {
346 let (data, scratch) = self.take_gglwe(infos);
347 (GLWEAutomorphismKey { key: data, p: 0 }, scratch)
348 }
349
350 fn take_glwe_automorphism_key_prepared<A, M>(
351 &mut self,
352 module: &M,
353 infos: &A,
354 ) -> (GLWEAutomorphismKeyPrepared<&mut [u8], B>, &mut Self)
355 where
356 A: GGLWEInfos,
357 M: ModuleN + VmpPMatBytesOf,
358 {
359 assert_eq!(module.n() as u32, infos.n());
360 let (data, scratch) = self.take_gglwe_prepared(module, infos);
361 (GLWEAutomorphismKeyPrepared { key: data, p: 0 }, scratch)
362 }
363
364 fn take_glwe_tensor_key<A, M>(&mut self, infos: &A) -> (GLWETensorKey<&mut [u8]>, &mut Self)
365 where
366 A: GGLWEInfos,
367 {
368 assert_eq!(
369 infos.rank_in(),
370 infos.rank_out(),
371 "rank_in != rank_out is not supported for GLWETensorKey"
372 );
373
374 let pairs: u32 = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1);
375 let mut ksk_infos: GGLWELayout = infos.gglwe_layout();
376 ksk_infos.rank_in = Rank(pairs);
377 let (data, scratch) = self.take_gglwe(&ksk_infos);
378 (GLWETensorKey(data), scratch)
379 }
380
381 fn take_glwe_tensor_key_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWETensorKeyPrepared<&mut [u8], B>, &mut Self)
382 where
383 A: GGLWEInfos,
384 M: ModuleN + VmpPMatBytesOf,
385 {
386 assert_eq!(module.n() as u32, infos.n());
387 assert_eq!(
388 infos.rank_in(),
389 infos.rank_out(),
390 "rank_in != rank_out is not supported for GGLWETensorKeyPrepared"
391 );
392
393 let pairs: u32 = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1);
394 let mut ksk_infos: GGLWELayout = infos.gglwe_layout();
395 ksk_infos.rank_in = Rank(pairs);
396 let (data, scratch) = self.take_gglwe_prepared(module, &ksk_infos);
397 (GLWETensorKeyPrepared(data), scratch)
398 }
399}
400
401impl<B: Backend> ScratchTakeCore<B> for Scratch<B> where Self: ScratchTakeBasic + ScratchAvailable + ScratchFromBytes<B> {}