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: usize = 12;
34 let k_in: usize = 54;
35 let dsize: usize = k_in.div_ceil(base2k);
36 for rank in 1_usize..3 {
37 for di in 1..dsize + 1 {
38 let k_ksk: usize = k_in + base2k * di;
39 let k_tsk: usize = k_ksk;
40 let k_out: usize = k_ksk; // Better capture noise.
41
42 let n: usize = module.n();
43 let dnum: usize = k_in.div_ceil(di * base2k);
44
45 let dsize_in: usize = 1;
46
47 let ggsw_in_infos: GGSWLayout = GGSWLayout {
48 n: n.into(),
49 base2k: base2k.into(),
50 k: k_in.into(),
51 dnum: dnum.into(),
52 dsize: dsize_in.into(),
53 rank: rank.into(),
54 };
55
56 let ggsw_out_infos: GGSWLayout = GGSWLayout {
57 n: n.into(),
58 base2k: base2k.into(),
59 k: k_out.into(),
60 dnum: dnum.into(),
61 dsize: dsize_in.into(),
62 rank: rank.into(),
63 };
64
65 let tsk_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
66 n: n.into(),
67 base2k: base2k.into(),
68 k: k_tsk.into(),
69 dnum: dnum.into(),
70 dsize: di.into(),
71 rank: rank.into(),
72 };
73
74 let ksk_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
75 n: n.into(),
76 base2k: base2k.into(),
77 k: k_ksk.into(),
78 dnum: dnum.into(),
79 dsize: di.into(),
80 rank_in: rank.into(),
81 rank_out: rank.into(),
82 };
83
84 let mut ggsw_in: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_in_infos);
85 let mut ggsw_out: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_out_infos);
86 let mut tsk: GGLWEToGGSWKey<Vec<u8>> = GGLWEToGGSWKey::alloc_from_infos(&tsk_infos);
87 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_apply_infos);
88 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
89
90 let mut source_xs: Source = Source::new([0u8; 32]);
91 let mut source_xe: Source = Source::new([0u8; 32]);
92 let mut source_xa: Source = Source::new([0u8; 32]);
93
94 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
95 GGSW::encrypt_sk_tmp_bytes(module, &ggsw_in_infos)
96 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_apply_infos)
97 | GGLWEToGGSWKey::encrypt_sk_tmp_bytes(module, &tsk_infos)
98 | GGSW::keyswitch_tmp_bytes(
99 module,
100 &ggsw_out_infos,
101 &ggsw_in_infos,
102 &ksk_apply_infos,
103 &tsk_infos,
104 ),
105 );
106
107 let var_xs: f64 = 0.5;
108
109 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
110 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
111
112 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
113 sk_in_prepared.prepare(module, &sk_in);
114
115 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
116 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
117
118 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
119 sk_out_prepared.prepare(module, &sk_out);
120
121 ksk.encrypt_sk(
122 module,
123 &sk_in,
124 &sk_out,
125 &mut source_xa,
126 &mut source_xe,
127 scratch.borrow(),
128 );
129 tsk.encrypt_sk(
130 module,
131 &sk_out,
132 &mut source_xa,
133 &mut source_xe,
134 scratch.borrow(),
135 );
136
137 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
138
139 ggsw_in.encrypt_sk(
140 module,
141 &pt_scalar,
142 &sk_in_prepared,
143 &mut source_xa,
144 &mut source_xe,
145 scratch.borrow(),
146 );
147
148 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
149 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
150 ksk_prepared.prepare(module, &ksk, scratch.borrow());
151
152 let mut tsk_prepared: GGLWEToGGSWKeyPrepared<Vec<u8>, BE> = GGLWEToGGSWKeyPrepared::alloc_from_infos(module, &tsk);
153 tsk_prepared.prepare(module, &tsk, scratch.borrow());
154
155 ggsw_out.keyswitch(
156 module,
157 &ggsw_in,
158 &ksk_prepared,
159 &tsk_prepared,
160 scratch.borrow(),
161 );
162
163 let max_noise = |col_j: usize| -> f64 {
164 noise_ggsw_keyswitch(
165 n as f64,
166 base2k * di,
167 col_j,
168 var_xs,
169 0f64,
170 SIGMA * SIGMA,
171 0f64,
172 rank as f64,
173 k_in,
174 k_ksk,
175 k_tsk,
176 ) + 0.5
177 };
178
179 ggsw_out.assert_noise(module, &sk_out_prepared, &pt_scalar, &max_noise);
180 }
181 }
182}
183
184#[allow(clippy::too_many_arguments)]
185pub fn test_ggsw_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
186where
187 Module<BE>: GGSWEncryptSk<BE>
188 + GLWESwitchingKeyEncryptSk<BE>
189 + GGLWEToGGSWKeyEncryptSk<BE>
190 + GGSWKeyswitch<BE>
191 + GLWESecretPreparedFactory<BE>
192 + GGLWEToGGSWKeyPreparedFactory<BE>
193 + GLWESwitchingKeyPreparedFactory<BE>
194 + GGSWNoise<BE>,
195 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
196 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
197{
198 let base2k: usize = 12;
199 let k_out: usize = 54;
200 let dsize: usize = k_out.div_ceil(base2k);
201 for rank in 1_usize..3 {
202 for di in 1..dsize + 1 {
203 let k_ksk: usize = k_out + base2k * di;
204 let k_tsk: usize = k_ksk;
205
206 let n: usize = module.n();
207 let dnum: usize = k_out.div_ceil(di * base2k);
208
209 let dsize_in: usize = 1;
210
211 let ggsw_out_infos: GGSWLayout = GGSWLayout {
212 n: n.into(),
213 base2k: base2k.into(),
214 k: k_out.into(),
215 dnum: dnum.into(),
216 dsize: dsize_in.into(),
217 rank: rank.into(),
218 };
219
220 let tsk_infos: GLWETensorKeyLayout = GLWETensorKeyLayout {
221 n: n.into(),
222 base2k: base2k.into(),
223 k: k_tsk.into(),
224 dnum: dnum.into(),
225 dsize: di.into(),
226 rank: rank.into(),
227 };
228
229 let ksk_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
230 n: n.into(),
231 base2k: base2k.into(),
232 k: k_ksk.into(),
233 dnum: dnum.into(),
234 dsize: di.into(),
235 rank_in: rank.into(),
236 rank_out: rank.into(),
237 };
238
239 let mut ggsw_out: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_out_infos);
240 let mut tsk: GGLWEToGGSWKey<Vec<u8>> = GGLWEToGGSWKey::alloc_from_infos(&tsk_infos);
241 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_apply_infos);
242 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
243
244 let mut source_xs: Source = Source::new([0u8; 32]);
245 let mut source_xe: Source = Source::new([0u8; 32]);
246 let mut source_xa: Source = Source::new([0u8; 32]);
247
248 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
249 GGSW::encrypt_sk_tmp_bytes(module, &ggsw_out_infos)
250 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_apply_infos)
251 | GGLWEToGGSWKey::encrypt_sk_tmp_bytes(module, &tsk_infos)
252 | GGSW::keyswitch_tmp_bytes(
253 module,
254 &ggsw_out_infos,
255 &ggsw_out_infos,
256 &ksk_apply_infos,
257 &tsk_infos,
258 ),
259 );
260
261 let var_xs: f64 = 0.5;
262
263 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
264 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
265
266 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
267 sk_in_prepared.prepare(module, &sk_in);
268
269 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
270 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
271
272 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
273 sk_out_prepared.prepare(module, &sk_out);
274
275 ksk.encrypt_sk(
276 module,
277 &sk_in,
278 &sk_out,
279 &mut source_xa,
280 &mut source_xe,
281 scratch.borrow(),
282 );
283 tsk.encrypt_sk(
284 module,
285 &sk_out,
286 &mut source_xa,
287 &mut source_xe,
288 scratch.borrow(),
289 );
290
291 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
292
293 ggsw_out.encrypt_sk(
294 module,
295 &pt_scalar,
296 &sk_in_prepared,
297 &mut source_xa,
298 &mut source_xe,
299 scratch.borrow(),
300 );
301
302 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
303 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
304 ksk_prepared.prepare(module, &ksk, scratch.borrow());
305
306 let mut tsk_prepared: GGLWEToGGSWKeyPrepared<Vec<u8>, BE> = GGLWEToGGSWKeyPrepared::alloc_from_infos(module, &tsk);
307 tsk_prepared.prepare(module, &tsk, scratch.borrow());
308
309 ggsw_out.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
310
311 let max_noise = |col_j: usize| -> f64 {
312 noise_ggsw_keyswitch(
313 n as f64,
314 base2k * di,
315 col_j,
316 var_xs,
317 0f64,
318 SIGMA * SIGMA,
319 0f64,
320 rank as f64,
321 k_out,
322 k_ksk,
323 k_tsk,
324 ) + 0.5
325 };
326
327 ggsw_out.assert_noise(module, &sk_out_prepared, &pt_scalar, &max_noise);
328 }
329 }
330}