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