poulpy/core/src/tests/generics/keyswitch/gglwe_ct.rs
core/src/tests/generics/keyswitch/gglwe_ct.rsRaw
1use backend::hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxCopy, VecZnxSubScalarInplace, VecZnxSwithcDegree,
4 VmpPMatAlloc, VmpPMatPrepare,
5 },
6 layouts::{Backend, Module, ScratchOwned},
7 oep::{
8 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
9 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl,
10 },
11};
12use sampling::source::Source;
13
14use crate::{
15 layouts::{
16 GGLWESwitchingKey, GLWESecret,
17 prepared::{GGLWESwitchingKeyPrepared, GLWESecretPrepared, PrepareAlloc},
18 },
19 noise::log2_std_noise_gglwe_product,
20 trait_families::{GLWEDecryptFamily, GLWEKeyswitchFamily},
21};
22
23use crate::trait_families::{GGLWEEncryptSkFamily, GGLWESwitchingKeyEncryptSkFamily, GLWESecretPreparedModuleFamily};
24
25pub fn test_gglwe_switching_key_keyswitch<B: Backend>(
26 module: &Module<B>,
27 basek: usize,
28 k_out: usize,
29 k_in: usize,
30 k_ksk: usize,
31 digits: usize,
32 rank_in_s0s1: usize,
33 rank_out_s0s1: usize,
34 rank_out_s1s2: usize,
35 sigma: f64,
36) where
37 Module<B>: GGLWEEncryptSkFamily<B>
38 + GLWESecretPreparedModuleFamily<B>
39 + GLWEDecryptFamily<B>
40 + VecZnxSwithcDegree
41 + VecZnxAddScalarInplace
42 + VecZnxSubScalarInplace
43 + VecZnxCopy
44 + VmpPMatAlloc<B>
45 + VmpPMatPrepare<B>
46 + GGLWEEncryptSkFamily<B>
47 + GLWEDecryptFamily<B>
48 + GLWEKeyswitchFamily<B>,
49 B: TakeVecZnxDftImpl<B>
50 + TakeVecZnxBigImpl<B>
51 + TakeSvpPPolImpl<B>
52 + ScratchOwnedAllocImpl<B>
53 + ScratchOwnedBorrowImpl<B>
54 + ScratchAvailableImpl<B>
55 + TakeScalarZnxImpl<B>
56 + TakeVecZnxImpl<B>
57 + VecZnxDftAllocBytesImpl<B>
58 + VecZnxBigAllocBytesImpl<B>
59 + TakeSvpPPolImpl<B>,
60{
61 let n: usize = module.n();
62 let rows: usize = k_in.div_ceil(basek * digits);
63 let digits_in: usize = 1;
64
65 let mut ct_gglwe_s0s1: GGLWESwitchingKey<Vec<u8>> =
66 GGLWESwitchingKey::alloc(n, basek, k_in, rows, digits_in, rank_in_s0s1, rank_out_s0s1);
67 let mut ct_gglwe_s1s2: GGLWESwitchingKey<Vec<u8>> =
68 GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_out_s0s1, rank_out_s1s2);
69 let mut ct_gglwe_s0s2: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(
70 n,
71 basek,
72 k_out,
73 rows,
74 digits_in,
75 rank_in_s0s1,
76 rank_out_s1s2,
77 );
78
79 let mut source_xs: Source = Source::new([0u8; 32]);
80 let mut source_xe: Source = Source::new([0u8; 32]);
81 let mut source_xa: Source = Source::new([0u8; 32]);
82
83 let mut scratch_enc: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::encrypt_sk_scratch_space(
84 module,
85 n,
86 basek,
87 k_ksk,
88 rank_in_s0s1 | rank_out_s0s1,
89 rank_out_s0s1 | rank_out_s1s2,
90 ));
91 let mut scratch_apply: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::keyswitch_scratch_space(
92 module,
93 n,
94 basek,
95 k_out,
96 k_in,
97 k_ksk,
98 digits,
99 ct_gglwe_s1s2.rank_in(),
100 ct_gglwe_s1s2.rank_out(),
101 ));
102
103 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in_s0s1);
104 sk0.fill_ternary_prob(0.5, &mut source_xs);
105
106 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out_s0s1);
107 sk1.fill_ternary_prob(0.5, &mut source_xs);
108
109 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out_s1s2);
110 sk2.fill_ternary_prob(0.5, &mut source_xs);
111 let sk2_prepared: GLWESecretPrepared<Vec<u8>, B> = sk2.prepare_alloc(module, scratch_apply.borrow());
112
113 // gglwe_{s1}(s0) = s0 -> s1
114 ct_gglwe_s0s1.encrypt_sk(
115 module,
116 &sk0,
117 &sk1,
118 &mut source_xa,
119 &mut source_xe,
120 sigma,
121 scratch_enc.borrow(),
122 );
123
124 // gglwe_{s2}(s1) -> s1 -> s2
125 ct_gglwe_s1s2.encrypt_sk(
126 module,
127 &sk1,
128 &sk2,
129 &mut source_xa,
130 &mut source_xe,
131 sigma,
132 scratch_enc.borrow(),
133 );
134
135 let ct_gglwe_s1s2_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> =
136 ct_gglwe_s1s2.prepare_alloc(module, scratch_apply.borrow());
137
138 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
139 ct_gglwe_s0s2.keyswitch(
140 module,
141 &ct_gglwe_s0s1,
142 &ct_gglwe_s1s2_prepared,
143 scratch_apply.borrow(),
144 );
145
146 let max_noise: f64 = log2_std_noise_gglwe_product(
147 n as f64,
148 basek * digits,
149 0.5,
150 0.5,
151 0f64,
152 sigma * sigma,
153 0f64,
154 rank_out_s0s1 as f64,
155 k_in,
156 k_ksk,
157 );
158
159 ct_gglwe_s0s2
160 .key
161 .assert_noise(module, &sk2_prepared, &sk0.data, max_noise + 0.5);
162}
163
164pub fn test_gglwe_switching_key_keyswitch_inplace<B: Backend>(
165 module: &Module<B>,
166 basek: usize,
167 k_ct: usize,
168 k_ksk: usize,
169 digits: usize,
170 rank_in: usize,
171 rank_out: usize,
172 sigma: f64,
173) where
174 Module<B>: GGLWEEncryptSkFamily<B>
175 + GLWESecretPreparedModuleFamily<B>
176 + GLWEDecryptFamily<B>
177 + VecZnxSwithcDegree
178 + VecZnxAddScalarInplace
179 + VecZnxSubScalarInplace
180 + VecZnxCopy
181 + VmpPMatAlloc<B>
182 + VmpPMatPrepare<B>
183 + GGLWESwitchingKeyEncryptSkFamily<B>
184 + GLWEKeyswitchFamily<B>
185 + GLWEDecryptFamily<B>,
186 B: TakeVecZnxDftImpl<B>
187 + TakeVecZnxBigImpl<B>
188 + TakeSvpPPolImpl<B>
189 + ScratchOwnedAllocImpl<B>
190 + ScratchOwnedBorrowImpl<B>
191 + ScratchAvailableImpl<B>
192 + TakeScalarZnxImpl<B>
193 + TakeVecZnxImpl<B>
194 + VecZnxDftAllocBytesImpl<B>
195 + VecZnxBigAllocBytesImpl<B>
196 + TakeSvpPPolImpl<B>,
197{
198 let n: usize = module.n();
199 let rows: usize = k_ct.div_ceil(basek * digits);
200 let digits_in: usize = 1;
201
202 let mut ct_gglwe_s0s1: GGLWESwitchingKey<Vec<u8>> =
203 GGLWESwitchingKey::alloc(n, basek, k_ct, rows, digits_in, rank_in, rank_out);
204 let mut ct_gglwe_s1s2: GGLWESwitchingKey<Vec<u8>> =
205 GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_out, rank_out);
206
207 let mut source_xs: Source = Source::new([0u8; 32]);
208 let mut source_xe: Source = Source::new([0u8; 32]);
209 let mut source_xa: Source = Source::new([0u8; 32]);
210
211 let mut scratch_enc: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::encrypt_sk_scratch_space(
212 module,
213 n,
214 basek,
215 k_ksk,
216 rank_in | rank_out,
217 rank_out,
218 ));
219 let mut scratch_apply: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::keyswitch_inplace_scratch_space(
220 module, n, basek, k_ct, k_ksk, digits, rank_out,
221 ));
222
223 let var_xs: f64 = 0.5;
224
225 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
226 sk0.fill_ternary_prob(var_xs, &mut source_xs);
227
228 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
229 sk1.fill_ternary_prob(var_xs, &mut source_xs);
230
231 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
232 sk2.fill_ternary_prob(var_xs, &mut source_xs);
233 let sk2_prepared: GLWESecretPrepared<Vec<u8>, B> = sk2.prepare_alloc(module, scratch_apply.borrow());
234
235 // gglwe_{s1}(s0) = s0 -> s1
236 ct_gglwe_s0s1.encrypt_sk(
237 module,
238 &sk0,
239 &sk1,
240 &mut source_xa,
241 &mut source_xe,
242 sigma,
243 scratch_enc.borrow(),
244 );
245
246 // gglwe_{s2}(s1) -> s1 -> s2
247 ct_gglwe_s1s2.encrypt_sk(
248 module,
249 &sk1,
250 &sk2,
251 &mut source_xa,
252 &mut source_xe,
253 sigma,
254 scratch_enc.borrow(),
255 );
256
257 let ct_gglwe_s1s2_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> =
258 ct_gglwe_s1s2.prepare_alloc(module, scratch_apply.borrow());
259
260 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
261 ct_gglwe_s0s1.keyswitch_inplace(module, &ct_gglwe_s1s2_prepared, scratch_apply.borrow());
262
263 let ct_gglwe_s0s2: GGLWESwitchingKey<Vec<u8>> = ct_gglwe_s0s1;
264
265 let max_noise: f64 = log2_std_noise_gglwe_product(
266 n as f64,
267 basek * digits,
268 var_xs,
269 var_xs,
270 0f64,
271 sigma * sigma,
272 0f64,
273 rank_out as f64,
274 k_ct,
275 k_ksk,
276 );
277
278 ct_gglwe_s0s2
279 .key
280 .assert_noise(module, &sk2_prepared, &sk0.data, max_noise + 0.5);
281}