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