poulpy/poulpy-core/src/tests/generics/keyswitch/gglwe_ct.rs
poulpy-core/src/tests/generics/keyswitch/gglwe_ct.rsRaw
1use poulpy_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 source::Source,
15};
16
17use crate::{
18 encryption::SIGMA,
19 layouts::{
20 GGLWESwitchingKey, GLWESecret,
21 prepared::{GGLWESwitchingKeyPrepared, GLWESecretPrepared, PrepareAlloc},
22 },
23 noise::log2_std_noise_gglwe_product,
24};
25
26#[allow(clippy::too_many_arguments)]
27pub fn test_gglwe_switching_key_keyswitch<B>(
28 module: &Module<B>,
29 basek: usize,
30 k_out: usize,
31 k_in: usize,
32 k_ksk: usize,
33 digits: usize,
34 rank_in_s0s1: usize,
35 rank_out_s0s1: usize,
36 rank_out_s1s2: usize,
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 scratch_enc.borrow(),
136 );
137
138 // gglwe_{s2}(s1) -> s1 -> s2
139 ct_gglwe_s1s2.encrypt_sk(
140 module,
141 &sk1,
142 &sk2,
143 &mut source_xa,
144 &mut source_xe,
145 scratch_enc.borrow(),
146 );
147
148 let ct_gglwe_s1s2_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> =
149 ct_gglwe_s1s2.prepare_alloc(module, scratch_apply.borrow());
150
151 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
152 ct_gglwe_s0s2.keyswitch(
153 module,
154 &ct_gglwe_s0s1,
155 &ct_gglwe_s1s2_prepared,
156 scratch_apply.borrow(),
157 );
158
159 let max_noise: f64 = log2_std_noise_gglwe_product(
160 n as f64,
161 basek * digits,
162 0.5,
163 0.5,
164 0f64,
165 SIGMA * SIGMA,
166 0f64,
167 rank_out_s0s1 as f64,
168 k_in,
169 k_ksk,
170 );
171
172 ct_gglwe_s0s2
173 .key
174 .assert_noise(module, &sk2_prepared, &sk0.data, max_noise + 0.5);
175}
176
177#[allow(clippy::too_many_arguments)]
178pub fn test_gglwe_switching_key_keyswitch_inplace<B>(
179 module: &Module<B>,
180 basek: usize,
181 k_ct: usize,
182 k_ksk: usize,
183 digits: usize,
184 rank_in: usize,
185 rank_out: usize,
186) where
187 Module<B>: VecZnxDftAllocBytes
188 + VecZnxBigNormalize<B>
189 + VecZnxDftFromVecZnx<B>
190 + SvpApplyInplace<B>
191 + VecZnxDftToVecZnxBigConsume<B>
192 + VecZnxFillUniform
193 + VecZnxSubABInplace
194 + VecZnxAddInplace
195 + VecZnxNormalizeInplace<B>
196 + VecZnxAddNormal
197 + VecZnxNormalize<B>
198 + VecZnxSub
199 + SvpPrepare<B>
200 + SvpPPolAllocBytes
201 + SvpPPolAlloc<B>
202 + VecZnxBigAllocBytes
203 + VecZnxBigAddInplace<B>
204 + VecZnxBigAddSmallInplace<B>
205 + VecZnxNormalizeTmpBytes
206 + VecZnxAddScalarInplace
207 + VmpPMatAlloc<B>
208 + VmpPrepare<B>
209 + VmpApplyTmpBytes
210 + VmpApply<B>
211 + VmpApplyAdd<B>
212 + VecZnxBigNormalizeTmpBytes
213 + VecZnxSwithcDegree
214 + VecZnxSubScalarInplace,
215 B: Backend
216 + TakeVecZnxDftImpl<B>
217 + TakeVecZnxBigImpl<B>
218 + TakeSvpPPolImpl<B>
219 + ScratchOwnedAllocImpl<B>
220 + ScratchOwnedBorrowImpl<B>
221 + ScratchAvailableImpl<B>
222 + TakeScalarZnxImpl<B>
223 + TakeVecZnxImpl<B>,
224{
225 let n: usize = module.n();
226 let rows: usize = k_ct.div_ceil(basek * digits);
227 let digits_in: usize = 1;
228
229 let mut ct_gglwe_s0s1: GGLWESwitchingKey<Vec<u8>> =
230 GGLWESwitchingKey::alloc(n, basek, k_ct, rows, digits_in, rank_in, rank_out);
231 let mut ct_gglwe_s1s2: GGLWESwitchingKey<Vec<u8>> =
232 GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_out, rank_out);
233
234 let mut source_xs: Source = Source::new([0u8; 32]);
235 let mut source_xe: Source = Source::new([0u8; 32]);
236 let mut source_xa: Source = Source::new([0u8; 32]);
237
238 let mut scratch_enc: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::encrypt_sk_scratch_space(
239 module,
240 n,
241 basek,
242 k_ksk,
243 rank_in | rank_out,
244 rank_out,
245 ));
246 let mut scratch_apply: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::keyswitch_inplace_scratch_space(
247 module, n, basek, k_ct, k_ksk, digits, rank_out,
248 ));
249
250 let var_xs: f64 = 0.5;
251
252 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
253 sk0.fill_ternary_prob(var_xs, &mut source_xs);
254
255 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
256 sk1.fill_ternary_prob(var_xs, &mut source_xs);
257
258 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
259 sk2.fill_ternary_prob(var_xs, &mut source_xs);
260 let sk2_prepared: GLWESecretPrepared<Vec<u8>, B> = sk2.prepare_alloc(module, scratch_apply.borrow());
261
262 // gglwe_{s1}(s0) = s0 -> s1
263 ct_gglwe_s0s1.encrypt_sk(
264 module,
265 &sk0,
266 &sk1,
267 &mut source_xa,
268 &mut source_xe,
269 scratch_enc.borrow(),
270 );
271
272 // gglwe_{s2}(s1) -> s1 -> s2
273 ct_gglwe_s1s2.encrypt_sk(
274 module,
275 &sk1,
276 &sk2,
277 &mut source_xa,
278 &mut source_xe,
279 scratch_enc.borrow(),
280 );
281
282 let ct_gglwe_s1s2_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> =
283 ct_gglwe_s1s2.prepare_alloc(module, scratch_apply.borrow());
284
285 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
286 ct_gglwe_s0s1.keyswitch_inplace(module, &ct_gglwe_s1s2_prepared, scratch_apply.borrow());
287
288 let ct_gglwe_s0s2: GGLWESwitchingKey<Vec<u8>> = ct_gglwe_s0s1;
289
290 let max_noise: f64 = log2_std_noise_gglwe_product(
291 n as f64,
292 basek * digits,
293 var_xs,
294 var_xs,
295 0f64,
296 SIGMA * SIGMA,
297 0f64,
298 rank_out as f64,
299 k_ct,
300 k_ksk,
301 );
302
303 ct_gglwe_s0s2
304 .key
305 .assert_noise(module, &sk2_prepared, &sk0.data, max_noise + 0.5);
306}