1#[cfg(test)]
 2mod tests {
 3    use base2k::{
 4        FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDftOps, ScratchOwned, Stats, VecZnxBig, VecZnxBigAlloc, VecZnxBigOps,
 5        VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, ZnxZero,
 6    };
 7    use sampling::source::Source;
 8
 9    use crate::{
10        elem::Infos,
11        elem_rgsw::RGSWCt,
12        elem_rlwe::{RLWECtDft, RLWEPt},
13        keys::{SecretKey, SecretKeyDft},
14    };
15
16    #[test]
17    fn encrypt_rgsw_sk() {
18        let module: Module<FFT64> = Module::<FFT64>::new(2048);
19        let log_base2k: usize = 8;
20        let log_k_ct: usize = 54;
21        let rows: usize = 4;
22
23        let sigma: f64 = 3.2;
24        let bound: f64 = sigma * 6.0;
25
26        let mut ct: RGSWCt<Vec<u8>, FFT64> = RGSWCt::new(&module, log_base2k, log_k_ct, rows);
27        let mut pt_have: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
28        let mut pt_want: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_ct);
29        let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
30
31        let mut source_xs: Source = Source::new([0u8; 32]);
32        let mut source_xe: Source = Source::new([0u8; 32]);
33        let mut source_xa: Source = Source::new([0u8; 32]);
34
35        pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
36
37        let mut scratch: ScratchOwned = ScratchOwned::new(
38            RGSWCt::encrypt_sk_scratch_space(&module, ct.size()) | RLWECtDft::decrypt_scratch_space(&module, ct.size()),
39        );
40
41        let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
42        sk.fill_ternary_prob(0.5, &mut source_xs);
43
44        let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
45        sk_dft.dft(&module, &sk);
46
47        ct.encrypt_sk(
48            &module,
49            &pt_scalar,
50            &sk_dft,
51            &mut source_xa,
52            &mut source_xe,
53            sigma,
54            bound,
55            scratch.borrow(),
56        );
57
58        let mut ct_rlwe_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_ct);
59        let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
60        let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
61
62        (0..ct.cols()).for_each(|col_j| {
63            (0..ct.rows()).for_each(|row_i| {
64                module.vec_znx_add_scalar_inplace(&mut pt_want, 0, row_i, &pt_scalar, 0);
65
66                if col_j == 1 {
67                    module.vec_znx_dft(&mut pt_dft, 0, &pt_want, 0);
68                    module.svp_apply_inplace(&mut pt_dft, 0, &sk_dft, 0);
69                    module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
70                    module.vec_znx_big_normalize(log_base2k, &mut pt_want, 0, &pt_big, 0, scratch.borrow());
71                }
72
73                ct.get_row(&module, row_i, col_j, &mut ct_rlwe_dft);
74
75                ct_rlwe_dft.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
76
77                module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
78
79                let std_pt: f64 = pt_have.data.std(0, log_base2k) * (log_k_ct as f64).exp2();
80                assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
81
82                pt_want.data.zero();
83            });
84        });
85
86        module.free();
87    }
88}