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