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::encrypt_sk_scratch_space(
34        &module,
35        rank,
36        tensor_key.size(),
37    ));
38
39    let mut sk: SecretKey<Vec<u8>> = SecretKey::alloc(&module, rank);
40    sk.fill_ternary_prob(0.5, &mut source_xs);
41
42    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::alloc(&module, rank);
43    sk_dft.dft(&module, &sk);
44
45    tensor_key.encrypt_sk(
46        &module,
47        &sk_dft,
48        &mut source_xa,
49        &mut source_xe,
50        sigma,
51        scratch.borrow(),
52    );
53
54    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k, rank);
55    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
56
57    (0..rank).for_each(|i| {
58        (0..rank).for_each(|j| {
59            let mut sk_ij_dft: ScalarZnxDft<Vec<u8>, FFT64> = module.new_scalar_znx_dft(1);
60            module.svp_apply(&mut sk_ij_dft, 0, &sk_dft.data, i, &sk_dft.data, j);
61            let sk_ij: ScalarZnx<Vec<u8>> = module
62                .vec_znx_idft_consume(sk_ij_dft.as_vec_znx_dft())
63                .to_vec_znx_small()
64                .to_scalar_znx();
65
66            (0..tensor_key.rank_in()).for_each(|col_i| {
67                (0..tensor_key.rows()).for_each(|row_i| {
68                    tensor_key
69                        .at(i, j)
70                        .get_row(&module, row_i, col_i, &mut ct_glwe_fourier);
71                    ct_glwe_fourier.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
72                    module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk_ij, col_i);
73                    let std_pt: f64 = pt.data.std(0, basek) * (k as f64).exp2();
74                    assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
75                });
76            });
77        })
78    })
79}