1use base2k::{
 2    AddNormal, Backend, FillUniform, Module, VecZnxDftOps, ScalarZnxDftOps, ScalarZnxDftToRef, VecZnxBigOps, Scratch, VecZnx, VecZnxToMut, VecZnxToRef, ZnxInfos, FFT64
 3};
 4
 5use sampling::source::Source;
 6
 7use crate::{
 8    elem::{CipherVecZnx, Plaintext},
 9    keys::SecretKey,
10};
11
12pub trait EncryptSk<B: Backend, D> {
13    fn encrypt<P, S>(
14        module: &Module<B>,
15        res: &mut D,
16        pt: Option<&Plaintext<P>>,
17        sk: &SecretKey<S>,
18        source_xa: &mut Source,
19        source_xe: &mut Source,
20        scratch: &mut Scratch,
21        sigma: f64,
22        bound: f64,
23    ) where
24        P: VecZnxToRef,
25        S: ScalarZnxDftToRef<B>;
26}
27
28impl<C> EncryptSk<FFT64, CipherVecZnx<C>> for CipherVecZnx<C>
29where
30    VecZnx<C>: VecZnxToMut + VecZnxToRef,
31{
32    fn encrypt<P, S>(
33        module: &Module<FFT64>,
34        ct: &mut CipherVecZnx<C>,
35        pt: Option<&Plaintext<P>>,
36        sk: &SecretKey<S>,
37        source_xa: &mut Source,
38        source_xe: &mut Source,
39        scratch: &mut Scratch,
40        sigma: f64,
41        bound: f64,
42    ) where
43        P: VecZnxToRef,
44        S: ScalarZnxDftToRef<FFT64>,
45    {
46        let log_base2k: usize = ct.log_base2k();
47        let log_q: usize = ct.log_q();
48        let mut ct_mut: VecZnx<&mut [u8]> = ct.data_mut().to_mut();
49        let size: usize = ct_mut.size();
50
51        // c1 = a
52        ct_mut.fill_uniform(log_base2k, 1, size, source_xa);
53
54        let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
55
56        {
57            let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
58
59            // c0_dft = DFT(a) * DFT(s)
60            module.svp_apply(&mut c0_dft, 0, &sk.data().to_ref(), 0, &ct_mut, 1);
61
62            // c0_big = IDFT(c0_dft)
63            module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
64        }
65
66        // c0_big = m - c0_big
67        if let Some(pt) = pt {
68            module.vec_znx_big_sub_small_b_inplace(&mut c0_big, 0, &pt.data().to_ref(), 0);
69        }
70        // c0_big += e
71        c0_big.add_normal(log_base2k, 0, log_q, source_xe, sigma, bound);
72
73        // c0 = norm(c0_big = -as + m + e)
74        module.vec_znx_big_normalize(log_base2k, &mut ct_mut, 0, &c0_big, 0, scratch_1);
75    }
76}