poulpy/poulpy-core/src/tests/test_suite/keyswitch/gglwe_ct.rs
poulpy-core/src/tests/test_suite/keyswitch/gglwe_ct.rsRaw
1use poulpy_hal::{
2 api::{ScratchAvailable, ScratchOwnedAlloc, ScratchOwnedBorrow},
3 layouts::{Backend, Module, Scratch, ScratchOwned},
4 source::Source,
5};
6
7use crate::{
8 GGLWEKeyswitch, GGLWENoise, GLWESwitchingKeyEncryptSk, ScratchTakeCore,
9 encryption::SIGMA,
10 layouts::{
11 GGLWEInfos, GLWESecret, GLWESecretPreparedFactory, GLWESwitchingKey, GLWESwitchingKeyLayout,
12 GLWESwitchingKeyPreparedFactory,
13 prepared::{GLWESecretPrepared, GLWESwitchingKeyPrepared},
14 },
15 noise::log2_std_noise_gglwe_product,
16 var_noise_gglwe_product_v2,
17};
18
19pub fn test_gglwe_switching_key_keyswitch<BE: Backend>(module: &Module<BE>)
20where
21 Module<BE>: GLWESwitchingKeyEncryptSk<BE>
22 + GGLWEKeyswitch<BE>
23 + GLWESwitchingKeyPreparedFactory<BE>
24 + GLWESecretPreparedFactory<BE>
25 + GGLWENoise<BE>,
26 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
27 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
28{
29 let in_base2k: usize = 17;
30 let key_base2k: usize = 13;
31 let out_base2k: usize = in_base2k; // MUST BE SAME
32 let k_in: usize = 102;
33 let max_dsize: usize = k_in.div_ceil(key_base2k);
34
35 for rank_in_s0s1 in 1_usize..2 {
36 for rank_out_s0s1 in 1_usize..3 {
37 for rank_out_s1s2 in 1_usize..3 {
38 for dsize in 1_usize..max_dsize + 1 {
39 let k_ksk: usize = k_in + key_base2k * dsize;
40 let k_out: usize = k_ksk; // Better capture noise.
41
42 let n: usize = module.n();
43 let dsize_in: usize = 1;
44 let dnum_in: usize = k_in / in_base2k;
45 let dnum_ksk: usize = k_in.div_ceil(key_base2k * dsize);
46
47 let gglwe_s0s1_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
48 n: n.into(),
49 base2k: in_base2k.into(),
50 k: k_in.into(),
51 dnum: dnum_in.into(),
52 dsize: dsize_in.into(),
53 rank_in: rank_in_s0s1.into(),
54 rank_out: rank_out_s0s1.into(),
55 };
56
57 let gglwe_s1s2_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
58 n: n.into(),
59 base2k: key_base2k.into(),
60 k: k_ksk.into(),
61 dnum: dnum_ksk.into(),
62 dsize: dsize.into(),
63 rank_in: rank_out_s0s1.into(),
64 rank_out: rank_out_s1s2.into(),
65 };
66
67 let gglwe_s0s2_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
68 n: n.into(),
69 base2k: out_base2k.into(),
70 k: k_out.into(),
71 dnum: dnum_in.into(),
72 dsize: dsize_in.into(),
73 rank_in: rank_in_s0s1.into(),
74 rank_out: rank_out_s1s2.into(),
75 };
76
77 let mut gglwe_s0s1: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s0s1_infos);
78 let mut gglwe_s1s2: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s1s2_infos);
79 let mut gglwe_s0s2: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s0s2_infos);
80
81 let mut source_xs: Source = Source::new([0u8; 32]);
82 let mut source_xe: Source = Source::new([0u8; 32]);
83 let mut source_xa: Source = Source::new([0u8; 32]);
84
85 let mut scratch_enc: ScratchOwned<BE> = ScratchOwned::alloc(
86 GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s0s1_infos)
87 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s1s2_infos)
88 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s0s2_infos),
89 );
90 let mut scratch_apply: ScratchOwned<BE> = ScratchOwned::alloc(GLWESwitchingKey::keyswitch_tmp_bytes(
91 module,
92 &gglwe_s0s2_infos,
93 &gglwe_s0s1_infos,
94 &gglwe_s1s2_infos,
95 ));
96
97 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in_s0s1.into());
98 sk0.fill_ternary_prob(0.5, &mut source_xs);
99
100 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out_s0s1.into());
101 sk1.fill_ternary_prob(0.5, &mut source_xs);
102
103 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out_s1s2.into());
104 sk2.fill_ternary_prob(0.5, &mut source_xs);
105
106 let mut sk2_prepared: GLWESecretPrepared<Vec<u8>, BE> =
107 GLWESecretPrepared::alloc(module, rank_out_s1s2.into());
108 sk2_prepared.prepare(module, &sk2);
109
110 // gglwe_{s1}(s0) = s0 -> s1
111 gglwe_s0s1.encrypt_sk(module, &sk0, &sk1, &mut source_xa, &mut source_xe, scratch_enc.borrow());
112
113 // gglwe_{s2}(s1) -> s1 -> s2
114 gglwe_s1s2.encrypt_sk(module, &sk1, &sk2, &mut source_xa, &mut source_xe, scratch_enc.borrow());
115
116 let mut gglwe_s1s2_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
117 GLWESwitchingKeyPrepared::alloc_from_infos(module, &gglwe_s1s2);
118 gglwe_s1s2_prepared.prepare(module, &gglwe_s1s2, scratch_apply.borrow());
119
120 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
121 gglwe_s0s2.keyswitch(module, &gglwe_s0s1, &gglwe_s1s2_prepared, scratch_apply.borrow());
122
123 let max_noise: f64 = var_noise_gglwe_product_v2(
124 module.n() as f64,
125 k_ksk,
126 dnum_ksk,
127 dsize,
128 key_base2k,
129 0.5,
130 0.5,
131 0f64,
132 SIGMA * SIGMA,
133 0f64,
134 rank_out_s0s1 as f64,
135 )
136 .sqrt()
137 .log2();
138
139 for row in 0..gglwe_s0s2.dnum().as_usize() {
140 for col in 0..gglwe_s0s2.rank_in().as_usize() {
141 let noise: f64 = gglwe_s0s2
142 .key
143 .noise(module, row, col, &sk0.data, &sk2_prepared, scratch_apply.borrow())
144 .std()
145 .log2();
146 assert!(noise <= max_noise + 0.5, "{noise} > {max_noise}",)
147 }
148 }
149 }
150 }
151 }
152 }
153}
154
155#[allow(clippy::too_many_arguments)]
156pub fn test_gglwe_switching_key_keyswitch_inplace<BE: Backend>(module: &Module<BE>)
157where
158 Module<BE>: GLWESwitchingKeyEncryptSk<BE>
159 + GGLWEKeyswitch<BE>
160 + GLWESecretPreparedFactory<BE>
161 + GGLWENoise<BE>
162 + GLWESwitchingKeyPreparedFactory<BE>,
163 ScratchOwned<BE>: ScratchOwnedAlloc<BE> + ScratchOwnedBorrow<BE>,
164 Scratch<BE>: ScratchAvailable + ScratchTakeCore<BE>,
165{
166 let out_base2k: usize = 17;
167 let key_base2k: usize = 13;
168 let k_out: usize = 102;
169 let max_dsize: usize = k_out.div_ceil(key_base2k);
170
171 for rank_in in 1_usize..3 {
172 for rank_out in 1_usize..3 {
173 for dsize in 1_usize..max_dsize + 1 {
174 let k_ksk: usize = k_out + key_base2k * dsize;
175
176 let n: usize = module.n();
177 let dsize_in: usize = 1;
178
179 let dnum_in: usize = k_out / out_base2k;
180 let dnum_ksk: usize = k_out.div_ceil(key_base2k * dsize);
181
182 let gglwe_s0s1_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
183 n: n.into(),
184 base2k: out_base2k.into(),
185 k: k_out.into(),
186 dnum: dnum_in.into(),
187 dsize: dsize_in.into(),
188 rank_in: rank_in.into(),
189 rank_out: rank_out.into(),
190 };
191
192 let gglwe_s1s2_infos: GLWESwitchingKeyLayout = GLWESwitchingKeyLayout {
193 n: n.into(),
194 base2k: key_base2k.into(),
195 k: k_ksk.into(),
196 dnum: dnum_ksk.into(),
197 dsize: dsize.into(),
198 rank_in: rank_out.into(),
199 rank_out: rank_out.into(),
200 };
201
202 let mut gglwe_s0s1: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s0s1_infos);
203 let mut gglwe_s1s2: GLWESwitchingKey<Vec<u8>> = GLWESwitchingKey::alloc_from_infos(&gglwe_s1s2_infos);
204
205 let mut source_xs: Source = Source::new([0u8; 32]);
206 let mut source_xe: Source = Source::new([0u8; 32]);
207 let mut source_xa: Source = Source::new([0u8; 32]);
208
209 let mut scratch_enc: ScratchOwned<BE> = ScratchOwned::alloc(
210 GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s0s1_infos)
211 | GLWESwitchingKey::encrypt_sk_tmp_bytes(module, &gglwe_s1s2_infos),
212 );
213 let mut scratch_apply: ScratchOwned<BE> = ScratchOwned::alloc(GLWESwitchingKey::keyswitch_tmp_bytes(
214 module,
215 &gglwe_s0s1_infos,
216 &gglwe_s0s1_infos,
217 &gglwe_s1s2_infos,
218 ));
219
220 let var_xs: f64 = 0.5;
221
222 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_in.into());
223 sk0.fill_ternary_prob(var_xs, &mut source_xs);
224
225 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
226 sk1.fill_ternary_prob(var_xs, &mut source_xs);
227
228 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n.into(), rank_out.into());
229 sk2.fill_ternary_prob(var_xs, &mut source_xs);
230
231 let mut sk2_prepared: GLWESecretPrepared<Vec<u8>, BE> = GLWESecretPrepared::alloc(module, rank_out.into());
232 sk2_prepared.prepare(module, &sk2);
233
234 // gglwe_{s1}(s0) = s0 -> s1
235 gglwe_s0s1.encrypt_sk(module, &sk0, &sk1, &mut source_xa, &mut source_xe, scratch_enc.borrow());
236
237 // gglwe_{s2}(s1) -> s1 -> s2
238 gglwe_s1s2.encrypt_sk(module, &sk1, &sk2, &mut source_xa, &mut source_xe, scratch_enc.borrow());
239
240 let mut gglwe_s1s2_prepared: GLWESwitchingKeyPrepared<Vec<u8>, BE> =
241 GLWESwitchingKeyPrepared::alloc_from_infos(module, &gglwe_s1s2);
242 gglwe_s1s2_prepared.prepare(module, &gglwe_s1s2, scratch_apply.borrow());
243
244 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
245 gglwe_s0s1.keyswitch_inplace(module, &gglwe_s1s2_prepared, scratch_apply.borrow());
246
247 let gglwe_s0s2: GLWESwitchingKey<Vec<u8>> = gglwe_s0s1;
248
249 let max_noise: f64 = log2_std_noise_gglwe_product(
250 n as f64,
251 key_base2k * dsize,
252 var_xs,
253 var_xs,
254 0f64,
255 SIGMA * SIGMA,
256 0f64,
257 rank_out as f64,
258 k_out,
259 k_ksk,
260 );
261
262 for row in 0..gglwe_s0s2.dnum().as_usize() {
263 for col in 0..gglwe_s0s2.rank_in().as_usize() {
264 let noise = gglwe_s0s2
265 .key
266 .noise(module, row, col, &sk0.data, &sk2_prepared, scratch_apply.borrow())
267 .std()
268 .log2();
269 assert!(noise <= max_noise + 0.5, "{noise} > {max_noise}")
270 }
271 }
272 }
273 }
274 }
275}