poulpy/core/src/tests/generics/automorphism/ggsw_ct.rs
core/src/tests/generics/automorphism/ggsw_ct.rsRaw
1use backend::hal::{
2 api::{
3 ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxCopy,
4 VecZnxStd, VecZnxSubABInplace, VecZnxSwithcDegree, 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 GGLWEAutomorphismKey, GGLWETensorKey, GGSWCiphertext, GLWESecret,
17 prepared::{GGLWEAutomorphismKeyExec, GLWESecretExec, GGLWETensorKeyExec},
18 },
19 noise::noise_ggsw_keyswitch,
20 trait_families::GGSWAssertNoiseFamily,
21};
22
23use crate::trait_families::{
24 GGLWESwitchingKeyEncryptSkFamily, GGLWETensorKeyEncryptSkFamily, GGSWKeySwitchFamily, GLWESecretExecModuleFamily,
25};
26
27pub fn test_ggsw_automorphism<B: Backend>(
28 p: i64,
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>: GGSWAssertNoiseFamily<B>
40 + GLWESecretExecModuleFamily<B>
41 + VecZnxAddScalarInplace
42 + VecZnxCopy
43 + VecZnxStd
44 + VecZnxSubABInplace
45 + VmpPMatAlloc<B>
46 + VmpPMatPrepare<B>
47 + GGSWKeySwitchFamily<B>
48 + GGLWESwitchingKeyEncryptSkFamily<B>
49 + GGLWETensorKeyEncryptSkFamily<B>
50 + VecZnxSwithcDegree
51 + VecZnxAutomorphismInplace
52 + VecZnxAutomorphismInplace
53 + VecZnxAutomorphism,
54 B: TakeVecZnxDftImpl<B>
55 + TakeVecZnxBigImpl<B>
56 + TakeSvpPPolImpl<B>
57 + ScratchOwnedAllocImpl<B>
58 + ScratchOwnedBorrowImpl<B>
59 + ScratchAvailableImpl<B>
60 + TakeScalarZnxImpl<B>
61 + TakeVecZnxImpl<B>
62 + VecZnxDftAllocBytesImpl<B>
63 + VecZnxBigAllocBytesImpl<B>
64 + TakeSvpPPolImpl<B>,
65{
66 let n: usize = module.n();
67 let rows: usize = k_in.div_ceil(basek * digits);
68 let rows_in: usize = k_in.div_euclid(basek * digits);
69
70 let digits_in: usize = 1;
71
72 let mut ct_in: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_in, rows_in, digits_in, rank);
73 let mut ct_out: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_out, rows_in, digits_in, rank);
74 let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
75 let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, 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 | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
85 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
86 | GGSWCiphertext::automorphism_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: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
94 sk.fill_ternary_prob(var_xs, &mut source_xs);
95 let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
96
97 auto_key.encrypt_sk(
98 module,
99 p,
100 &sk,
101 &mut source_xa,
102 &mut source_xe,
103 sigma,
104 scratch.borrow(),
105 );
106 tensor_key.encrypt_sk(
107 module,
108 &sk,
109 &mut source_xa,
110 &mut source_xe,
111 sigma,
112 scratch.borrow(),
113 );
114
115 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
116
117 ct_in.encrypt_sk(
118 module,
119 &pt_scalar,
120 &sk_exec,
121 &mut source_xa,
122 &mut source_xe,
123 sigma,
124 scratch.borrow(),
125 );
126
127 let mut auto_key_exec: GGLWEAutomorphismKeyExec<Vec<u8>, B> =
128 GGLWEAutomorphismKeyExec::alloc(module, n, basek, k_ksk, rows, digits, rank);
129 auto_key_exec.prepare(module, &auto_key, scratch.borrow());
130
131 let mut tsk_exec: GGLWETensorKeyExec<Vec<u8>, B> = GGLWETensorKeyExec::alloc(module, n, basek, k_tsk, rows, digits, rank);
132 tsk_exec.prepare(module, &tensor_key, scratch.borrow());
133
134 ct_out.automorphism(module, &ct_in, &auto_key_exec, &tsk_exec, scratch.borrow());
135
136 module.vec_znx_automorphism_inplace(p, &mut pt_scalar.as_vec_znx_mut(), 0);
137
138 let max_noise = |col_j: usize| -> f64 {
139 noise_ggsw_keyswitch(
140 n as f64,
141 basek * digits,
142 col_j,
143 var_xs,
144 0f64,
145 sigma * sigma,
146 0f64,
147 rank as f64,
148 k_in,
149 k_ksk,
150 k_tsk,
151 ) + 0.5
152 };
153
154 ct_out.assert_noise(module, &sk_exec, &pt_scalar, &max_noise);
155}
156
157pub fn test_ggsw_automorphism_inplace<B: Backend>(
158 p: i64,
159 module: &Module<B>,
160 basek: usize,
161 k_ct: usize,
162 k_ksk: usize,
163 k_tsk: usize,
164 digits: usize,
165 rank: usize,
166 sigma: f64,
167) where
168 Module<B>: GGSWAssertNoiseFamily<B>
169 + GLWESecretExecModuleFamily<B>
170 + VecZnxAddScalarInplace
171 + VecZnxCopy
172 + VecZnxStd
173 + VecZnxSubABInplace
174 + VmpPMatAlloc<B>
175 + VmpPMatPrepare<B>
176 + GGSWKeySwitchFamily<B>
177 + GGLWESwitchingKeyEncryptSkFamily<B>
178 + GGLWETensorKeyEncryptSkFamily<B>
179 + VecZnxSwithcDegree
180 + VecZnxAutomorphismInplace
181 + VecZnxAutomorphismInplace
182 + VecZnxAutomorphism,
183 B: TakeVecZnxDftImpl<B>
184 + TakeVecZnxBigImpl<B>
185 + TakeSvpPPolImpl<B>
186 + ScratchOwnedAllocImpl<B>
187 + ScratchOwnedBorrowImpl<B>
188 + ScratchAvailableImpl<B>
189 + TakeScalarZnxImpl<B>
190 + TakeVecZnxImpl<B>
191 + VecZnxDftAllocBytesImpl<B>
192 + VecZnxBigAllocBytesImpl<B>
193 + TakeSvpPPolImpl<B>,
194{
195 let n: usize = module.n();
196 let rows: usize = k_ct.div_ceil(digits * basek);
197 let rows_in: usize = k_ct.div_euclid(basek * digits);
198 let digits_in: usize = 1;
199
200 let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(n, basek, k_ct, rows_in, digits_in, rank);
201 let mut tensor_key: GGLWETensorKey<Vec<u8>> = GGLWETensorKey::alloc(n, basek, k_tsk, rows, digits, rank);
202 let mut auto_key: GGLWEAutomorphismKey<Vec<u8>> = GGLWEAutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
203 let mut pt_scalar: ScalarZnx<Vec<u8>> = ScalarZnx::alloc(n, 1);
204
205 let mut source_xs: Source = Source::new([0u8; 32]);
206 let mut source_xe: Source = Source::new([0u8; 32]);
207 let mut source_xa: Source = Source::new([0u8; 32]);
208
209 let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
210 GGSWCiphertext::encrypt_sk_scratch_space(module, n, basek, k_ct, rank)
211 | GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
212 | GGLWETensorKey::encrypt_sk_scratch_space(module, n, basek, k_tsk, rank)
213 | GGSWCiphertext::automorphism_inplace_scratch_space(module, n, basek, k_ct, k_ksk, digits, k_tsk, digits, rank),
214 );
215
216 let var_xs: f64 = 0.5;
217
218 let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
219 sk.fill_ternary_prob(var_xs, &mut source_xs);
220 let sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
221
222 auto_key.encrypt_sk(
223 module,
224 p,
225 &sk,
226 &mut source_xa,
227 &mut source_xe,
228 sigma,
229 scratch.borrow(),
230 );
231 tensor_key.encrypt_sk(
232 module,
233 &sk,
234 &mut source_xa,
235 &mut source_xe,
236 sigma,
237 scratch.borrow(),
238 );
239
240 pt_scalar.fill_ternary_hw(0, n, &mut source_xs);
241
242 ct.encrypt_sk(
243 module,
244 &pt_scalar,
245 &sk_exec,
246 &mut source_xa,
247 &mut source_xe,
248 sigma,
249 scratch.borrow(),
250 );
251
252 let mut auto_key_exec: GGLWEAutomorphismKeyExec<Vec<u8>, B> =
253 GGLWEAutomorphismKeyExec::alloc(module, n, basek, k_ksk, rows, digits, rank);
254 auto_key_exec.prepare(module, &auto_key, scratch.borrow());
255
256 let mut tsk_exec: GGLWETensorKeyExec<Vec<u8>, B> = GGLWETensorKeyExec::alloc(module, n, basek, k_tsk, rows, digits, rank);
257 tsk_exec.prepare(module, &tensor_key, scratch.borrow());
258
259 ct.automorphism_inplace(module, &auto_key_exec, &tsk_exec, scratch.borrow());
260
261 module.vec_znx_automorphism_inplace(p, &mut pt_scalar.as_vec_znx_mut(), 0);
262
263 let max_noise = |col_j: usize| -> f64 {
264 noise_ggsw_keyswitch(
265 n as f64,
266 basek * digits,
267 col_j,
268 var_xs,
269 0f64,
270 sigma * sigma,
271 0f64,
272 rank as f64,
273 k_ct,
274 k_ksk,
275 k_tsk,
276 ) + 0.5
277 };
278
279 ct.assert_noise(module, &sk_exec, &pt_scalar, &max_noise);
280}