1use backend::{FFT64, Module, ScalarZnxDftOps, ScratchOwned, Stats, VecZnxOps};
 2use sampling::source::Source;
 3
 4use crate::{FourierGLWECiphertext, FourierGLWESecret, GLWEPlaintext, GLWESecret, GLWETensorKey, GetRow, Infos};
 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: GLWETensorKey<Vec<u8>, FFT64> = GLWETensorKey::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(GLWETensorKey::encrypt_sk_scratch_space(
27        &module,
28        basek,
29        tensor_key.k(),
30        rank,
31    ));
32
33    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, rank);
34    sk.fill_ternary_prob(0.5, &mut source_xs);
35    let sk_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::from(&module, &sk);
36
37    tensor_key.encrypt_sk(
38        &module,
39        &sk_dft,
40        &mut source_xa,
41        &mut source_xe,
42        sigma,
43        scratch.borrow(),
44    );
45
46    let mut ct_glwe_fourier: FourierGLWECiphertext<Vec<u8>, FFT64> = FourierGLWECiphertext::alloc(&module, basek, k, rank);
47    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
48
49    let mut sk_ij_dft: FourierGLWESecret<Vec<u8>, FFT64> = FourierGLWESecret::alloc(&module, 1);
50    let mut sk_ij: GLWESecret<Vec<u8>> = GLWESecret::alloc(&module, 1);
51
52    (0..rank).for_each(|i| {
53        (0..rank).for_each(|j| {
54            module.svp_apply(&mut sk_ij_dft.data, 0, &sk_dft.data, i, &sk_dft.data, j);
55            module.scalar_znx_idft(&mut sk_ij.data, 0, &sk_ij_dft.data, 0, scratch.borrow());
56            (0..tensor_key.rank_in()).for_each(|col_i| {
57                (0..tensor_key.rows()).for_each(|row_i| {
58                    tensor_key
59                        .at(i, j)
60                        .get_row(&module, row_i, col_i, &mut ct_glwe_fourier);
61                    ct_glwe_fourier.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
62                    module.vec_znx_sub_scalar_inplace(&mut pt.data, 0, row_i, &sk_ij.data, col_i);
63                    let std_pt: f64 = pt.data.std(0, basek) * (k as f64).exp2();
64                    assert!((sigma - std_pt).abs() <= 0.5, "{} {}", sigma, std_pt);
65                });
66            });
67        })
68    })
69}