poulpy/core/src/tests/generics/keyswitch/ggsw_ct.rs
core/src/tests/generics/keyswitch/ggsw_ct.rsRaw
1use backend::hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, SvpApply, SvpApplyInplace, SvpPPolAlloc, SvpPPolAllocBytes, SvpPrepare,
4 VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxBigAddInplace, VecZnxBigAddSmallInplace, VecZnxBigAlloc,
5 VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftAddInplace, VecZnxDftAlloc,
6 VecZnxDftAllocBytes, VecZnxDftCopy, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxDftToVecZnxBigTmpA,
7 VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace,
8 VecZnxSwithcDegree, VmpApply, VmpApplyAdd, VmpApplyTmpBytes, VmpPMatAlloc, VmpPrepare,
9 },
10 layouts::{Backend, Module, ScalarZnx, ScratchOwned},
11 oep::{
12 ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
13 TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
14 },
15};
16use sampling::source::Source;
17
18use crate::{
19 layouts::{
20 GGLWESwitchingKey, GGLWETensorKey, GGSWCiphertext, GLWESecret,
21 prepared::{GGLWESwitchingKeyPrepared, GGLWETensorKeyPrepared, GLWESecretPrepared, PrepareAlloc},
22 },
23 noise::noise_ggsw_keyswitch,
24};
25
26#[allow(clippy::too_many_arguments)]
27pub fn test_ggsw_keyswitch<B>(
28 module: &Module<B>,
29 basek: usize,
30 k_out: usize,
31 k_in: usize,
32 k_ksk: usize,
33 k_tsk: usize,
34 digits: usize,
35 rank: 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 + SvpApply<B>
66 + VecZnxDftToVecZnxBigTmpA<B>
67 + VecZnxDftCopy<B>
68 + VecZnxDftAddInplace<B>
69 + VecZnxBigAlloc<B>
70 + VecZnxDftAlloc<B>,
71 B: Backend
72 + TakeVecZnxDftImpl<B>
73 + TakeVecZnxBigImpl<B>
74 + TakeSvpPPolImpl<B>
75 + ScratchOwnedAllocImpl<B>
76 + ScratchOwnedBorrowImpl<B>
77 + ScratchAvailableImpl<B>
78 + TakeScalarZnxImpl<B>
79 + TakeVecZnxImpl<B>,
80{
81 let n: usize = module.n();
82 let rows: usize = k_in.div_ceil(digits * basek);
83
84 let digits_in: usize = 1;
85
86 let mut ct_in: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_in, rows, digits_in, rank);
87 let mut ct_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_out, rows, digits_in, rank);
88 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_ksk, rows, digits, rank);
89 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
90 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
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: ScratchOwned<B> = ScratchOwned::alloc(
97 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_in, rank)
98 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank, rank)
99 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
100 | GGSWCiphertext::keyswitch_scratch_space(
101 module, n, basek, k_out, k_in, k_ksk, digits, k_tsk, digits, rank,
102 ),
103 );
104
105 let var_xs: f64 = 0.5;
106
107 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
108 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
109 let sk_in_dft: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
110
111 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
112 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
113 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
114
115 ksk.encrypt_sk(
116 module,
117 &sk_in,
118 &sk_out,
119 &mut source_xa,
120 &mut source_xe,
121 sigma,
122 scratch.borrow(),
123 );
124 tsk.encrypt_sk(
125 module,
126 &sk_out,
127 &mut source_xa,
128 &mut source_xe,
129 sigma,
130 scratch.borrow(),
131 );
132
133 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
134
135 ct_in.encrypt_sk(
136 module,
137 &pt_scalar,
138 &sk_in_dft,
139 &mut source_xa,
140 &mut source_xe,
141 sigma,
142 scratch.borrow(),
143 );
144
145 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
146 let tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> = tsk.prepare_alloc(module, scratch.borrow());
147
148 ct_out.keyswitch(
149 module,
150 &ct_in,
151 &ksk_prepared,
152 &tsk_prepared,
153 scratch.borrow(),
154 );
155
156 let max_noise = |col_j: usize| -> f64 {
157 noise_ggsw_keyswitch(
158 n as f64,
159 basek * digits,
160 col_j,
161 var_xs,
162 0f64,
163 sigma * sigma,
164 0f64,
165 rank as f64,
166 k_in,
167 k_ksk,
168 k_tsk,
169 ) + 0.5
170 };
171
172 ct_out.assert_noise(module, &sk_out_prepared, &pt_scalar, max_noise);
173}
174
175#[allow(clippy::too_many_arguments)]
176pub fn test_ggsw_keyswitch_inplace<B>(
177 module: &Module<B>,
178 basek: usize,
179 k_ct: usize,
180 k_ksk: usize,
181 k_tsk: usize,
182 digits: usize,
183 rank: usize,
184 sigma: f64,
185) where
186 Module<B>: VecZnxDftAllocBytes
187 + VecZnxBigNormalize<B>
188 + VecZnxDftFromVecZnx<B>
189 + SvpApplyInplace<B>
190 + VecZnxDftToVecZnxBigConsume<B>
191 + VecZnxFillUniform
192 + VecZnxSubABInplace
193 + VecZnxAddInplace
194 + VecZnxNormalizeInplace<B>
195 + VecZnxAddNormal
196 + VecZnxNormalize<B>
197 + VecZnxSub
198 + SvpPrepare<B>
199 + SvpPPolAllocBytes
200 + SvpPPolAlloc<B>
201 + VecZnxBigAllocBytes
202 + VecZnxBigAddInplace<B>
203 + VecZnxBigAddSmallInplace<B>
204 + VecZnxNormalizeTmpBytes
205 + VecZnxAddScalarInplace
206 + VmpPMatAlloc<B>
207 + VmpPrepare<B>
208 + VmpApplyTmpBytes
209 + VmpApply<B>
210 + VmpApplyAdd<B>
211 + VecZnxBigNormalizeTmpBytes
212 + VecZnxSwithcDegree
213 + SvpApply<B>
214 + VecZnxDftToVecZnxBigTmpA<B>
215 + VecZnxDftCopy<B>
216 + VecZnxDftAddInplace<B>
217 + VecZnxBigAlloc<B>
218 + VecZnxDftAlloc<B>,
219 B: Backend
220 + TakeVecZnxDftImpl<B>
221 + TakeVecZnxBigImpl<B>
222 + TakeSvpPPolImpl<B>
223 + ScratchOwnedAllocImpl<B>
224 + ScratchOwnedBorrowImpl<B>
225 + ScratchAvailableImpl<B>
226 + TakeScalarZnxImpl<B>
227 + TakeVecZnxImpl<B>,
228{
229 let n: usize = module.n();
230 let rows: usize = k_ct.div_ceil(digits * basek);
231
232 let digits_in: usize = 1;
233
234 let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ct, rows, digits_in, rank);
235 let mut tsk: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
236 let mut ksk: GGLWESwitchingKey<Vec<u8>> = GGLWESwitchingKey::alloc(n, basek, k_ksk, rows, digits, rank, rank);
237 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
238
239 let mut source_xs: Source = Source::new([0u8; 32]);
240 let mut source_xe: Source = Source::new([0u8; 32]);
241 let mut source_xa: Source = Source::new([0u8; 32]);
242
243 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
244 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_ct, rank)
245 | GGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank, rank)
246 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
247 | GGSWCiphertext::keyswitch_inplace_scratch_space(module, n, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
248 );
249
250 let var_xs: f64 = 0.5;
251
252 let mut sk_in: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
253 sk_in.fill_ternary_prob(var_xs, &mut source_xs);
254 let sk_in_dft: GLWESecretPrepared<Vec<u8>, B> = sk_in.prepare_alloc(module, scratch.borrow());
255
256 let mut sk_out: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
257 sk_out.fill_ternary_prob(var_xs, &mut source_xs);
258 let sk_out_prepared: GLWESecretPrepared<Vec<u8>, B> = sk_out.prepare_alloc(module, scratch.borrow());
259
260 ksk.encrypt_sk(
261 module,
262 &sk_in,
263 &sk_out,
264 &mut source_xa,
265 &mut source_xe,
266 sigma,
267 scratch.borrow(),
268 );
269 tsk.encrypt_sk(
270 module,
271 &sk_out,
272 &mut source_xa,
273 &mut source_xe,
274 sigma,
275 scratch.borrow(),
276 );
277
278 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
279
280 ct.encrypt_sk(
281 module,
282 &pt_scalar,
283 &sk_in_dft,
284 &mut source_xa,
285 &mut source_xe,
286 sigma,
287 scratch.borrow(),
288 );
289
290 let ksk_prepared: GGLWESwitchingKeyPrepared<Vec<u8>, B> = ksk.prepare_alloc(module, scratch.borrow());
291 let tsk_prepared: GGLWETensorKeyPrepared<Vec<u8>, B> = tsk.prepare_alloc(module, scratch.borrow());
292
293 ct.keyswitch_inplace(module, &ksk_prepared, &tsk_prepared, scratch.borrow());
294
295 let max_noise = |col_j: usize| -> f64 {
296 noise_ggsw_keyswitch(
297 n as f64,
298 basek * digits,
299 col_j,
300 var_xs,
301 0f64,
302 sigma * sigma,
303 0f64,
304 rank as f64,
305 k_ct,
306 k_ksk,
307 k_tsk,
308 ) + 0.5
309 };
310
311 ct.assert_noise(module, &sk_out_prepared, &pt_scalar, max_noise);
312}