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