poulpy/core/src/lib.rs
core/src/lib.rsRaw
1pub mod blind_rotation;
2pub mod dist;
3pub mod elem;
4pub mod fourier_glwe;
5pub mod gglwe;
6pub mod ggsw;
7pub mod glwe;
8pub mod lwe;
9pub mod noise;
10
11use backend::Backend;
12use backend::FFT64;
13use backend::Module;
14pub use elem::{GetRow, Infos, SetMetaData, SetRow};
15pub use fourier_glwe::{FourierGLWECiphertext, FourierGLWESecret};
16pub use gglwe::{GGLWECiphertext, GLWEAutomorphismKey, GLWESwitchingKey, GLWETensorKey};
17pub use ggsw::GGSWCiphertext;
18pub use glwe::{GLWECiphertext, GLWEOps, GLWEPacker, GLWEPlaintext, GLWEPublicKey, GLWESecret};
19pub(crate) use glwe::{GLWECiphertextToMut, GLWECiphertextToRef};
20pub use lwe::{LWECiphertext, LWESecret};
21
22pub use backend;
23pub use backend::Scratch;
24pub use backend::ScratchOwned;
25
26use crate::dist::Distribution;
27
28pub(crate) const SIX_SIGMA: f64 = 6.0;
29
30pub trait ScratchCore<B: Backend> {
31 fn tmp_glwe_ct(&mut self, module: &Module<B>, basek: usize, k: usize, rank: usize) -> (GLWECiphertext<&mut [u8]>, &mut Self);
32 fn tmp_vec_glwe_ct(
33 &mut self,
34 size: usize,
35 module: &Module<B>,
36 basek: usize,
37 k: usize,
38 rank: usize,
39 ) -> (Vec<GLWECiphertext<&mut [u8]>>, &mut Self);
40 fn tmp_glwe_pt(&mut self, module: &Module<B>, basek: usize, k: usize) -> (GLWEPlaintext<&mut [u8]>, &mut Self);
41 fn tmp_gglwe(
42 &mut self,
43 module: &Module<B>,
44 basek: usize,
45 k: usize,
46 rows: usize,
47 digits: usize,
48 rank_in: usize,
49 rank_out: usize,
50 ) -> (GGLWECiphertext<&mut [u8], B>, &mut Self);
51 fn tmp_ggsw(
52 &mut self,
53 module: &Module<B>,
54 basek: usize,
55 k: usize,
56 rows: usize,
57 digits: usize,
58 rank: usize,
59 ) -> (GGSWCiphertext<&mut [u8], B>, &mut Self);
60 fn tmp_fourier_glwe_ct(
61 &mut self,
62 module: &Module<B>,
63 basek: usize,
64 k: usize,
65 rank: usize,
66 ) -> (FourierGLWECiphertext<&mut [u8], B>, &mut Self);
67 fn tmp_vec_fourier_glwe_ct(
68 &mut self,
69 size: usize,
70 module: &Module<B>,
71 basek: usize,
72 k: usize,
73 rank: usize,
74 ) -> (Vec<FourierGLWECiphertext<&mut [u8], B>>, &mut Self);
75 fn tmp_sk(&mut self, module: &Module<B>, rank: usize) -> (GLWESecret<&mut [u8]>, &mut Self);
76 fn tmp_fourier_sk(&mut self, module: &Module<B>, rank: usize) -> (FourierGLWESecret<&mut [u8], B>, &mut Self);
77 fn tmp_glwe_pk(
78 &mut self,
79 module: &Module<B>,
80 basek: usize,
81 k: usize,
82 rank: usize,
83 ) -> (GLWEPublicKey<&mut [u8], B>, &mut Self);
84 fn tmp_glwe_ksk(
85 &mut self,
86 module: &Module<B>,
87 basek: usize,
88 k: usize,
89 rows: usize,
90 digits: usize,
91 rank_in: usize,
92 rank_out: usize,
93 ) -> (GLWESwitchingKey<&mut [u8], B>, &mut Self);
94 fn tmp_tsk(
95 &mut self,
96 module: &Module<B>,
97 basek: usize,
98 k: usize,
99 rows: usize,
100 digits: usize,
101 rank: usize,
102 ) -> (GLWETensorKey<&mut [u8], B>, &mut Self);
103 fn tmp_autokey(
104 &mut self,
105 module: &Module<B>,
106 basek: usize,
107 k: usize,
108 rows: usize,
109 digits: usize,
110 rank: usize,
111 ) -> (GLWEAutomorphismKey<&mut [u8], B>, &mut Self);
112}
113
114impl ScratchCore<FFT64> for Scratch {
115 fn tmp_glwe_ct(
116 &mut self,
117 module: &Module<FFT64>,
118 basek: usize,
119 k: usize,
120 rank: usize,
121 ) -> (GLWECiphertext<&mut [u8]>, &mut Self) {
122 let (data, scratch) = self.tmp_vec_znx(module, rank + 1, k.div_ceil(basek));
123 (GLWECiphertext { data, basek, k }, scratch)
124 }
125
126 fn tmp_vec_glwe_ct(
127 &mut self,
128 size: usize,
129 module: &Module<FFT64>,
130 basek: usize,
131 k: usize,
132 rank: usize,
133 ) -> (Vec<GLWECiphertext<&mut [u8]>>, &mut Self) {
134 let mut scratch: &mut Scratch = self;
135 let mut cts: Vec<GLWECiphertext<&mut [u8]>> = Vec::with_capacity(size);
136 for _ in 0..size {
137 let (ct, new_scratch) = scratch.tmp_glwe_ct(module, basek, k, rank);
138 scratch = new_scratch;
139 cts.push(ct);
140 }
141 (cts, scratch)
142 }
143
144 fn tmp_glwe_pt(&mut self, module: &Module<FFT64>, basek: usize, k: usize) -> (GLWEPlaintext<&mut [u8]>, &mut Self) {
145 let (data, scratch) = self.tmp_vec_znx(module, 1, k.div_ceil(basek));
146 (GLWEPlaintext { data, basek, k }, scratch)
147 }
148
149 fn tmp_gglwe(
150 &mut self,
151 module: &Module<FFT64>,
152 basek: usize,
153 k: usize,
154 rows: usize,
155 digits: usize,
156 rank_in: usize,
157 rank_out: usize,
158 ) -> (GGLWECiphertext<&mut [u8], FFT64>, &mut Self) {
159 let (data, scratch) = self.tmp_mat_znx_dft(
160 module,
161 rows.div_ceil(digits),
162 rank_in,
163 rank_out + 1,
164 k.div_ceil(basek),
165 );
166 (
167 GGLWECiphertext {
168 data: data,
169 basek: basek,
170 k,
171 digits,
172 },
173 scratch,
174 )
175 }
176
177 fn tmp_ggsw(
178 &mut self,
179 module: &Module<FFT64>,
180 basek: usize,
181 k: usize,
182 rows: usize,
183 digits: usize,
184 rank: usize,
185 ) -> (GGSWCiphertext<&mut [u8], FFT64>, &mut Self) {
186 let (data, scratch) = self.tmp_mat_znx_dft(
187 module,
188 rows.div_ceil(digits),
189 rank + 1,
190 rank + 1,
191 k.div_ceil(basek),
192 );
193 (
194 GGSWCiphertext {
195 data,
196 basek,
197 k,
198 digits,
199 },
200 scratch,
201 )
202 }
203
204 fn tmp_fourier_glwe_ct(
205 &mut self,
206 module: &Module<FFT64>,
207 basek: usize,
208 k: usize,
209 rank: usize,
210 ) -> (FourierGLWECiphertext<&mut [u8], FFT64>, &mut Self) {
211 let (data, scratch) = self.tmp_vec_znx_dft(module, rank + 1, k.div_ceil(basek));
212 (FourierGLWECiphertext { data, basek, k }, scratch)
213 }
214
215 fn tmp_vec_fourier_glwe_ct(
216 &mut self,
217 size: usize,
218 module: &Module<FFT64>,
219 basek: usize,
220 k: usize,
221 rank: usize,
222 ) -> (Vec<FourierGLWECiphertext<&mut [u8], FFT64>>, &mut Self) {
223 let mut scratch: &mut Scratch = self;
224 let mut cts: Vec<FourierGLWECiphertext<&mut [u8], FFT64>> = Vec::with_capacity(size);
225 for _ in 0..size {
226 let (ct, new_scratch) = scratch.tmp_fourier_glwe_ct(module, basek, k, rank);
227 scratch = new_scratch;
228 cts.push(ct);
229 }
230 (cts, scratch)
231 }
232
233 fn tmp_glwe_pk(
234 &mut self,
235 module: &Module<FFT64>,
236 basek: usize,
237 k: usize,
238 rank: usize,
239 ) -> (GLWEPublicKey<&mut [u8], FFT64>, &mut Self) {
240 let (data, scratch) = self.tmp_fourier_glwe_ct(module, basek, k, rank);
241 (
242 GLWEPublicKey {
243 data,
244 dist: Distribution::NONE,
245 },
246 scratch,
247 )
248 }
249
250 fn tmp_sk(&mut self, module: &Module<FFT64>, rank: usize) -> (GLWESecret<&mut [u8]>, &mut Self) {
251 let (data, scratch) = self.tmp_scalar_znx(module, rank);
252 (
253 GLWESecret {
254 data,
255 dist: Distribution::NONE,
256 },
257 scratch,
258 )
259 }
260
261 fn tmp_fourier_sk(&mut self, module: &Module<FFT64>, rank: usize) -> (FourierGLWESecret<&mut [u8], FFT64>, &mut Self) {
262 let (data, scratch) = self.tmp_scalar_znx_dft(module, rank);
263 (
264 FourierGLWESecret {
265 data,
266 dist: Distribution::NONE,
267 },
268 scratch,
269 )
270 }
271
272 fn tmp_glwe_ksk(
273 &mut self,
274 module: &Module<FFT64>,
275 basek: usize,
276 k: usize,
277 rows: usize,
278 digits: usize,
279 rank_in: usize,
280 rank_out: usize,
281 ) -> (GLWESwitchingKey<&mut [u8], FFT64>, &mut Self) {
282 let (data, scratch) = self.tmp_gglwe(module, basek, k, rows, digits, rank_in, rank_out);
283 (GLWESwitchingKey(data), scratch)
284 }
285
286 fn tmp_autokey(
287 &mut self,
288 module: &Module<FFT64>,
289 basek: usize,
290 k: usize,
291 rows: usize,
292 digits: usize,
293 rank: usize,
294 ) -> (GLWEAutomorphismKey<&mut [u8], FFT64>, &mut Self) {
295 let (data, scratch) = self.tmp_glwe_ksk(module, basek, k, rows, digits, rank, rank);
296 (GLWEAutomorphismKey { key: data, p: 0 }, scratch)
297 }
298
299 fn tmp_tsk(
300 &mut self,
301 module: &Module<FFT64>,
302 basek: usize,
303 k: usize,
304 rows: usize,
305 digits: usize,
306 rank: usize,
307 ) -> (GLWETensorKey<&mut [u8], FFT64>, &mut Self) {
308 let mut keys: Vec<GLWESwitchingKey<&mut [u8], FFT64>> = Vec::new();
309 let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
310
311 let mut scratch: &mut Scratch = self;
312
313 if pairs != 0 {
314 let (gglwe, s) = scratch.tmp_glwe_ksk(module, basek, k, rows, digits, 1, rank);
315 scratch = s;
316 keys.push(gglwe);
317 }
318 for _ in 1..pairs {
319 let (gglwe, s) = scratch.tmp_glwe_ksk(module, basek, k, rows, digits, 1, rank);
320 scratch = s;
321 keys.push(gglwe);
322 }
323 (GLWETensorKey { keys }, scratch)
324 }
325}