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