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