1use base2k::{
  2    AddNormal, Backend, FFT64, FillUniform, Module, ScalarZnxDftOps, ScalarZnxDftToRef, Scratch, VecZnx, VecZnxAlloc,
  3    VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef,
  4    VecZnxToMut, VecZnxToRef, ZnxInfos,
  5};
  6
  7use sampling::source::Source;
  8
  9use crate::{
 10    elem::{CtVecZnx, CtVecZnxDft, PtVecZnx},
 11    keys::SecretKey,
 12};
 13
 14pub trait EncryptSk<B: Backend, D, P> {
 15    fn encrypt<S>(
 16        module: &Module<B>,
 17        res: &mut D,
 18        pt: Option<&P>,
 19        sk: &SecretKey<S>,
 20        source_xa: &mut Source,
 21        source_xe: &mut Source,
 22        scratch: &mut Scratch,
 23        sigma: f64,
 24        bound: f64,
 25    ) where
 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, P> EncryptSk<FFT64, CtVecZnx<C>, PtVecZnx<P>> for CtVecZnx<C>
 34where
 35    VecZnx<C>: VecZnxToMut + VecZnxToRef,
 36    VecZnx<P>: VecZnxToRef,
 37{
 38    fn encrypt<S>(
 39        module: &Module<FFT64>,
 40        ct: &mut CtVecZnx<C>,
 41        pt: Option<&PtVecZnx<P>>,
 42        sk: &SecretKey<S>,
 43        source_xa: &mut Source,
 44        source_xe: &mut Source,
 45        scratch: &mut Scratch,
 46        sigma: f64,
 47        bound: f64,
 48    ) where
 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            module.vec_znx_dft(&mut c0_dft, 0, &ct_mut, 1);
 64
 65            // c0_dft = DFT(a) * DFT(s)
 66            module.svp_apply_inplace(&mut c0_dft, 0, &sk.data().to_ref(), 0);
 67
 68            // c0_big = IDFT(c0_dft)
 69            module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
 70        }
 71
 72        // c0_big = m - c0_big
 73        if let Some(pt) = pt {
 74            module.vec_znx_big_sub_small_b_inplace(&mut c0_big, 0, &pt.data().to_ref(), 0);
 75        }
 76        // c0_big += e
 77        c0_big.add_normal(log_base2k, 0, log_q, source_xe, sigma, bound);
 78
 79        // c0 = norm(c0_big = -as + m + e)
 80        module.vec_znx_big_normalize(log_base2k, &mut ct_mut, 0, &c0_big, 0, scratch_1);
 81    }
 82}
 83
 84pub trait EncryptZeroSk<B: Backend, D> {
 85    fn encrypt_zero<S>(
 86        module: &Module<B>,
 87        res: &mut D,
 88        sk: &SecretKey<S>,
 89        source_xa: &mut Source,
 90        source_xe: &mut Source,
 91        scratch: &mut Scratch,
 92        sigma: f64,
 93        bound: f64,
 94    ) where
 95        S: ScalarZnxDftToRef<B>;
 96
 97    fn encrypt_zero_tmp_bytes(module: &Module<B>, size: usize) -> usize {
 98        (module.bytes_of_vec_znx(1, size) | module.bytes_of_vec_znx_dft(1, size))
 99            + module.bytes_of_vec_znx_big(1, size)
100            + module.bytes_of_vec_znx(1, size)
101            + module.vec_znx_big_normalize_tmp_bytes()
102    }
103}
104
105impl<C> EncryptZeroSk<FFT64, CtVecZnxDft<C, FFT64>> for CtVecZnxDft<C, FFT64>
106where
107    VecZnxDft<C, FFT64>: VecZnxDftToMut<FFT64> + VecZnxDftToRef<FFT64>,
108{
109    fn encrypt_zero<S>(
110        module: &Module<FFT64>,
111        ct: &mut CtVecZnxDft<C, FFT64>,
112        sk: &SecretKey<S>,
113        source_xa: &mut Source,
114        source_xe: &mut Source,
115        scratch: &mut Scratch,
116        sigma: f64,
117        bound: f64,
118    ) where
119        S: ScalarZnxDftToRef<FFT64>,
120    {
121        let log_base2k: usize = ct.log_base2k();
122        let log_q: usize = ct.log_q();
123        let mut ct_mut: VecZnxDft<&mut [u8], FFT64> = ct.data_mut().to_mut();
124        let size: usize = ct_mut.size();
125
126        // ct[1] = DFT(a)
127        {
128            let (mut tmp_znx, _) = scratch.tmp_vec_znx(module, 1, size);
129            tmp_znx.fill_uniform(log_base2k, 1, size, source_xa);
130            module.vec_znx_dft(&mut ct_mut, 1, &tmp_znx, 0);
131        }
132
133        let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
134
135        {
136            let (mut tmp_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
137            // c0_dft = DFT(a) * DFT(s)
138            module.svp_apply(&mut tmp_dft, 0, &sk.data().to_ref(), 0, &ct_mut, 1);
139            // c0_big = IDFT(c0_dft)
140            module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut tmp_dft, 0);
141        }
142
143        // c0_big += e
144        c0_big.add_normal(log_base2k, 0, log_q, source_xe, sigma, bound);
145
146        // c0 = norm(c0_big = -as + e)
147        let (mut tmp_znx, scratch_2) = scratch_1.tmp_vec_znx(module, 1, size);
148        module.vec_znx_big_normalize(log_base2k, &mut tmp_znx, 0, &c0_big, 0, scratch_2);
149        // ct[0] = DFT(-as + e)
150        module.vec_znx_dft(&mut ct_mut, 0, &tmp_znx, 0);
151    }
152}