poulpy/poulpy-core/src/tests/test_suite/keyswitch/ggsw_ct.rs
poulpy-core/src/tests/test_suite/keyswitch/ggsw_ct.rsRaw
1use poulpy_hal::{
2 api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
3 layouts::{Backend, Module, ScalarZnx, Scratch, ScratchOwned},
4 source::Source,
5};
6
7use crate::{
8 GGLWEToGGSWKeyEncryptSk, GGSWEncryptSk, GGSWKeyswitch, GGSWNoise, GLWESwitchingKeyEncryptSk, ScratchTakeCore,
9 encryption::SIGMA,
10 layouts::{
11 GGLWEToGGSWKey, GGLWEToGGSWKeyPrepared, GGLWEToGGSWKeyPreparedFactory, GGSW, GGSWInfos, GGSWLayout, GLWEInfos,
12 GLWESecret, GLWESecretPreparedFactory, GLWESwitchingKey, GLWESwitchingKeyLayout, GLWESwitchingKeyPreparedFactory,
13 GLWETensorKeyLayout,
14 prepared::{GLWESecretPrepared, GLWESwitchingKeyPrepared},
15 },
16 noise::noise_ggsw_keyswitch,
17};
18
19#[allow(clippy::too_many_arguments)]
20pub fn test_ggsw_keyswitch<BE: Backend>(module: &Module<BE>)
21where
22 Module<BE>: GGSWEncryptSk<BE>
23 + GLWESwitchingKeyEncryptSk<BE>
24 + GGLWEToGGSWKeyEncryptSk<BE>
25 + GGSWKeyswitch<BE>
26 + GLWESecretPreparedFactory<BE>
27 + GGLWEToGGSWKeyPreparedFactory<BE>
28 + GLWESwitchingKeyPreparedFactory<BE>
29 + GGSWNoise<BE>,
30 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
31 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
32{
33 let base2k_in: usize = 17;
34 let base2k_key: usize = 13;
35 let base2k_out: usize = base2k_in; // MUST BE SAME
36 let k_in: usize = 102;
37 let max_dsize: usize = k_in.div_ceil(base2k_key);
38
39 for rank in 1_usize..3 {
40 for dsize in 1..max_dsize + 1 {
41 let k_ksk: usize = k_in + base2k_key * dsize;
42 let k_tsk: usize = k_ksk;
43 let k_out: usize = k_ksk; // Better capture noise.
44
45 let n: usize = module.n();
46 let dnum_in: usize = k_in / base2k_in;
47 let dnum_ksk: usize = k_in.div_ceil(base2k_key * dsize);
48
49 let dsize_in: usize = 1;
50
51 let ggsw_in_infos: GGSWLayout = GGSWLayout {
52 n: n.into(),
53 base2k: base2k_in.into(),
54 k: k_in.into(),
55 dnum: dnum_in.into(),
56 dsize: dsize_in.into(),
57 rank: rank.into(),
58 };
59
60 let ggsw_out_infos: GGSWLayout = GGSWLayout {
61 n: n.into(),
62 base2k: base2k_out.into(),
63 k: k_out.into(),
64 dnum: dnum_in.into(),
65 dsize: dsize_in.into(),
66 rank: rank.into(),
67 };
68
69 let tsk_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
70 n: n.into(),
71 base2k: base2k_key.into(),
72 k: k_tsk.into(),
73 dnum: dnum_ksk.into(),
74 dsize: dsize.into(),
75 rank: rank.into(),
76 };
77
78 let ksk_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
79 n: n.into(),
80 base2k: base2k_key.into(),
81 k: k_ksk.into(),
82 dnum: dnum_ksk.into(),
83 dsize: dsize.into(),
84 rank_in: rank.into(),
85 rank_out: rank.into(),
86 };
87
88 let mut ggsw_in: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_in_infos);
89 let mut ggsw_out: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_out_infos);
90 let mut tsk: GGLWEToGGSWKey<Vec<u8>> = GGLWEToGGSWKey::alloc_from_infos(&tsk_infos);
91 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_apply_infos);
92 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
93
94 let mut source_xs: Source = Source::new([0u8; 32]);
95 let mut source_xe: Source = Source::new([0u8; 32]);
96 let mut source_xa: Source = Source::new([0u8; 32]);
97
98 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
99 GGSW::encrypt_sk_tmp_bytes(module, &ggsw_in_infos)
100 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_apply_infos)
101 | GGLWEToGGSWKey::encrypt_sk_tmp_bytes(module, &tsk_infos)
102 | GGSW::keyswitch_tmp_bytes(
103 module,
104 &ggsw_out_infos,
105 &ggsw_in_infos,
106 &ksk_apply_infos,
107 &tsk_infos,
108 ),
109 );
110
111 let var_xs: f64 = 0.5;
112
113 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
114 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
115
116 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
117 sk_in_prepared.prepare(module, &sk_in);
118
119 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
120 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
121
122 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
123 sk_out_prepared.prepare(module, &sk_out);
124
125 ksk.encrypt_sk(
126 module,
127 &sk_in,
128 &sk_out,
129 &mut source_xa,
130 &mut source_xe,
131 scratch.borrow(),
132 );
133 tsk.encrypt_sk(
134 module,
135 &sk_out,
136 &mut source_xa,
137 &mut source_xe,
138 scratch.borrow(),
139 );
140
141 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
142
143 ggsw_in.encrypt_sk(
144 module,
145 &pt_scalar,
146 &sk_in_prepared,
147 &mut source_xa,
148 &mut source_xe,
149 scratch.borrow(),
150 );
151
152 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
153 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
154 ksk_prepared.prepare(module, &ksk, scratch.borrow());
155
156 let mut tsk_prepared: GGLWEToGGSWKeyPrepared<Vec<u8>, BE> = GGLWEToGGSWKeyPrepared::alloc_from_infos(module, &tsk);
157 tsk_prepared.prepare(module, &tsk, scratch.borrow());
158
159 ggsw_out.keyswitch(
160 module,
161 &ggsw_in,
162 &ksk_prepared,
163 &tsk_prepared,
164 scratch.borrow(),
165 );
166
167 let max_noise = |col_j: usize| -> f64 {
168 noise_ggsw_keyswitch(
169 n as f64,
170 base2k_key * dsize,
171 col_j,
172 var_xs,
173 0f64,
174 SIGMA * SIGMA,
175 0f64,
176 rank as f64,
177 k_in,
178 k_ksk,
179 k_tsk,
180 ) + 0.5
181 };
182
183 for row in 0..ggsw_out.dnum().as_usize() {
184 for col in 0..ggsw_out.rank().as_usize() + 1 {
185 assert!(
186 ggsw_out
187 .noise(
188 module,
189 row,
190 col,
191 &pt_scalar,
192 &sk_out_prepared,
193 scratch.borrow()
194 )
195 .std()
196 .log2()
197 <= max_noise(col)
198 )
199 }
200 }
201 }
202 }
203}
204
205#[allow(clippy::too_many_arguments)]
206pub fn test_ggsw_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
207where
208 Module<BE>: GGSWEncryptSk<BE>
209 + GLWESwitchingKeyEncryptSk<BE>
210 + GGLWEToGGSWKeyEncryptSk<BE>
211 + GGSWKeyswitch<BE>
212 + GLWESecretPreparedFactory<BE>
213 + GGLWEToGGSWKeyPreparedFactory<BE>
214 + GLWESwitchingKeyPreparedFactory<BE>
215 + GGSWNoise<BE>,
216 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
217 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
218{
219 let base2k_out: usize = 17;
220 let base2k_key: usize = 13;
221 let k_out: usize = 102;
222 let max_dsize: usize = k_out.div_ceil(base2k_key);
223
224 for rank in 1_usize..3 {
225 for dsize in 1..max_dsize + 1 {
226 let k_ksk: usize = k_out + base2k_key * dsize;
227 let k_tsk: usize = k_ksk;
228
229 let n: usize = module.n();
230 let dnum_in: usize = k_out / base2k_out;
231 let dnum_ksk: usize = k_out.div_ceil(base2k_key * dsize);
232 let dsize_in: usize = 1;
233
234 let ggsw_out_infos: GGSWLayout = GGSWLayout {
235 n: n.into(),
236 base2k: base2k_out.into(),
237 k: k_out.into(),
238 dnum: dnum_in.into(),
239 dsize: dsize_in.into(),
240 rank: rank.into(),
241 };
242
243 let tsk_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
244 n: n.into(),
245 base2k: base2k_key.into(),
246 k: k_tsk.into(),
247 dnum: dnum_ksk.into(),
248 dsize: dsize.into(),
249 rank: rank.into(),
250 };
251
252 let ksk_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
253 n: n.into(),
254 base2k: base2k_key.into(),
255 k: k_ksk.into(),
256 dnum: dnum_ksk.into(),
257 dsize: dsize.into(),
258 rank_in: rank.into(),
259 rank_out: rank.into(),
260 };
261
262 let mut ggsw_out: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_out_infos);
263 let mut tsk: GGLWEToGGSWKey<Vec<u8>> = GGLWEToGGSWKey::alloc_from_infos(&tsk_infos);
264 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_apply_infos);
265 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
266
267 let mut source_xs: Source = Source::new([0u8; 32]);
268 let mut source_xe: Source = Source::new([0u8; 32]);
269 let mut source_xa: Source = Source::new([0u8; 32]);
270
271 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
272 GGSW::encrypt_sk_tmp_bytes(module, &ggsw_out_infos)
273 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_apply_infos)
274 | GGLWEToGGSWKey::encrypt_sk_tmp_bytes(module, &tsk_infos)
275 | GGSW::keyswitch_tmp_bytes(
276 module,
277 &ggsw_out_infos,
278 &ggsw_out_infos,
279 &ksk_apply_infos,
280 &tsk_infos,
281 ),
282 );
283
284 let var_xs: f64 = 0.5;
285
286 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
287 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
288
289 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
290 sk_in_prepared.prepare(module, &sk_in);
291
292 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
293 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
294
295 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
296 sk_out_prepared.prepare(module, &sk_out);
297
298 ksk.encrypt_sk(
299 module,
300 &sk_in,
301 &sk_out,
302 &mut source_xa,
303 &mut source_xe,
304 scratch.borrow(),
305 );
306 tsk.encrypt_sk(
307 module,
308 &sk_out,
309 &mut source_xa,
310 &mut source_xe,
311 scratch.borrow(),
312 );
313
314 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
315
316 ggsw_out.encrypt_sk(
317 module,
318 &pt_scalar,
319 &sk_in_prepared,
320 &mut source_xa,
321 &mut source_xe,
322 scratch.borrow(),
323 );
324
325 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
326 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
327 ksk_prepared.prepare(module, &ksk, scratch.borrow());
328
329 let mut tsk_prepared: GGLWEToGGSWKeyPrepared<Vec<u8>, BE> = GGLWEToGGSWKeyPrepared::alloc_from_infos(module, &tsk);
330 tsk_prepared.prepare(module, &tsk, scratch.borrow());
331
332 ggsw_out.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
333
334 let max_noise = |col_j: usize| -> f64 {
335 noise_ggsw_keyswitch(
336 n as f64,
337 base2k_key * dsize,
338 col_j,
339 var_xs,
340 0f64,
341 SIGMA * SIGMA,
342 0f64,
343 rank as f64,
344 k_out,
345 k_ksk,
346 k_tsk,
347 ) + 0.5
348 };
349
350 for row in 0..ggsw_out.dnum().as_usize() {
351 for col in 0..ggsw_out.rank().as_usize() + 1 {
352 assert!(
353 ggsw_out
354 .noise(
355 module,
356 row,
357 col,
358 &pt_scalar,
359 &sk_out_prepared,
360 scratch.borrow()
361 )
362 .std()
363 .log2()
364 <= max_noise(col)
365 )
366 }
367 }
368 }
369 }
370}