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, GGSWLayout, GLWESecret,
12 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 ggsw_out.assert_noise(module, &sk_out_prepared, &pt_scalar, &max_noise);
184 }
185 }
186}
187
188#[allow(clippy::too_many_arguments)]
189pub fn test_ggsw_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
190where
191 Module<BE>: GGSWEncryptSk<BE>
192 + GLWESwitchingKeyEncryptSk<BE>
193 + GGLWEToGGSWKeyEncryptSk<BE>
194 + GGSWKeyswitch<BE>
195 + GLWESecretPreparedFactory<BE>
196 + GGLWEToGGSWKeyPreparedFactory<BE>
197 + GLWESwitchingKeyPreparedFactory<BE>
198 + GGSWNoise<BE>,
199 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
200 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
201{
202 let base2k_out: usize = 17;
203 let base2k_key: usize = 13;
204 let k_out: usize = 102;
205 let max_dsize: usize = k_out.div_ceil(base2k_key);
206
207 for rank in 1_usize..3 {
208 for dsize in 1..max_dsize + 1 {
209 let k_ksk: usize = k_out + base2k_key * dsize;
210 let k_tsk: usize = k_ksk;
211
212 let n: usize = module.n();
213 let dnum_in: usize = k_out / base2k_out;
214 let dnum_ksk: usize = k_out.div_ceil(base2k_key * dsize);
215 let dsize_in: usize = 1;
216
217 let ggsw_out_infos: GGSWLayout = GGSWLayout {
218 n: n.into(),
219 base2k: base2k_out.into(),
220 k: k_out.into(),
221 dnum: dnum_in.into(),
222 dsize: dsize_in.into(),
223 rank: rank.into(),
224 };
225
226 let tsk_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
227 n: n.into(),
228 base2k: base2k_key.into(),
229 k: k_tsk.into(),
230 dnum: dnum_ksk.into(),
231 dsize: dsize.into(),
232 rank: rank.into(),
233 };
234
235 let ksk_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
236 n: n.into(),
237 base2k: base2k_key.into(),
238 k: k_ksk.into(),
239 dnum: dnum_ksk.into(),
240 dsize: dsize.into(),
241 rank_in: rank.into(),
242 rank_out: rank.into(),
243 };
244
245 let mut ggsw_out: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_out_infos);
246 let mut tsk: GGLWEToGGSWKey<Vec<u8>> = GGLWEToGGSWKey::alloc_from_infos(&tsk_infos);
247 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_apply_infos);
248 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
249
250 let mut source_xs: Source = Source::new([0u8; 32]);
251 let mut source_xe: Source = Source::new([0u8; 32]);
252 let mut source_xa: Source = Source::new([0u8; 32]);
253
254 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
255 GGSW::encrypt_sk_tmp_bytes(module, &ggsw_out_infos)
256 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_apply_infos)
257 | GGLWEToGGSWKey::encrypt_sk_tmp_bytes(module, &tsk_infos)
258 | GGSW::keyswitch_tmp_bytes(
259 module,
260 &ggsw_out_infos,
261 &ggsw_out_infos,
262 &ksk_apply_infos,
263 &tsk_infos,
264 ),
265 );
266
267 let var_xs: f64 = 0.5;
268
269 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
270 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
271
272 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
273 sk_in_prepared.prepare(module, &sk_in);
274
275 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
276 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
277
278 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
279 sk_out_prepared.prepare(module, &sk_out);
280
281 ksk.encrypt_sk(
282 module,
283 &sk_in,
284 &sk_out,
285 &mut source_xa,
286 &mut source_xe,
287 scratch.borrow(),
288 );
289 tsk.encrypt_sk(
290 module,
291 &sk_out,
292 &mut source_xa,
293 &mut source_xe,
294 scratch.borrow(),
295 );
296
297 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
298
299 ggsw_out.encrypt_sk(
300 module,
301 &pt_scalar,
302 &sk_in_prepared,
303 &mut source_xa,
304 &mut source_xe,
305 scratch.borrow(),
306 );
307
308 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
309 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
310 ksk_prepared.prepare(module, &ksk, scratch.borrow());
311
312 let mut tsk_prepared: GGLWEToGGSWKeyPrepared<Vec<u8>, BE> = GGLWEToGGSWKeyPrepared::alloc_from_infos(module, &tsk);
313 tsk_prepared.prepare(module, &tsk, scratch.borrow());
314
315 ggsw_out.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
316
317 let max_noise = |col_j: usize| -> f64 {
318 noise_ggsw_keyswitch(
319 n as f64,
320 base2k_key * dsize,
321 col_j,
322 var_xs,
323 0f64,
324 SIGMA * SIGMA,
325 0f64,
326 rank as f64,
327 k_out,
328 k_ksk,
329 k_tsk,
330 ) + 0.5
331 };
332
333 ggsw_out.assert_noise(module, &sk_out_prepared, &pt_scalar, &max_noise);
334 }
335 }
336}