1use backend::{FFT64, Module, ScalarZnxDftOps, ScratchOwned, Stats, VecZnxOps};
 2use sampling::source::Source;
 3
 4use crate::{GLWECiphertextFourier, GLWEPlaintext, GLWESecret, GetRow, Infos, TensorKey};
 5
 6#[test]
 7fn encrypt_sk() {
 8    (1..4).for_each(|rank| {
 9        println!("test encrypt_sk rank: {}", rank);
10        test_encrypt_sk(12, 16, 54, 3.2, rank);
11    });
12}
13
14fn test_encrypt_sk(log_n: usize, basek: usize, k: usize, sigma: f64, rank: usize) {
15    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
16
17    let rows: usize = (k + basek - 1) / basek;
18
19    let mut tensor_key: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k, rows, rank);
20
21    let mut source_xs: Source = Source::new([0u8; 32]);
22    let mut source_xe: Source = Source::new([0u8; 32]);
23    let mut source_xa: Source = Source::new([0u8; 32]);
24
25    let mut scratch: ScratchOwned = ScratchOwned::new(TensorKey::generate_from_sk_scratch_space(
26        &module,
27        basek,
28        tensor_key.k(),
29        rank,
30    ));
31
32    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
33    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
34
35    tensor_key.generate_from_sk(
36        &module,
37        &sk,
38        &mut source_xa,
39        &mut source_xe,
40        sigma,
41        scratch.borrow(),
42    );
43
44    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k, rank);
45    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
46
47    let mut sk_ij = GLWESecret::alloc(&module, 1);
48
49    (0..rank).for_each(|i| {
50        (0..rank).for_each(|j| {
51            module.svp_apply(
52                &mut sk_ij.data_fourier,
53                0,
54                &sk.data_fourier,
55                i,
56                &sk.data_fourier,
57                j,
58            );
59            module.svp_idft(&mut sk_ij.data, 0, &sk_ij.data_fourier, 0, scratch.borrow());
60            (0..tensor_key.rank_in()).for_each(|col_i| {
61                (0..tensor_key.rows()).for_each(|row_i| {
62                    tensor_key
63                        .at(i, j)
64                        .get_row(&module, row_i, col_i, &mut ct_glwe_fourier);
65                    ct_glwe_fourier.decrypt(&module, &mut pt, &sk, scratch.borrow());
66                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
67                    let std_pt: f64 = pt.data.std(0, basek) * (k as f64).exp2();
68                    assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
69                });
70            });
71        })
72    })
73}