1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxCopy,
  4        VecZnxStd, VecZnxSubScalarInplace, VecZnxSwithcDegree,
  5    },
  6    layouts::{Backend, Module, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    AutomorphismKey, AutomorphismKeyCompressed, AutomorphismKeyEncryptSkFamily, AutomorphismKeyExec, GGLWEExecLayoutFamily,
 16    GLWEDecryptFamily, GLWEKeyswitchFamily, GLWEPlaintext, GLWESecret, GLWESecretExec, Infos,
 17    noise::log2_std_noise_gglwe_product,
 18};
 19
 20pub(crate) trait AutomorphismTestModuleFamily<B: Backend> = AutomorphismKeyEncryptSkFamily<B>
 21    + GLWEKeyswitchFamily<B>
 22    + VecZnxAutomorphism
 23    + GGLWEExecLayoutFamily<B>
 24    + VecZnxSwithcDegree
 25    + VecZnxAddScalarInplace
 26    + VecZnxAutomorphism
 27    + VecZnxAutomorphismInplace
 28    + GLWEDecryptFamily<B>
 29    + VecZnxSubScalarInplace
 30    + VecZnxStd
 31    + VecZnxCopy;
 32pub(crate) trait AutomorphismTestScratchFamily<B: Backend> = ScratchOwnedAllocImpl<B>
 33    + ScratchOwnedBorrowImpl<B>
 34    + ScratchAvailableImpl<B>
 35    + TakeScalarZnxImpl<B>
 36    + TakeVecZnxDftImpl<B>
 37    + TakeVecZnxImpl<B>
 38    + TakeSvpPPolImpl<B>
 39    + TakeVecZnxBigImpl<B>;
 40
 41pub(crate) fn test_automorphisk_key_encrypt_sk<B: Backend>(
 42    module: &Module<B>,
 43    basek: usize,
 44    k_ksk: usize,
 45    digits: usize,
 46    rank: usize,
 47    sigma: f64,
 48) where
 49    Module<B>: AutomorphismTestModuleFamily<B>,
 50    B: AutomorphismTestScratchFamily<B>,
 51{
 52    let n: usize = TakeScalarZnx;
 53    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
 54
 55    let mut atk: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
 56
 57    let mut source_xs: Source = Source::new([0u8; 32]);
 58    let mut source_xe: Source = Source::new([0u8; 32]);
 59    let mut source_xa: Source = Source::new([0u8; 32]);
 60
 61    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(AutomorphismKey::encrypt_sk_scratch_space(
 62        module, n, basek, k_ksk, rank,
 63    ));
 64
 65    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 66    sk.fill_ternary_prob(0.5, &mut source_xs);
 67
 68    let p = -5;
 69
 70    atk.encrypt_sk(
 71        module,
 72        p,
 73        &sk,
 74        &mut source_xa,
 75        &mut source_xe,
 76        sigma,
 77        scratch.borrow(),
 78    );
 79
 80    let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
 81    (0..atk.rank()).for_each(|i| {
 82        module.vec_znx_automorphism(
 83            module.galois_element_inv(p),
 84            &mut sk_out.data.as_vec_znx_mut(),
 85            i,
 86            &sk.data.as_vec_znx(),
 87            i,
 88        );
 89    });
 90    let sk_out_exec = GLWESecretExec::from(module, &sk_out);
 91
 92    atk.key
 93        .key
 94        .assert_noise(module, &sk_out_exec, &sk.data, sigma);
 95}
 96
 97pub(crate) fn test_automorphisk_key_encrypt_sk_compressed<B: Backend>(
 98    module: &Module<B>,
 99    basek: usize,
100    k_ksk: usize,
101    digits: usize,
102    rank: usize,
103    sigma: f64,
104) where
105    Module<B>: AutomorphismTestModuleFamily<B>,
106    B: AutomorphismTestScratchFamily<B>,
107{
108    let n: usize = TakeScalarZnx;
109    let rows: usize = (k_ksk - digits * basek) / (digits * basek);
110
111    let mut atk_compressed: AutomorphismKeyCompressed<Vec<u8>> =
112        AutomorphismKeyCompressed::alloc(n, basek, k_ksk, rows, digits, rank);
113
114    let mut source_xs: Source = Source::new([0u8; 32]);
115    let mut source_xe: Source = Source::new([0u8; 32]);
116
117    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(AutomorphismKey::encrypt_sk_scratch_space(
118        module, n, basek, k_ksk, rank,
119    ));
120
121    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
122    sk.fill_ternary_prob(0.5, &mut source_xs);
123
124    let p = -5;
125
126    let seed_xa: [u8; 32] = [1u8; 32];
127
128    atk_compressed.encrypt_sk(
129        module,
130        p,
131        &sk,
132        seed_xa,
133        &mut source_xe,
134        sigma,
135        scratch.borrow(),
136    );
137
138    let mut sk_out: GLWESecret<Vec<u8>> = sk.clone();
139    (0..atk_compressed.rank()).for_each(|i| {
140        module.vec_znx_automorphism(
141            module.galois_element_inv(p),
142            &mut sk_out.data.as_vec_znx_mut(),
143            i,
144            &sk.data.as_vec_znx(),
145            i,
146        );
147    });
148    let sk_out_exec = GLWESecretExec::from(module, &sk_out);
149
150    let mut atk: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_ksk, rows, digits, rank);
151    atk.decompress(module, &atk_compressed);
152
153    atk.key
154        .key
155        .assert_noise(module, &sk_out_exec, &sk.data, sigma);
156}
157
158pub(crate) fn test_gglwe_automorphism<B: Backend>(
159    module: &Module<B>,
160    p0: i64,
161    p1: i64,
162    basek: usize,
163    digits: usize,
164    k_in: usize,
165    k_out: usize,
166    k_apply: usize,
167    sigma: f64,
168    rank: usize,
169) where
170    Module<B>: AutomorphismTestModuleFamily<B>,
171    B: AutomorphismTestScratchFamily<B>,
172{
173    let n: usize = TakeScalarZnx;
174    let digits_in: usize = 1;
175
176    let rows_in: usize = k_in / (basek * digits);
177    let rows_apply: usize = k_in.div_ceil(basek * digits);
178
179    let mut auto_key_in: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
180    let mut auto_key_out: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_out, rows_in, digits_in, rank);
181    let mut auto_key_apply: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
182
183    let mut source_xs: Source = Source::new([0u8; 32]);
184    let mut source_xe: Source = Source::new([0u8; 32]);
185    let mut source_xa: Source = Source::new([0u8; 32]);
186
187    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
188        AutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
189            | AutomorphismKey::automorphism_scratch_space(module, n, basek, k_out, k_in, k_apply, digits, rank),
190    );
191
192    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(TakeScalarZnx, rank);
193    sk.fill_ternary_prob(0.5, &mut source_xs);
194
195    // gglwe_{s1}(s0) = s0 -> s1
196    auto_key_in.encrypt_sk(
197        module,
198        p0,
199        &sk,
200        &mut source_xa,
201        &mut source_xe,
202        sigma,
203        scratch.borrow(),
204    );
205
206    // gglwe_{s2}(s1) -> s1 -> s2
207    auto_key_apply.encrypt_sk(
208        module,
209        p1,
210        &sk,
211        &mut source_xa,
212        &mut source_xe,
213        sigma,
214        scratch.borrow(),
215    );
216
217    let mut auto_key_apply_exec: AutomorphismKeyExec<Vec<u8>, B> =
218        AutomorphismKeyExec::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
219
220    auto_key_apply_exec.prepare(module, &auto_key_apply, scratch.borrow());
221
222    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
223    auto_key_out.automorphism(module, &auto_key_in, &auto_key_apply_exec, scratch.borrow());
224
225    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(TakeScalarZnx, basek, k_out);
226
227    let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(TakeScalarZnx, rank);
228    sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
229    (0..rank).for_each(|i| {
230        module.vec_znx_automorphism(
231            module.galois_element_inv(p0 * p1),
232            &mut sk_auto.data.as_vec_znx_mut(),
233            i,
234            &sk.data.as_vec_znx(),
235            i,
236        );
237    });
238
239    let sk_auto_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_auto);
240
241    (0..auto_key_out.rank_in()).for_each(|col_i| {
242        (0..auto_key_out.rows()).for_each(|row_i| {
243            auto_key_out
244                .at(row_i, col_i)
245                .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
246
247            module.vec_znx_sub_scalar_inplace(
248                &mut pt.data,
249                0,
250                (digits_in - 1) + row_i * digits_in,
251                &sk.data,
252                col_i,
253            );
254
255            let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
256            let noise_want: f64 = log2_std_noise_gglwe_product(
257                TakeScalarZnx as f64,
258                basek * digits,
259                0.5,
260                0.5,
261                0f64,
262                sigma * sigma,
263                0f64,
264                rank as f64,
265                k_out,
266                k_apply,
267            );
268
269            assert!(
270                noise_have < noise_want + 0.5,
271                "{} {}",
272                noise_have,
273                noise_want
274            );
275        });
276    });
277}
278
279pub(crate) fn test_gglwe_automorphism_inplace<B: Backend>(
280    module: &Module<B>,
281    p0: i64,
282    p1: i64,
283    basek: usize,
284    digits: usize,
285    k_in: usize,
286    k_apply: usize,
287    sigma: f64,
288    rank: usize,
289) where
290    Module<B>: AutomorphismTestModuleFamily<B>,
291    B: AutomorphismTestScratchFamily<B>,
292{
293    let n: usize = TakeScalarZnx;
294    let digits_in: usize = 1;
295
296    let rows_in: usize = k_in / (basek * digits);
297    let rows_apply: usize = k_in.div_ceil(basek * digits);
298
299    let mut auto_key: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_in, rows_in, digits_in, rank);
300    let mut auto_key_apply: AutomorphismKey<Vec<u8>> = AutomorphismKey::alloc(n, basek, k_apply, rows_apply, digits, rank);
301
302    let mut source_xs: Source = Source::new([0u8; 32]);
303    let mut source_xe: Source = Source::new([0u8; 32]);
304    let mut source_xa: Source = Source::new([0u8; 32]);
305
306    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
307        AutomorphismKey::encrypt_sk_scratch_space(module, n, basek, k_apply, rank)
308            | AutomorphismKey::automorphism_inplace_scratch_space(module, n, basek, k_in, k_apply, digits, rank),
309    );
310
311    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(TakeScalarZnx, rank);
312    sk.fill_ternary_prob(0.5, &mut source_xs);
313
314    // gglwe_{s1}(s0) = s0 -> s1
315    auto_key.encrypt_sk(
316        module,
317        p0,
318        &sk,
319        &mut source_xa,
320        &mut source_xe,
321        sigma,
322        scratch.borrow(),
323    );
324
325    // gglwe_{s2}(s1) -> s1 -> s2
326    auto_key_apply.encrypt_sk(
327        module,
328        p1,
329        &sk,
330        &mut source_xa,
331        &mut source_xe,
332        sigma,
333        scratch.borrow(),
334    );
335
336    let mut auto_key_apply_exec: AutomorphismKeyExec<Vec<u8>, B> =
337        AutomorphismKeyExec::alloc(module, n, basek, k_apply, rows_apply, digits, rank);
338
339    auto_key_apply_exec.prepare(module, &auto_key_apply, scratch.borrow());
340
341    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
342    auto_key.automorphism_inplace(module, &auto_key_apply_exec, scratch.borrow());
343
344    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(TakeScalarZnx, basek, k_in);
345
346    let mut sk_auto: GLWESecret<Vec<u8>> = GLWESecret::alloc(TakeScalarZnx, rank);
347    sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
348
349    (0..rank).for_each(|i| {
350        module.vec_znx_automorphism(
351            module.galois_element_inv(p0 * p1),
352            &mut sk_auto.data.as_vec_znx_mut(),
353            i,
354            &sk.data.as_vec_znx(),
355            i,
356        );
357    });
358
359    let sk_auto_dft: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_auto);
360
361    (0..auto_key.rank_in()).for_each(|col_i| {
362        (0..auto_key.rows()).for_each(|row_i| {
363            auto_key
364                .at(row_i, col_i)
365                .decrypt(module, &mut pt, &sk_auto_dft, scratch.borrow());
366            module.vec_znx_sub_scalar_inplace(
367                &mut pt.data,
368                0,
369                (digits_in - 1) + row_i * digits_in,
370                &sk.data,
371                col_i,
372            );
373
374            let noise_have: f64 = module.vec_znx_std(basek, &pt.data, 0).log2();
375            let noise_want: f64 = log2_std_noise_gglwe_product(
376                TakeScalarZnx as f64,
377                basek * digits,
378                0.5,
379                0.5,
380                0f64,
381                sigma * sigma,
382                0f64,
383                rank as f64,
384                k_in,
385                k_apply,
386            );
387
388            assert!(
389                noise_have < noise_want + 0.5,
390                "{} {}",
391                noise_have,
392                noise_want
393            );
394        });
395    });
396}