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>(&mut self, size: usize, infos: &A) -> (Vec<GGSW<&mut [u8]>>, &mut Self)
159 where
160 A: GGSWInfos,
161 {
162 let mut scratch: &mut Self = self;
163 let mut cts: Vec<GGSW<&mut [u8]>> = Vec::with_capacity(size);
164 for _ in 0..size {
165 let (ct, new_scratch) = scratch.take_ggsw(infos);
166 scratch = new_scratch;
167 cts.push(ct)
168 }
169 (cts, scratch)
170 }
171
172 fn take_ggsw_prepared_slice<A, M>(
173 &mut self,
174 module: &M,
175 size: usize,
176 infos: &A,
177 ) -> (Vec<GGSWPrepared<&mut [u8], B>>, &mut Self)
178 where
179 A: GGSWInfos,
180 M: ModuleN + VmpPMatBytesOf,
181 {
182 let mut scratch: &mut Self = self;
183 let mut cts: Vec<GGSWPrepared<&mut [u8], B>> = Vec::with_capacity(size);
184 for _ in 0..size {
185 let (ct, new_scratch) = scratch.take_ggsw_prepared(module, infos);
186 scratch = new_scratch;
187 cts.push(ct)
188 }
189 (cts, scratch)
190 }
191
192 fn take_glwe_public_key<A>(&mut self, infos: &A) -> (GLWEPublicKey<&mut [u8]>, &mut Self)
193 where
194 A: GLWEInfos,
195 {
196 let (data, scratch) = self.take_glwe(infos);
197 (
198 GLWEPublicKey {
199 dist: Distribution::NONE,
200 key: data,
201 },
202 scratch,
203 )
204 }
205
206 fn take_glwe_public_key_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPublicKeyPrepared<&mut [u8], B>, &mut Self)
207 where
208 A: GLWEInfos,
209 M: ModuleN + VecZnxDftBytesOf,
210 {
211 let (data, scratch) = self.take_glwe_prepared(module, infos);
212 (
213 GLWEPublicKeyPrepared {
214 dist: Distribution::NONE,
215 key: data,
216 },
217 scratch,
218 )
219 }
220
221 fn take_glwe_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWEPrepared<&mut [u8], B>, &mut Self)
222 where
223 A: GLWEInfos,
224 M: ModuleN + VecZnxDftBytesOf,
225 {
226 assert_eq!(module.n() as u32, infos.n());
227 let (data, scratch) = self.take_vec_znx_dft(module, (infos.rank() + 1).into(), infos.size());
228 (
229 GLWEPrepared {
230 k: infos.k(),
231 base2k: infos.base2k(),
232 data,
233 },
234 scratch,
235 )
236 }
237
238 fn take_glwe_secret(&mut self, n: Degree, rank: Rank) -> (GLWESecret<&mut [u8]>, &mut Self) {
239 let (data, scratch) = self.take_scalar_znx(n.into(), rank.into());
240 (
241 GLWESecret {
242 data,
243 dist: Distribution::NONE,
244 },
245 scratch,
246 )
247 }
248
249 fn take_glwe_secret_tensor(&mut self, n: Degree, rank: Rank) -> (GLWESecretTensor<&mut [u8]>, &mut Self) {
250 let (data, scratch) = self.take_scalar_znx(n.into(), GLWESecretTensor::pairs(rank.into()));
251 (
252 GLWESecretTensor {
253 data,
254 rank,
255 dist: Distribution::NONE,
256 },
257 scratch,
258 )
259 }
260
261 fn take_glwe_secret_prepared<M>(&mut self, module: &M, rank: Rank) -> (GLWESecretPrepared<&mut [u8], B>, &mut Self)
262 where
263 M: ModuleN + SvpPPolBytesOf,
264 {
265 let (data, scratch) = self.take_svp_ppol(module, rank.into());
266 (
267 GLWESecretPrepared {
268 data,
269 dist: Distribution::NONE,
270 },
271 scratch,
272 )
273 }
274
275 fn take_glwe_switching_key<A>(&mut self, infos: &A) -> (GLWESwitchingKey<&mut [u8]>, &mut Self)
276 where
277 A: GGLWEInfos,
278 {
279 let (data, scratch) = self.take_gglwe(infos);
280 (
281 GLWESwitchingKey {
282 key: data,
283 input_degree: Degree(0),
284 output_degree: Degree(0),
285 },
286 scratch,
287 )
288 }
289
290 fn take_glwe_switching_key_prepared<A, M>(
291 &mut self,
292 module: &M,
293 infos: &A,
294 ) -> (GLWESwitchingKeyPrepared<&mut [u8], B>, &mut Self)
295 where
296 A: GGLWEInfos,
297 M: ModuleN + VmpPMatBytesOf,
298 {
299 assert_eq!(module.n() as u32, infos.n());
300 let (data, scratch) = self.take_gglwe_prepared(module, infos);
301 (
302 GLWESwitchingKeyPrepared {
303 key: data,
304 input_degree: Degree(0),
305 output_degree: Degree(0),
306 },
307 scratch,
308 )
309 }
310
311 fn take_glwe_automorphism_key<A>(&mut self, infos: &A) -> (GLWEAutomorphismKey<&mut [u8]>, &mut Self)
312 where
313 A: GGLWEInfos,
314 {
315 let (data, scratch) = self.take_gglwe(infos);
316 (GLWEAutomorphismKey { key: data, p: 0 }, scratch)
317 }
318
319 fn take_glwe_automorphism_key_prepared<A, M>(
320 &mut self,
321 module: &M,
322 infos: &A,
323 ) -> (GLWEAutomorphismKeyPrepared<&mut [u8], B>, &mut Self)
324 where
325 A: GGLWEInfos,
326 M: ModuleN + VmpPMatBytesOf,
327 {
328 assert_eq!(module.n() as u32, infos.n());
329 let (data, scratch) = self.take_gglwe_prepared(module, infos);
330 (GLWEAutomorphismKeyPrepared { key: data, p: 0 }, scratch)
331 }
332
333 fn take_glwe_tensor_key<A, M>(&mut self, infos: &A) -> (GLWETensorKey<&mut [u8]>, &mut Self)
334 where
335 A: GGLWEInfos,
336 {
337 assert_eq!(
338 infos.rank_in(),
339 infos.rank_out(),
340 "rank_in != rank_out is not supported for GLWETensorKey"
341 );
342
343 let pairs: u32 = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1);
344 let mut ksk_infos: GGLWELayout = infos.gglwe_layout();
345 ksk_infos.rank_in = Rank(pairs);
346 let (data, scratch) = self.take_gglwe(infos);
347 (GLWETensorKey(data), scratch)
348 }
349
350 fn take_glwe_tensor_key_prepared<A, M>(&mut self, module: &M, infos: &A) -> (GLWETensorKeyPrepared<&mut [u8], B>, &mut Self)
351 where
352 A: GGLWEInfos,
353 M: ModuleN + VmpPMatBytesOf,
354 {
355 assert_eq!(module.n() as u32, infos.n());
356 assert_eq!(
357 infos.rank_in(),
358 infos.rank_out(),
359 "rank_in != rank_out is not supported for GGLWETensorKeyPrepared"
360 );
361
362 let pairs: u32 = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1);
363 let mut ksk_infos: GGLWELayout = infos.gglwe_layout();
364 ksk_infos.rank_in = Rank(pairs);
365 let (data, scratch) = self.take_gglwe_prepared(module, infos);
366 (GLWETensorKeyPrepared(data), scratch)
367 }
368}
369
370impl<B: Backend> ScratchTakeCore<B> for Scratch<B> where Self: ScratchTakeBasic + ScratchAvailable {}