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