poulpy/poulpy-core/src/scratch.rs
poulpy-core/src/scratch.rsRaw
1use poulpy_hal::{
2 api::{ScratchAvailable, ScratchTakeBasic},
3 layouts::{Backend, Module, Scratch},
4};
5
6use crate::{
7 dist::Distribution,
8 layouts::{
9 AutomorphismKey, Degree, GGLWE, GGLWEInfos, GGSW, GGSWInfos, GLWE, GLWEInfos, GLWEPlaintext, GLWEPublicKey, GLWESecret,
10 GLWESwitchingKey, GetDegree, Rank, TensorKey,
11 prepared::{
12 AutomorphismKeyPrepared, GGLWEPrepared, GGSWPrepared, GLWEPublicKeyPrepared, GLWESecretPrepared,
13 GLWESwitchingKeyPrepared, TensorKeyPrepared,
14 },
15 },
16};
17
18pub trait ScratchTakeCore<B: Backend>
19where
20 Self: ScratchTakeBasic<B> + ScratchAvailable,
21{
22 fn take_glwe_ct<A>(&mut self, module: &Module<B>, infos: &A) -> (GLWE<&mut [u8]>, &mut Self)
23 where
24 A: GLWEInfos,
25 {
26 let (data, scratch) = self.take_vec_znx(module, (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_ct_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 Scratch<B> = 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_ct(infos);
45 scratch = new_scratch;
46 cts.push(ct);
47 }
48 (cts, scratch)
49 }
50
51 fn take_glwe_pt<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>(&mut self, infos: &A) -> (GGLWEPrepared<&mut [u8], B>, &mut Self)
89 where
90 A: GGLWEInfos,
91 {
92 let (data, scratch) = self.take_vmp_pmat(
93 infos.n().into(),
94 infos.dnum().into(),
95 infos.rank_in().into(),
96 (infos.rank_out() + 1).into(),
97 infos.size(),
98 );
99 (
100 GGLWEPrepared {
101 k: infos.k(),
102 base2k: infos.base2k(),
103 dsize: infos.dsize(),
104 data,
105 },
106 scratch,
107 )
108 }
109
110 fn take_ggsw<A>(&mut self, infos: &A) -> (GGSW<&mut [u8]>, &mut Self)
111 where
112 A: GGSWInfos,
113 {
114 let (data, scratch) = self.take_mat_znx(
115 infos.n().into(),
116 infos.dnum().into(),
117 (infos.rank() + 1).into(),
118 (infos.rank() + 1).into(),
119 infos.size(),
120 );
121 (
122 GGSW {
123 k: infos.k(),
124 base2k: infos.base2k(),
125 dsize: infos.dsize(),
126 data,
127 },
128 scratch,
129 )
130 }
131
132 fn take_ggsw_prepared<A>(&mut self, infos: &A) -> (GGSWPrepared<&mut [u8], B>, &mut Self)
133 where
134 A: GGSWInfos,
135 {
136 let (data, scratch) = self.take_vmp_pmat(
137 infos.n().into(),
138 infos.dnum().into(),
139 (infos.rank() + 1).into(),
140 (infos.rank() + 1).into(),
141 infos.size(),
142 );
143 (
144 GGSWPrepared {
145 k: infos.k(),
146 base2k: infos.base2k(),
147 dsize: infos.dsize(),
148 data,
149 },
150 scratch,
151 )
152 }
153
154 fn take_ggsw_prepared_slice<A>(&mut self, size: usize, infos: &A) -> (Vec<GGSWPrepared<&mut [u8], B>>, &mut Self)
155 where
156 A: GGSWInfos,
157 {
158 let mut scratch: &mut Scratch<B> = self;
159 let mut cts: Vec<GGSWPrepared<&mut [u8], B>> = Vec::with_capacity(size);
160 for _ in 0..size {
161 let (ct, new_scratch) = scratch.take_ggsw_prepared(infos);
162 scratch = new_scratch;
163 cts.push(ct)
164 }
165 (cts, scratch)
166 }
167
168 fn take_glwe_pk<A>(&mut self, infos: &A) -> (GLWEPublicKey<&mut [u8]>, &mut Self)
169 where
170 A: GLWEInfos,
171 {
172 let (data, scratch) = self.take_vec_znx(infos.n().into(), (infos.rank() + 1).into(), infos.size());
173 (
174 GLWEPublicKey {
175 k: infos.k(),
176 dist: Distribution::NONE,
177 base2k: infos.base2k(),
178 data,
179 },
180 scratch,
181 )
182 }
183
184 fn take_glwe_pk_prepared<A>(&mut self, infos: &A) -> (GLWEPublicKeyPrepared<&mut [u8], B>, &mut Self)
185 where
186 A: GLWEInfos,
187 {
188 let (data, scratch) = self.take_vec_znx_dft(infos.n().into(), (infos.rank() + 1).into(), infos.size());
189 (
190 GLWEPublicKeyPrepared {
191 k: infos.k(),
192 dist: Distribution::NONE,
193 base2k: infos.base2k(),
194 data,
195 },
196 scratch,
197 )
198 }
199
200 fn take_glwe_secret(&mut self, n: Degree, rank: Rank) -> (GLWESecret<&mut [u8]>, &mut Self) {
201 let (data, scratch) = self.take_scalar_znx(n.into(), rank.into());
202 (
203 GLWESecret {
204 data,
205 dist: Distribution::NONE,
206 },
207 scratch,
208 )
209 }
210
211 fn take_glwe_secret_prepared(&mut self, n: Degree, rank: Rank) -> (GLWESecretPrepared<&mut [u8], B>, &mut Self) {
212 let (data, scratch) = self.take_svp_ppol(n.into(), rank.into());
213 (
214 GLWESecretPrepared {
215 data,
216 dist: Distribution::NONE,
217 },
218 scratch,
219 )
220 }
221
222 fn take_glwe_switching_key<A>(&mut self, infos: &A) -> (GLWESwitchingKey<&mut [u8]>, &mut Self)
223 where
224 A: GGLWEInfos,
225 {
226 let (data, scratch) = self.take_gglwe(infos);
227 (
228 GLWESwitchingKey {
229 key: data,
230 sk_in_n: 0,
231 sk_out_n: 0,
232 },
233 scratch,
234 )
235 }
236
237 fn take_gglwe_switching_key_prepared<A>(&mut self, infos: &A) -> (GLWESwitchingKeyPrepared<&mut [u8], B>, &mut Self)
238 where
239 A: GGLWEInfos,
240 {
241 let (data, scratch) = self.take_gglwe_prepared(infos);
242 (
243 GLWESwitchingKeyPrepared {
244 key: data,
245 sk_in_n: 0,
246 sk_out_n: 0,
247 },
248 scratch,
249 )
250 }
251
252 fn take_gglwe_automorphism_key<A>(&mut self, infos: &A) -> (AutomorphismKey<&mut [u8]>, &mut Self)
253 where
254 A: GGLWEInfos,
255 {
256 let (data, scratch) = self.take_glwe_switching_key(infos);
257 (AutomorphismKey { key: data, p: 0 }, scratch)
258 }
259
260 fn take_gglwe_automorphism_key_prepared<A>(&mut self, infos: &A) -> (AutomorphismKeyPrepared<&mut [u8], B>, &mut Self)
261 where
262 A: GGLWEInfos,
263 {
264 let (data, scratch) = self.take_gglwe_switching_key_prepared(infos);
265 (AutomorphismKeyPrepared { key: data, p: 0 }, scratch)
266 }
267
268 fn take_tensor_key<A>(&mut self, infos: &A) -> (TensorKey<&mut [u8]>, &mut Self)
269 where
270 A: GGLWEInfos,
271 {
272 assert_eq!(
273 infos.rank_in(),
274 infos.rank_out(),
275 "rank_in != rank_out is not supported for GGLWETensorKey"
276 );
277 let mut keys: Vec<GLWESwitchingKey<&mut [u8]>> = Vec::new();
278 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
279
280 let mut scratch: &mut Scratch<B> = self;
281
282 let mut ksk_infos: crate::layouts::GGLWECiphertextLayout = infos.gglwe_layout();
283 ksk_infos.rank_in = Rank(1);
284
285 if pairs != 0 {
286 let (gglwe, s) = scratch.take_glwe_switching_key(&ksk_infos);
287 scratch = s;
288 keys.push(gglwe);
289 }
290 for _ in 1..pairs {
291 let (gglwe, s) = scratch.take_glwe_switching_key(&ksk_infos);
292 scratch = s;
293 keys.push(gglwe);
294 }
295 (TensorKey { keys }, scratch)
296 }
297
298 fn take_gglwe_tensor_key_prepared<A>(&mut self, infos: &A) -> (TensorKeyPrepared<&mut [u8], B>, &mut Self)
299 where
300 A: GGLWEInfos,
301 {
302 assert_eq!(
303 infos.rank_in(),
304 infos.rank_out(),
305 "rank_in != rank_out is not supported for GGLWETensorKeyPrepared"
306 );
307
308 let mut keys: Vec<GLWESwitchingKeyPrepared<&mut [u8], B>> = Vec::new();
309 let pairs: usize = (((infos.rank_out().0 + 1) * infos.rank_out().0) >> 1).max(1) as usize;
310
311 let mut scratch: &mut Scratch<B> = self;
312
313 let mut ksk_infos: crate::layouts::GGLWECiphertextLayout = infos.gglwe_layout();
314 ksk_infos.rank_in = Rank(1);
315
316 if pairs != 0 {
317 let (gglwe, s) = scratch.take_gglwe_switching_key_prepared(&ksk_infos);
318 scratch = s;
319 keys.push(gglwe);
320 }
321 for _ in 1..pairs {
322 let (gglwe, s) = scratch.take_gglwe_switching_key_prepared(&ksk_infos);
323 scratch = s;
324 keys.push(gglwe);
325 }
326 (TensorKeyPrepared { keys }, scratch)
327 }
328}
329
330impl<B: Backend> ScratchTakeCore<B> for Scratch<B> where Self: ScratchTakeBasic<B> + ScratchAvailable {}