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