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