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