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