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