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