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, GGLWEToGGSWKeyLayout, GGLWEToGGSWKeyPrepared, GGLWEToGGSWKeyPreparedFactory, GGSW, GGSWInfos, GGSWLayout,
12 GLWEInfos, GLWESecret, GLWESecretPreparedFactory, GLWESwitchingKey, GLWESwitchingKeyLayout,
13 GLWESwitchingKeyPreparedFactory,
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 in_base2k: usize = 17;
34 let key_base2k: usize = 13;
35 let out_base2k: usize = in_base2k; // MUST BE SAME
36 let k_in: usize = 102;
37 let max_dsize: usize = k_in.div_ceil(key_base2k);
38
39 for rank in 1_usize..3 {
40 for dsize in 1..max_dsize + 1 {
41 let k_ksk: usize = k_in + key_base2k * 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 / in_base2k;
47 let dnum_ksk: usize = k_in.div_ceil(key_base2k * dsize);
48
49 let dsize_in: usize = 1;
50
51 let ggsw_in_infos: GGSWLayout = GGSWLayout {
52 n: n.into(),
53 base2k: in_base2k.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: out_base2k.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: GGLWEToGGSWKeyLayout = GGLWEToGGSWKeyLayout {
70 n: n.into(),
71 base2k: key_base2k.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: key_base2k.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(module, &ggsw_out_infos, &ggsw_in_infos, &ksk_apply_infos, &tsk_infos),
103 );
104
105 let var_xs: f64 = 0.5;
106
107 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
108 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
109
110 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
111 sk_in_prepared.prepare(module, &sk_in);
112
113 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
114 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
115
116 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
117 sk_out_prepared.prepare(module, &sk_out);
118
119 ksk.encrypt_sk(module, &sk_in, &sk_out, &mut source_xa, &mut source_xe, scratch.borrow());
120 tsk.encrypt_sk(module, &sk_out, &mut source_xa, &mut source_xe, scratch.borrow());
121
122 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
123
124 ggsw_in.encrypt_sk(
125 module,
126 &pt_scalar,
127 &sk_in_prepared,
128 &mut source_xa,
129 &mut source_xe,
130 scratch.borrow(),
131 );
132
133 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
134 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
135 ksk_prepared.prepare(module, &ksk, scratch.borrow());
136
137 let mut tsk_prepared: GGLWEToGGSWKeyPrepared<Vec<u8>, BE> = GGLWEToGGSWKeyPrepared::alloc_from_infos(module, &tsk);
138 tsk_prepared.prepare(module, &tsk, scratch.borrow());
139
140 ggsw_out.keyswitch(module, &ggsw_in, &ksk_prepared, &tsk_prepared, scratch.borrow());
141
142 let max_noise = |col_j: usize| -> f64 {
143 noise_ggsw_keyswitch(
144 n as f64,
145 key_base2k * dsize,
146 col_j,
147 var_xs,
148 0f64,
149 SIGMA * SIGMA,
150 0f64,
151 rank as f64,
152 k_in,
153 k_ksk,
154 k_tsk,
155 ) + 0.5
156 };
157
158 for row in 0..ggsw_out.dnum().as_usize() {
159 for col in 0..ggsw_out.rank().as_usize() + 1 {
160 assert!(
161 ggsw_out
162 .noise(module, row, col, &pt_scalar, &sk_out_prepared, scratch.borrow())
163 .std()
164 .log2()
165 <= max_noise(col)
166 )
167 }
168 }
169 }
170 }
171}
172
173#[allow(clippy::too_many_arguments)]
174pub fn test_ggsw_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
175where
176 Module<BE>: GGSWEncryptSk<BE>
177 + GLWESwitchingKeyEncryptSk<BE>
178 + GGLWEToGGSWKeyEncryptSk<BE>
179 + GGSWKeyswitch<BE>
180 + GLWESecretPreparedFactory<BE>
181 + GGLWEToGGSWKeyPreparedFactory<BE>
182 + GLWESwitchingKeyPreparedFactory<BE>
183 + GGSWNoise<BE>,
184 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
185 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
186{
187 let out_base2k: usize = 17;
188 let key_base2k: usize = 13;
189 let k_out: usize = 102;
190 let max_dsize: usize = k_out.div_ceil(key_base2k);
191
192 for rank in 1_usize..3 {
193 for dsize in 1..max_dsize + 1 {
194 let k_ksk: usize = k_out + key_base2k * dsize;
195 let k_tsk: usize = k_ksk;
196
197 let n: usize = module.n();
198 let dnum_in: usize = k_out / out_base2k;
199 let dnum_ksk: usize = k_out.div_ceil(key_base2k * dsize);
200 let dsize_in: usize = 1;
201
202 let ggsw_out_infos: GGSWLayout = GGSWLayout {
203 n: n.into(),
204 base2k: out_base2k.into(),
205 k: k_out.into(),
206 dnum: dnum_in.into(),
207 dsize: dsize_in.into(),
208 rank: rank.into(),
209 };
210
211 let tsk_infos: GGLWEToGGSWKeyLayout = GGLWEToGGSWKeyLayout {
212 n: n.into(),
213 base2k: key_base2k.into(),
214 k: k_tsk.into(),
215 dnum: dnum_ksk.into(),
216 dsize: dsize.into(),
217 rank: rank.into(),
218 };
219
220 let ksk_apply_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
221 n: n.into(),
222 base2k: key_base2k.into(),
223 k: k_ksk.into(),
224 dnum: dnum_ksk.into(),
225 dsize: dsize.into(),
226 rank_in: rank.into(),
227 rank_out: rank.into(),
228 };
229
230 let mut ggsw_out: GGSW<Vec<u8>> = GGSW::alloc_from_infos(&ggsw_out_infos);
231 let mut tsk: GGLWEToGGSWKey<Vec<u8>> = GGLWEToGGSWKey::alloc_from_infos(&tsk_infos);
232 let mut ksk: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&ksk_apply_infos);
233 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
234
235 let mut source_xs: Source = Source::new([0u8; 32]);
236 let mut source_xe: Source = Source::new([0u8; 32]);
237 let mut source_xa: Source = Source::new([0u8; 32]);
238
239 let mut scratch: ScratchOwned<BE> = ScratchOwned::alloc(
240 GGSW::encrypt_sk_tmp_bytes(module, &ggsw_out_infos)
241 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &ksk_apply_infos)
242 | GGLWEToGGSWKey::encrypt_sk_tmp_bytes(module, &tsk_infos)
243 | GGSW::keyswitch_tmp_bytes(module, &ggsw_out_infos, &ggsw_out_infos, &ksk_apply_infos, &tsk_infos),
244 );
245
246 let var_xs: f64 = 0.5;
247
248 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
249 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
250
251 let mut sk_in_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
252 sk_in_prepared.prepare(module, &sk_in);
253
254 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank.into());
255 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
256
257 let mut sk_out_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank.into());
258 sk_out_prepared.prepare(module, &sk_out);
259
260 ksk.encrypt_sk(module, &sk_in, &sk_out, &mut source_xa, &mut source_xe, scratch.borrow());
261 tsk.encrypt_sk(module, &sk_out, &mut source_xa, &mut source_xe, scratch.borrow());
262
263 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
264
265 ggsw_out.encrypt_sk(
266 module,
267 &pt_scalar,
268 &sk_in_prepared,
269 &mut source_xa,
270 &mut source_xe,
271 scratch.borrow(),
272 );
273
274 let mut ksk_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
275 GLWESwitchingKeyPrepared::alloc_from_infos(module, &ksk);
276 ksk_prepared.prepare(module, &ksk, scratch.borrow());
277
278 let mut tsk_prepared: GGLWEToGGSWKeyPrepared<Vec<u8>, BE> = GGLWEToGGSWKeyPrepared::alloc_from_infos(module, &tsk);
279 tsk_prepared.prepare(module, &tsk, scratch.borrow());
280
281 ggsw_out.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
282
283 let max_noise = |col_j: usize| -> f64 {
284 noise_ggsw_keyswitch(
285 n as f64,
286 key_base2k * dsize,
287 col_j,
288 var_xs,
289 0f64,
290 SIGMA * SIGMA,
291 0f64,
292 rank as f64,
293 k_out,
294 k_ksk,
295 k_tsk,
296 ) + 0.5
297 };
298
299 for row in 0..ggsw_out.dnum().as_usize() {
300 for col in 0..ggsw_out.rank().as_usize() + 1 {
301 assert!(
302 ggsw_out
303 .noise(module, row, col, &pt_scalar, &sk_out_prepared, scratch.borrow())
304 .std()
305 .log2()
306 <= max_noise(col)
307 )
308 }
309 }
310 }
311 }
312}