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::{
 10    elem::{Ciphertext, Infos, Plaintext},
 11    keys::SecretKey,
 12};
 13
 14pub trait EncryptSk<B: Backend, C, P> {
 15    fn encrypt<S>(
 16        module: &Module<B>,
 17        res: &mut Ciphertext<C>,
 18        pt: Option<&Plaintext<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_scratch_bytes(module: &Module<B>, size: usize) -> usize;
 29}
 30
 31impl<C, P> EncryptSk<FFT64, C, P> for Ciphertext<C>
 32where
 33    C: VecZnxToMut + ZnxInfos,
 34    P: VecZnxToRef + ZnxInfos,
 35{
 36    fn encrypt<S>(
 37        module: &Module<FFT64>,
 38        ct: &mut Ciphertext<C>,
 39        pt: Option<&Plaintext<P>>,
 40        sk: &SecretKey<S>,
 41        source_xa: &mut Source,
 42        source_xe: &mut Source,
 43        scratch: &mut Scratch,
 44        sigma: f64,
 45        bound: f64,
 46    ) where
 47        S: ScalarZnxDftToRef<FFT64>,
 48    {
 49        let log_base2k: usize = ct.log_base2k();
 50        let log_q: usize = ct.log_q();
 51        let mut ct_mut: VecZnx<&mut [u8]> = ct.to_mut();
 52        let size: usize = ct_mut.size();
 53
 54        // c1 = a
 55        ct_mut.fill_uniform(log_base2k, 1, size, source_xa);
 56
 57        let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
 58
 59        {
 60            let (mut c0_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
 61            module.vec_znx_dft(&mut c0_dft, 0, &ct_mut, 1);
 62
 63            // c0_dft = DFT(a) * DFT(s)
 64            module.svp_apply_inplace(&mut c0_dft, 0, &sk.data().to_ref(), 0);
 65
 66            // c0_big = IDFT(c0_dft)
 67            module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut c0_dft, 0);
 68        }
 69
 70        // c0_big = m - c0_big
 71        if let Some(pt) = pt {
 72            module.vec_znx_big_sub_small_b_inplace(&mut c0_big, 0, pt, 0);
 73        }
 74        // c0_big += e
 75        c0_big.add_normal(log_base2k, 0, log_q, source_xe, sigma, bound);
 76
 77        // c0 = norm(c0_big = -as + m + e)
 78        module.vec_znx_big_normalize(log_base2k, &mut ct_mut, 0, &c0_big, 0, scratch_1);
 79    }
 80
 81    fn encrypt_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize {
 82        (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
 83    }
 84}
 85
 86impl<C> Ciphertext<C>
 87where
 88    C: VecZnxToMut + ZnxInfos,
 89{
 90    pub fn encrypt_sk<P, S>(
 91        &mut self,
 92        module: &Module<FFT64>,
 93        pt: Option<&Plaintext<P>>,
 94        sk: &SecretKey<S>,
 95        source_xa: &mut Source,
 96        source_xe: &mut Source,
 97        scratch: &mut Scratch,
 98        sigma: f64,
 99        bound: f64,
100    ) where
101        P: VecZnxToRef + ZnxInfos,
102        S: ScalarZnxDftToRef<FFT64>,
103    {
104        <Self as EncryptSk<FFT64, _, _>>::encrypt(
105            module, self, pt, sk, source_xa, source_xe, scratch, sigma, bound,
106        );
107    }
108
109    pub fn encrypt_sk_scratch_bytes<P>(module: &Module<FFT64>, size: usize) -> usize
110    where
111        Self: EncryptSk<FFT64, C, P>,
112    {
113        <Self as EncryptSk<FFT64, C, P>>::encrypt_scratch_bytes(module, size)
114    }
115}
116
117pub trait EncryptZeroSk<B: Backend, D> {
118    fn encrypt_zero<S>(
119        module: &Module<B>,
120        res: &mut D,
121        sk: &SecretKey<S>,
122        source_xa: &mut Source,
123        source_xe: &mut Source,
124        scratch: &mut Scratch,
125        sigma: f64,
126        bound: f64,
127    ) where
128        S: ScalarZnxDftToRef<B>;
129
130    fn encrypt_zero_scratch_bytes(module: &Module<B>, size: usize) -> usize;
131}
132
133impl<C> EncryptZeroSk<FFT64, C> for C
134where
135    C: VecZnxDftToMut<FFT64> + ZnxInfos + Infos,
136{
137    fn encrypt_zero<S>(
138        module: &Module<FFT64>,
139        ct: &mut C,
140        sk: &SecretKey<S>,
141        source_xa: &mut Source,
142        source_xe: &mut Source,
143        scratch: &mut Scratch,
144        sigma: f64,
145        bound: f64,
146    ) where
147        S: ScalarZnxDftToRef<FFT64>,
148    {
149        let log_base2k: usize = ct.log_base2k();
150        let log_q: usize = ct.log_q();
151        let mut ct_mut: VecZnxDft<&mut [u8], FFT64> = ct.to_mut();
152        let size: usize = ct_mut.size();
153
154        // ct[1] = DFT(a)
155        {
156            let (mut tmp_znx, _) = scratch.tmp_vec_znx(module, 1, size);
157            tmp_znx.fill_uniform(log_base2k, 1, size, source_xa);
158            module.vec_znx_dft(&mut ct_mut, 1, &tmp_znx, 0);
159        }
160
161        let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, size);
162
163        {
164            let (mut tmp_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, size);
165            // c0_dft = DFT(a) * DFT(s)
166            module.svp_apply(&mut tmp_dft, 0, &sk.data().to_ref(), 0, &ct_mut, 1);
167            // c0_big = IDFT(c0_dft)
168            module.vec_znx_idft_tmp_a(&mut c0_big, 0, &mut tmp_dft, 0);
169        }
170
171        // c0_big += e
172        c0_big.add_normal(log_base2k, 0, log_q, source_xe, sigma, bound);
173
174        // c0 = norm(c0_big = -as + e)
175        let (mut tmp_znx, scratch_2) = scratch_1.tmp_vec_znx(module, 1, size);
176        module.vec_znx_big_normalize(log_base2k, &mut tmp_znx, 0, &c0_big, 0, scratch_2);
177        // ct[0] = DFT(-as + e)
178        module.vec_znx_dft(&mut ct_mut, 0, &tmp_znx, 0);
179    }
180
181    fn encrypt_zero_scratch_bytes(module: &Module<FFT64>, size: usize) -> usize{
182        (module.bytes_of_vec_znx(1, size) | module.bytes_of_vec_znx_dft(1, size))
183            + module.bytes_of_vec_znx_big(1, size)
184            + module.bytes_of_vec_znx(1, size)
185            + module.vec_znx_big_normalize_tmp_bytes()
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use base2k::{FFT64, Module, ScratchOwned, VecZnx, Scalar};
192    use sampling::source::Source;
193
194    use crate::{elem::{Ciphertext, Infos, Plaintext}, keys::SecretKey};
195
196    #[test]
197    fn encrypt_sk_vec_znx_fft64() {
198        let module: Module<FFT64> = Module::<FFT64>::new(32);
199        let log_base2k: usize = 8;
200        let log_q: usize = 54;
201
202        let sigma: f64 = 3.2;
203        let bound: f64 = sigma * 6;
204
205        let mut ct: Ciphertext<VecZnx<Vec<u8>>> = Ciphertext::<VecZnx<Vec<u8>>>::new(&module, log_base2k, log_q, 2);
206        let mut pt: Plaintext<VecZnx<Vec<u8>>> = Plaintext::<VecZnx<Vec<u8>>>::new(&module, log_base2k, log_q);
207
208        let mut source_xe = Source::new([0u8; 32]);
209        let mut source_xa: Source = Source::new([0u8; 32]);
210        
211
212        let mut scratch: ScratchOwned = ScratchOwned::new(ct.encrypt_encsk_scratch_bytes(&module, ct.size()));
213
214        let mut sk: SecretKey<Scalar<Vec<u8>>> = SecretKey::new(&module);
215        let mut sk_prep
216        sk.svp_prepare(&module, &mut sk_prep);
217
218        ct.encrypt_sk(
219            &module,
220            Some(&pt),
221            &sk_prep,
222            &mut source_xa,
223            &mut source_xe,
224            scratch.borrow(),
225            sigma,
226            bound,
227        );
228    }
229}