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
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>: VecZnxDftAllocBytes
38 + VecZnxBigNormalize<B>
39 + VecZnxDftFromVecZnx<B>
40 + SvpApplyInplace<B>
41 + VecZnxDftToVecZnxBigConsume<B>
42 + VecZnxFillUniform
43 + VecZnxSubABInplace
44 + VecZnxAddInplace
45 + VecZnxNormalizeInplace<B>
46 + VecZnxAddNormal
47 + VecZnxNormalize<B>
48 + VecZnxSub
49 + SvpPrepare<B>
50 + SvpPPolAllocBytes
51 + SvpPPolAlloc<B>
52 + VecZnxBigAllocBytes
53 + VecZnxBigAddInplace<B>
54 + VecZnxBigAddSmallInplace<B>
55 + VecZnxNormalizeTmpBytes
56 + VecZnxAddScalarInplace
57 + VmpPMatAlloc<B>
58 + VmpPrepare<B>
59 + VmpApplyTmpBytes
60 + VmpApply<B>
61 + VmpApplyAdd<B>
62 + VecZnxBigNormalizeTmpBytes
63 + VecZnxSwithcDegree
64 + VecZnxSubScalarInplace,
65 B: TakeVecZnxDftImpl<B>
66 + TakeVecZnxBigImpl<B>
67 + TakeSvpPPolImpl<B>
68 + ScratchOwnedAllocImpl<B>
69 + ScratchOwnedBorrowImpl<B>
70 + ScratchAvailableImpl<B>
71 + TakeScalarZnxImpl<B>
72 + TakeVecZnxImpl<B>,
73{
74 let n: usize = module.n();
75 let rows: usize = k_in.div_ceil(basek * digits);
76 let digits_in: usize = 1;
77
78 let mut ct_gglwe_s0s1: GGLWESwitchingKey<Vec<u8>> =
79 GGLWESwitchingKey::alloc(n, basek, k_in, rows, digits_in, rank_in_s0s1, rank_out_s0s1);
80 let mut ct_gglwe_s1s2: GGLWESwitchingKey<Vec<u8>> =
81 GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_out_s0s1, rank_out_s1s2);
82 let mut ct_gglwe_s0s2: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(
83 n,
84 basek,
85 k_out,
86 rows,
87 digits_in,
88 rank_in_s0s1,
89 rank_out_s1s2,
90 );
91
92 let mut source_xs: Source = Source::new([0u8; 32]);
93 let mut source_xe: Source = Source::new([0u8; 32]);
94 let mut source_xa: Source = Source::new([0u8; 32]);
95
96 let mut scratch_enc: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::encrypt_sk_scratch_space(
97 module,
98 n,
99 basek,
100 k_ksk,
101 rank_in_s0s1 | rank_out_s0s1,
102 rank_out_s0s1 | rank_out_s1s2,
103 ));
104 let mut scratch_apply: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::keyswitch_scratch_space(
105 module,
106 n,
107 basek,
108 k_out,
109 k_in,
110 k_ksk,
111 digits,
112 ct_gglwe_s1s2.rank_in(),
113 ct_gglwe_s1s2.rank_out(),
114 ));
115
116 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in_s0s1);
117 sk0.fill_ternary_prob(0.5, &mut source_xs);
118
119 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out_s0s1);
120 sk1.fill_ternary_prob(0.5, &mut source_xs);
121
122 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out_s1s2);
123 sk2.fill_ternary_prob(0.5, &mut source_xs);
124 let sk2_prepared: GLWESecretPrepared<Vec<u8>, B> = sk2.prepare_alloc(module, scratch_apply.borrow());
125
126 // gglwe_{s1}(s0) = s0 -> s1
127 ct_gglwe_s0s1.encrypt_sk(
128 module,
129 &sk0,
130 &sk1,
131 &mut source_xa,
132 &mut source_xe,
133 sigma,
134 scratch_enc.borrow(),
135 );
136
137 // gglwe_{s2}(s1) -> s1 -> s2
138 ct_gglwe_s1s2.encrypt_sk(
139 module,
140 &sk1,
141 &sk2,
142 &mut source_xa,
143 &mut source_xe,
144 sigma,
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
177pub fn test_gglwe_switching_key_keyswitch_inplace<B: Backend>(
178 module: &Module<B>,
179 basek: usize,
180 k_ct: usize,
181 k_ksk: usize,
182 digits: usize,
183 rank_in: usize,
184 rank_out: usize,
185 sigma: f64,
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: TakeVecZnxDftImpl<B>
216 + TakeVecZnxBigImpl<B>
217 + TakeSvpPPolImpl<B>
218 + ScratchOwnedAllocImpl<B>
219 + ScratchOwnedBorrowImpl<B>
220 + ScratchAvailableImpl<B>
221 + TakeScalarZnxImpl<B>
222 + TakeVecZnxImpl<B>,
223{
224 let n: usize = module.n();
225 let rows: usize = k_ct.div_ceil(basek * digits);
226 let digits_in: usize = 1;
227
228 let mut ct_gglwe_s0s1: GGLWESwitchingKey<Vec<u8>> =
229 GGLWESwitchingKey::alloc(n, basek, k_ct, rows, digits_in, rank_in, rank_out);
230 let mut ct_gglwe_s1s2: GGLWESwitchingKey<Vec<u8>> =
231 GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank_out, rank_out);
232
233 let mut source_xs: Source = Source::new([0u8; 32]);
234 let mut source_xe: Source = Source::new([0u8; 32]);
235 let mut source_xa: Source = Source::new([0u8; 32]);
236
237 let mut scratch_enc: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::encrypt_sk_scratch_space(
238 module,
239 n,
240 basek,
241 k_ksk,
242 rank_in | rank_out,
243 rank_out,
244 ));
245 let mut scratch_apply: ScratchOwned<B> = ScratchOwned::alloc(GGLWESwitchingKey::keyswitch_inplace_scratch_space(
246 module, n, basek, k_ct, k_ksk, digits, rank_out,
247 ));
248
249 let var_xs: f64 = 0.5;
250
251 let mut sk0: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_in);
252 sk0.fill_ternary_prob(var_xs, &mut source_xs);
253
254 let mut sk1: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
255 sk1.fill_ternary_prob(var_xs, &mut source_xs);
256
257 let mut sk2: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank_out);
258 sk2.fill_ternary_prob(var_xs, &mut source_xs);
259 let sk2_prepared: GLWESecretPrepared<Vec<u8>, B> = sk2.prepare_alloc(module, scratch_apply.borrow());
260
261 // gglwe_{s1}(s0) = s0 -> s1
262 ct_gglwe_s0s1.encrypt_sk(
263 module,
264 &sk0,
265 &sk1,
266 &mut source_xa,
267 &mut source_xe,
268 sigma,
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 sigma,
280 scratch_enc.borrow(),
281 );
282
283 let ct_gglwe_s1s2_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> =
284 ct_gglwe_s1s2.prepare_alloc(module, scratch_apply.borrow());
285
286 // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
287 ct_gglwe_s0s1.keyswitch_inplace(module, &ct_gglwe_s1s2_prepared, scratch_apply.borrow());
288
289 let ct_gglwe_s0s2: GGLWESwitchingKey<Vec<u8>> = ct_gglwe_s0s1;
290
291 let max_noise: f64 = log2_std_noise_gglwe_product(
292 n as f64,
293 basek * digits,
294 var_xs,
295 var_xs,
296 0f64,
297 sigma * sigma,
298 0f64,
299 rank_out as f64,
300 k_ct,
301 k_ksk,
302 );
303
304 ct_gglwe_s0s2
305 .key
306 .assert_noise(module, &sk2_prepared, &sk0.data, max_noise + 0.5);
307}