arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/core/src/glwe_ciphertext.rs
core/src/glwe_ciphertext.rsRaw
  1use backend::{
  2    AddNormal, Backend, FFT64, FillUniform, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnxAlloc, ScalarZnxDftAlloc,
  3    ScalarZnxDftOps, Scratch, VecZnx, VecZnxAlloc, VecZnxBig, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDftAlloc,
  4    VecZnxDftOps, VecZnxOps, VecZnxToMut, VecZnxToRef, ZnxZero,
  5};
  6use sampling::source::Source;
  7
  8use crate::{
  9    SIX_SIGMA,
 10    automorphism::AutomorphismKey,
 11    elem::{Infos, SetMetaData},
 12    ggsw_ciphertext::GGSWCiphertext,
 13    glwe_ciphertext_fourier::GLWECiphertextFourier,
 14    glwe_ops::GLWEOps,
 15    glwe_plaintext::GLWEPlaintext,
 16    keys::{GLWEPublicKey, SecretDistribution, SecretKeyFourier},
 17    keyswitch_key::GLWESwitchingKey,
 18    utils::derive_size,
 19};
 20
 21pub struct GLWECiphertext<C> {
 22    pub data: VecZnx<C>,
 23    pub basek: usize,
 24    pub k: usize,
 25}
 26
 27impl GLWECiphertext<Vec<u8>> {
 28    pub fn alloc<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
 29        Self {
 30            data: module.new_vec_znx(rank + 1, derive_size(basek, k)),
 31            basek,
 32            k,
 33        }
 34    }
 35
 36    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rank: usize) -> usize {
 37        module.bytes_of_vec_znx(rank + 1, derive_size(basek, k))
 38    }
 39}
 40
 41impl<T> Infos for GLWECiphertext<T> {
 42    type Inner = VecZnx<T>;
 43
 44    fn inner(&self) -> &Self::Inner {
 45        &self.data
 46    }
 47
 48    fn basek(&self) -> usize {
 49        self.basek
 50    }
 51
 52    fn k(&self) -> usize {
 53        self.k
 54    }
 55}
 56
 57impl<T> GLWECiphertext<T> {
 58    pub fn rank(&self) -> usize {
 59        self.cols() - 1
 60    }
 61}
 62
 63impl<C: AsRef<[u8]>> GLWECiphertext<C> {
 64    #[allow(dead_code)]
 65    pub(crate) fn dft<R: AsMut<[u8]> + AsRef<[u8]>>(&self, module: &Module<FFT64>, res: &mut GLWECiphertextFourier<R, FFT64>) {
 66        #[cfg(debug_assertions)]
 67        {
 68            assert_eq!(self.rank(), res.rank());
 69            assert_eq!(self.basek(), res.basek())
 70        }
 71
 72        (0..self.rank() + 1).for_each(|i| {
 73            module.vec_znx_dft(&mut res.data, i, &self.data, i);
 74        })
 75    }
 76}
 77
 78impl GLWECiphertext<Vec<u8>> {
 79    pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize) -> usize {
 80        let size: usize = derive_size(basek, k);
 81        module.vec_znx_big_normalize_tmp_bytes() + module.bytes_of_vec_znx_dft(1, size) + module.bytes_of_vec_znx(1, size)
 82    }
 83    pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, basek: usize, k: usize) -> usize {
 84        let size: usize = derive_size(basek, k);
 85        ((module.bytes_of_vec_znx_dft(1, size) + module.bytes_of_vec_znx_big(1, size)) | module.bytes_of_scalar_znx(1))
 86            + module.bytes_of_scalar_znx_dft(1)
 87            + module.vec_znx_big_normalize_tmp_bytes()
 88    }
 89
 90    pub fn decrypt_scratch_space(module: &Module<FFT64>, basek: usize, k: usize) -> usize {
 91        let size: usize = derive_size(basek, k);
 92        (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, size)) + module.bytes_of_vec_znx_big(1, size)
 93    }
 94
 95    pub fn keyswitch_scratch_space(
 96        module: &Module<FFT64>,
 97        basek: usize,
 98        out_k: usize,
 99        out_rank: usize,
100        in_k: usize,
101        in_rank: usize,
102        ksk_k: usize,
103    ) -> usize {
104        let res_dft: usize = GLWECiphertextFourier::bytes_of(module, basek, out_k, out_rank);
105        let in_size: usize = derive_size(basek, in_k);
106        let out_size: usize = derive_size(basek, out_k);
107        let ksk_size: usize = derive_size(basek, ksk_k);
108        let vmp: usize = module.vmp_apply_tmp_bytes(out_size, in_size, in_size, in_rank, out_rank + 1, ksk_size)
109            + module.bytes_of_vec_znx_dft(in_rank, in_size);
110        let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
111        return res_dft + (vmp | normalize);
112    }
113
114    pub fn keyswitch_from_fourier_scratch_space(
115        module: &Module<FFT64>,
116        basek: usize,
117        out_k: usize,
118        out_rank: usize,
119        in_k: usize,
120        in_rank: usize,
121        ksk_k: usize,
122    ) -> usize {
123        Self::keyswitch_scratch_space(module, basek, out_k, out_rank, in_k, in_rank, ksk_k)
124    }
125
126    pub fn keyswitch_inplace_scratch_space(
127        module: &Module<FFT64>,
128        basek: usize,
129        out_k: usize,
130        out_rank: usize,
131        ksk_k: usize,
132    ) -> usize {
133        Self::keyswitch_scratch_space(module, basek, out_k, out_rank, out_k, out_rank, ksk_k)
134    }
135
136    pub fn automorphism_scratch_space(
137        module: &Module<FFT64>,
138        basek: usize,
139        out_k: usize,
140        in_k: usize,
141        atk_k: usize,
142        rank: usize,
143    ) -> usize {
144        Self::keyswitch_scratch_space(module, basek, out_k, rank, in_k, rank, atk_k)
145    }
146
147    pub fn automorphism_inplace_scratch_space(
148        module: &Module<FFT64>,
149        basek: usize,
150        out_k: usize,
151        atk_k: usize,
152        rank: usize,
153    ) -> usize {
154        Self::keyswitch_scratch_space(module, basek, out_k, rank, out_k, rank, atk_k)
155    }
156
157    pub fn external_product_scratch_space(
158        module: &Module<FFT64>,
159        basek: usize,
160        out_k: usize,
161        in_k: usize,
162        ggsw_k: usize,
163        rank: usize,
164    ) -> usize {
165        let res_dft: usize = GLWECiphertextFourier::bytes_of(module, basek, out_k, rank);
166        let in_size: usize = derive_size(basek, in_k);
167        let out_size: usize = derive_size(basek, out_k);
168        let ggsw_size: usize = derive_size(basek, ggsw_k);
169        let vmp: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size)
170            + module.vmp_apply_tmp_bytes(
171                out_size,
172                in_size,
173                in_size,  // rows
174                rank + 1, // cols in
175                rank + 1, // cols out
176                ggsw_size,
177            );
178        let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
179        res_dft + (vmp | normalize)
180    }
181
182    pub fn external_product_inplace_scratch_space(
183        module: &Module<FFT64>,
184        basek: usize,
185        out_k: usize,
186        ggsw_k: usize,
187        rank: usize,
188    ) -> usize {
189        Self::external_product_scratch_space(module, basek, out_k, out_k, ggsw_k, rank)
190    }
191}
192
193impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> SetMetaData for GLWECiphertext<DataSelf> {
194    fn set_k(&mut self, k: usize) {
195        self.k = k
196    }
197
198    fn set_basek(&mut self, basek: usize) {
199        self.basek = basek
200    }
201}
202
203impl<DataSelf: AsRef<[u8]> + AsMut<[u8]>> GLWECiphertext<DataSelf> {
204    pub fn encrypt_sk<DataPt: AsRef<[u8]>, DataSk: AsRef<[u8]>>(
205        &mut self,
206        module: &Module<FFT64>,
207        pt: &GLWEPlaintext<DataPt>,
208        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
209        source_xa: &mut Source,
210        source_xe: &mut Source,
211        sigma: f64,
212        scratch: &mut Scratch,
213    ) {
214        self.encrypt_sk_private(
215            module,
216            Some((pt, 0)),
217            sk_dft,
218            source_xa,
219            source_xe,
220            sigma,
221            scratch,
222        );
223    }
224
225    pub fn encrypt_zero_sk<DataSk: AsRef<[u8]>>(
226        &mut self,
227        module: &Module<FFT64>,
228        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
229        source_xa: &mut Source,
230        source_xe: &mut Source,
231        sigma: f64,
232        scratch: &mut Scratch,
233    ) {
234        self.encrypt_sk_private(
235            module,
236            None::<(&GLWEPlaintext<Vec<u8>>, usize)>,
237            sk_dft,
238            source_xa,
239            source_xe,
240            sigma,
241            scratch,
242        );
243    }
244
245    pub fn encrypt_pk<DataPt: AsRef<[u8]>, DataPk: AsRef<[u8]>>(
246        &mut self,
247        module: &Module<FFT64>,
248        pt: &GLWEPlaintext<DataPt>,
249        pk: &GLWEPublicKey<DataPk, FFT64>,
250        source_xu: &mut Source,
251        source_xe: &mut Source,
252        sigma: f64,
253        scratch: &mut Scratch,
254    ) {
255        self.encrypt_pk_private(
256            module,
257            Some((pt, 0)),
258            pk,
259            source_xu,
260            source_xe,
261            sigma,
262            scratch,
263        );
264    }
265
266    pub fn encrypt_zero_pk<DataPk: AsRef<[u8]>>(
267        &mut self,
268        module: &Module<FFT64>,
269        pk: &GLWEPublicKey<DataPk, FFT64>,
270        source_xu: &mut Source,
271        source_xe: &mut Source,
272        sigma: f64,
273        scratch: &mut Scratch,
274    ) {
275        self.encrypt_pk_private(
276            module,
277            None::<(&GLWEPlaintext<Vec<u8>>, usize)>,
278            pk,
279            source_xu,
280            source_xe,
281            sigma,
282            scratch,
283        );
284    }
285
286    pub fn automorphism<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
287        &mut self,
288        module: &Module<FFT64>,
289        lhs: &GLWECiphertext<DataLhs>,
290        rhs: &AutomorphismKey<DataRhs, FFT64>,
291        scratch: &mut Scratch,
292    ) {
293        self.keyswitch(module, lhs, &rhs.key, scratch);
294        (0..self.rank() + 1).for_each(|i| {
295            module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i);
296        })
297    }
298
299    pub fn automorphism_inplace<DataRhs: AsRef<[u8]>>(
300        &mut self,
301        module: &Module<FFT64>,
302        rhs: &AutomorphismKey<DataRhs, FFT64>,
303        scratch: &mut Scratch,
304    ) {
305        self.keyswitch_inplace(module, &rhs.key, scratch);
306        (0..self.rank() + 1).for_each(|i| {
307            module.vec_znx_automorphism_inplace(rhs.p(), &mut self.data, i);
308        })
309    }
310
311    pub fn automorphism_add<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
312        &mut self,
313        module: &Module<FFT64>,
314        lhs: &GLWECiphertext<DataLhs>,
315        rhs: &AutomorphismKey<DataRhs, FFT64>,
316        scratch: &mut Scratch,
317    ) {
318        Self::keyswitch_private::<_, _, 1>(self, rhs.p(), module, lhs, &rhs.key, scratch);
319    }
320
321    pub fn automorphism_add_inplace<DataRhs: AsRef<[u8]>>(
322        &mut self,
323        module: &Module<FFT64>,
324        rhs: &AutomorphismKey<DataRhs, FFT64>,
325        scratch: &mut Scratch,
326    ) {
327        unsafe {
328            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
329            Self::keyswitch_private::<_, _, 1>(self, rhs.p(), module, &*self_ptr, &rhs.key, scratch);
330        }
331    }
332
333    pub fn automorphism_sub_ab<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
334        &mut self,
335        module: &Module<FFT64>,
336        lhs: &GLWECiphertext<DataLhs>,
337        rhs: &AutomorphismKey<DataRhs, FFT64>,
338        scratch: &mut Scratch,
339    ) {
340        Self::keyswitch_private::<_, _, 2>(self, rhs.p(), module, lhs, &rhs.key, scratch);
341    }
342
343    pub fn automorphism_sub_ab_inplace<DataRhs: AsRef<[u8]>>(
344        &mut self,
345        module: &Module<FFT64>,
346        rhs: &AutomorphismKey<DataRhs, FFT64>,
347        scratch: &mut Scratch,
348    ) {
349        unsafe {
350            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
351            Self::keyswitch_private::<_, _, 2>(self, rhs.p(), module, &*self_ptr, &rhs.key, scratch);
352        }
353    }
354
355    pub fn automorphism_sub_ba<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
356        &mut self,
357        module: &Module<FFT64>,
358        lhs: &GLWECiphertext<DataLhs>,
359        rhs: &AutomorphismKey<DataRhs, FFT64>,
360        scratch: &mut Scratch,
361    ) {
362        Self::keyswitch_private::<_, _, 3>(self, rhs.p(), module, lhs, &rhs.key, scratch);
363    }
364
365    pub fn automorphism_sub_ba_inplace<DataRhs: AsRef<[u8]>>(
366        &mut self,
367        module: &Module<FFT64>,
368        rhs: &AutomorphismKey<DataRhs, FFT64>,
369        scratch: &mut Scratch,
370    ) {
371        unsafe {
372            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
373            Self::keyswitch_private::<_, _, 3>(self, rhs.p(), module, &*self_ptr, &rhs.key, scratch);
374        }
375    }
376
377    pub(crate) fn keyswitch_from_fourier<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
378        &mut self,
379        module: &Module<FFT64>,
380        lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
381        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
382        scratch: &mut Scratch,
383    ) {
384        let basek: usize = self.basek();
385
386        #[cfg(debug_assertions)]
387        {
388            assert_eq!(lhs.rank(), rhs.rank_in());
389            assert_eq!(self.rank(), rhs.rank_out());
390            assert_eq!(self.basek(), basek);
391            assert_eq!(lhs.basek(), basek);
392            assert_eq!(rhs.n(), module.n());
393            assert_eq!(self.n(), module.n());
394            assert_eq!(lhs.n(), module.n());
395            assert!(
396                scratch.available()
397                    >= GLWECiphertext::keyswitch_from_fourier_scratch_space(
398                        module,
399                        self.basek(),
400                        self.k(),
401                        self.rank(),
402                        lhs.k(),
403                        lhs.rank(),
404                        rhs.k(),
405                    )
406            );
407        }
408
409        let cols_in: usize = rhs.rank_in();
410        let cols_out: usize = rhs.rank_out() + 1;
411
412        // Buffer of the result of VMP in DFT
413        let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
414
415        {
416            // Applies VMP
417            let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, lhs.size());
418            (0..cols_in).for_each(|col_i| {
419                module.vec_znx_dft_copy(&mut ai_dft, col_i, &lhs.data, col_i + 1);
420            });
421            module.vmp_apply(&mut res_dft, &ai_dft, &rhs.0.data, scratch2);
422        }
423
424        module.vec_znx_dft_add_inplace(&mut res_dft, 0, &lhs.data, 0);
425
426        // Switches result of VMP outside of DFT
427        let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
428
429        (0..cols_out).for_each(|i| {
430            module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
431        });
432    }
433
434    pub fn keyswitch<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
435        &mut self,
436        module: &Module<FFT64>,
437        lhs: &GLWECiphertext<DataLhs>,
438        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
439        scratch: &mut Scratch,
440    ) {
441        Self::keyswitch_private::<_, _, 0>(self, 0, module, lhs, rhs, scratch);
442    }
443
444    pub(crate) fn keyswitch_private<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>, const OP: u8>(
445        &mut self,
446        apply_auto: i64,
447        module: &Module<FFT64>,
448        lhs: &GLWECiphertext<DataLhs>,
449        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
450        scratch: &mut Scratch,
451    ) {
452        let basek: usize = self.basek();
453
454        #[cfg(debug_assertions)]
455        {
456            assert_eq!(lhs.rank(), rhs.rank_in());
457            assert_eq!(self.rank(), rhs.rank_out());
458            assert_eq!(self.basek(), basek);
459            assert_eq!(lhs.basek(), basek);
460            assert_eq!(rhs.n(), module.n());
461            assert_eq!(self.n(), module.n());
462            assert_eq!(lhs.n(), module.n());
463            assert!(
464                scratch.available()
465                    >= GLWECiphertext::keyswitch_scratch_space(
466                        module,
467                        self.basek(),
468                        self.k(),
469                        self.rank(),
470                        lhs.k(),
471                        lhs.rank(),
472                        rhs.k(),
473                    )
474            );
475        }
476
477        let cols_in: usize = rhs.rank_in();
478        let cols_out: usize = rhs.rank_out() + 1;
479
480        let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
481
482        {
483            let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, lhs.size());
484            (0..cols_in).for_each(|col_i| {
485                module.vec_znx_dft(&mut ai_dft, col_i, &lhs.data, col_i + 1);
486            });
487            module.vmp_apply(&mut res_dft, &ai_dft, &rhs.0.data, scratch2);
488        }
489
490        let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
491
492        module.vec_znx_big_add_small_inplace(&mut res_big, 0, &lhs.data, 0);
493
494        (0..cols_out).for_each(|i| {
495            if apply_auto != 0 {
496                module.vec_znx_big_automorphism_inplace(apply_auto, &mut res_big, i);
497            }
498
499            match OP {
500                1 => module.vec_znx_big_add_small_inplace(&mut res_big, i, &lhs.data, i),
501                2 => module.vec_znx_big_sub_small_a_inplace(&mut res_big, i, &lhs.data, i),
502                3 => module.vec_znx_big_sub_small_b_inplace(&mut res_big, i, &lhs.data, i),
503                _ => {}
504            }
505            module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
506        });
507    }
508
509    pub fn keyswitch_inplace<DataRhs: AsRef<[u8]>>(
510        &mut self,
511        module: &Module<FFT64>,
512        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
513        scratch: &mut Scratch,
514    ) {
515        unsafe {
516            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
517            self.keyswitch(&module, &*self_ptr, rhs, scratch);
518        }
519    }
520
521    pub fn external_product<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
522        &mut self,
523        module: &Module<FFT64>,
524        lhs: &GLWECiphertext<DataLhs>,
525        rhs: &GGSWCiphertext<DataRhs, FFT64>,
526        scratch: &mut Scratch,
527    ) {
528        let basek: usize = self.basek();
529
530        #[cfg(debug_assertions)]
531        {
532            assert_eq!(rhs.rank(), lhs.rank());
533            assert_eq!(rhs.rank(), self.rank());
534            assert_eq!(self.basek(), basek);
535            assert_eq!(lhs.basek(), basek);
536            assert_eq!(rhs.n(), module.n());
537            assert_eq!(self.n(), module.n());
538            assert_eq!(lhs.n(), module.n());
539        }
540
541        let cols: usize = rhs.rank() + 1;
542
543        let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size()); // Todo optimise
544
545        {
546            let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, lhs.size());
547            (0..cols).for_each(|col_i| {
548                module.vec_znx_dft(&mut a_dft, col_i, &lhs.data, col_i);
549            });
550            module.vmp_apply(&mut res_dft, &a_dft, &rhs.data, scratch2);
551        }
552
553        let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
554
555        (0..cols).for_each(|i| {
556            module.vec_znx_big_normalize(basek, &mut self.data, i, &res_big, i, scratch1);
557        });
558    }
559
560    pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
561        &mut self,
562        module: &Module<FFT64>,
563        rhs: &GGSWCiphertext<DataRhs, FFT64>,
564        scratch: &mut Scratch,
565    ) {
566        unsafe {
567            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
568            self.external_product(&module, &*self_ptr, rhs, scratch);
569        }
570    }
571
572    pub(crate) fn encrypt_sk_private<DataPt: AsRef<[u8]>, DataSk: AsRef<[u8]>>(
573        &mut self,
574        module: &Module<FFT64>,
575        pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
576        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
577        source_xa: &mut Source,
578        source_xe: &mut Source,
579        sigma: f64,
580        scratch: &mut Scratch,
581    ) {
582        #[cfg(debug_assertions)]
583        {
584            assert_eq!(self.rank(), sk_dft.rank());
585            assert_eq!(sk_dft.n(), module.n());
586            assert_eq!(self.n(), module.n());
587            if let Some((pt, col)) = pt {
588                assert_eq!(pt.n(), module.n());
589                assert!(col < self.rank() + 1);
590            }
591            assert!(
592                scratch.available() >= GLWECiphertext::encrypt_sk_scratch_space(module, self.basek(), self.k()),
593                "scratch.available(): {} < GLWECiphertext::encrypt_sk_scratch_space: {}",
594                scratch.available(),
595                GLWECiphertext::encrypt_sk_scratch_space(module, self.basek(), self.k())
596            )
597        }
598
599        let basek: usize = self.basek();
600        let k: usize = self.k();
601        let size: usize = self.size();
602        let cols: usize = self.rank() + 1;
603
604        let (mut c0_big, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
605        c0_big.zero();
606
607        {
608            // c[i] = uniform
609            // c[0] -= c[i] * s[i],
610            (1..cols).for_each(|i| {
611                let (mut ci_dft, scratch_2) = scratch_1.tmp_vec_znx_dft(module, 1, size);
612
613                // c[i] = uniform
614                self.data.fill_uniform(basek, i, size, source_xa);
615
616                // c[i] = norm(IDFT(DFT(c[i]) * DFT(s[i])))
617                module.vec_znx_dft(&mut ci_dft, 0, &self.data, i);
618                module.svp_apply_inplace(&mut ci_dft, 0, &sk_dft.data, i - 1);
619                let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
620
621                // use c[0] as buffer, which is overwritten later by the normalization step
622                module.vec_znx_big_normalize(basek, &mut self.data, 0, &ci_big, 0, scratch_2);
623
624                // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
625                module.vec_znx_sub_ab_inplace(&mut c0_big, 0, &self.data, 0);
626
627                // c[i] += m if col = i
628                if let Some((pt, col)) = pt {
629                    if i == col {
630                        module.vec_znx_add_inplace(&mut self.data, i, &pt.data, 0);
631                        module.vec_znx_normalize_inplace(basek, &mut self.data, i, scratch_2);
632                    }
633                }
634            });
635        }
636
637        // c[0] += e
638        c0_big.add_normal(basek, 0, k, source_xe, sigma, sigma * SIX_SIGMA);
639
640        // c[0] += m if col = 0
641        if let Some((pt, col)) = pt {
642            if col == 0 {
643                module.vec_znx_add_inplace(&mut c0_big, 0, &pt.data, 0);
644            }
645        }
646
647        // c[0] = norm(c[0])
648        module.vec_znx_normalize(basek, &mut self.data, 0, &c0_big, 0, scratch_1);
649    }
650
651    pub(crate) fn encrypt_pk_private<DataPt: AsRef<[u8]>, DataPk: AsRef<[u8]>>(
652        &mut self,
653        module: &Module<FFT64>,
654        pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
655        pk: &GLWEPublicKey<DataPk, FFT64>,
656        source_xu: &mut Source,
657        source_xe: &mut Source,
658        sigma: f64,
659        scratch: &mut Scratch,
660    ) {
661        #[cfg(debug_assertions)]
662        {
663            assert_eq!(self.basek(), pk.basek());
664            assert_eq!(self.n(), module.n());
665            assert_eq!(pk.n(), module.n());
666            assert_eq!(self.rank(), pk.rank());
667            if let Some((pt, _)) = pt {
668                assert_eq!(pt.basek(), pk.basek());
669                assert_eq!(pt.n(), module.n());
670            }
671        }
672
673        let basek: usize = pk.basek();
674        let size_pk: usize = pk.size();
675        let cols: usize = self.rank() + 1;
676
677        // Generates u according to the underlying secret distribution.
678        let (mut u_dft, scratch_1) = scratch.tmp_scalar_znx_dft(module, 1);
679
680        {
681            let (mut u, _) = scratch_1.tmp_scalar_znx(module, 1);
682            match pk.dist {
683                SecretDistribution::NONE => panic!(
684                    "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through \
685                     Self::generate"
686                ),
687                SecretDistribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
688                SecretDistribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
689                SecretDistribution::ZERO => {}
690            }
691
692            module.svp_prepare(&mut u_dft, 0, &u, 0);
693        }
694
695        // ct[i] = pk[i] * u + ei (+ m if col = i)
696        (0..cols).for_each(|i| {
697            let (mut ci_dft, scratch_2) = scratch_1.tmp_vec_znx_dft(module, 1, size_pk);
698            // ci_dft = DFT(u) * DFT(pk[i])
699            module.svp_apply(&mut ci_dft, 0, &u_dft, 0, &pk.data.data, i);
700
701            // ci_big = u * p[i]
702            let mut ci_big = module.vec_znx_idft_consume(ci_dft);
703
704            // ci_big = u * pk[i] + e
705            ci_big.add_normal(basek, 0, pk.k(), source_xe, sigma, sigma * SIX_SIGMA);
706
707            // ci_big = u * pk[i] + e + m (if col = i)
708            if let Some((pt, col)) = pt {
709                if col == i {
710                    module.vec_znx_big_add_small_inplace(&mut ci_big, 0, &pt.data, 0);
711                }
712            }
713
714            // ct[i] = norm(ci_big)
715            module.vec_znx_big_normalize(basek, &mut self.data, i, &ci_big, 0, scratch_2);
716        });
717    }
718}
719
720impl<DataSelf: AsRef<[u8]>> GLWECiphertext<DataSelf> {
721    pub fn clone(&self) -> GLWECiphertext<Vec<u8>> {
722        GLWECiphertext {
723            data: self.data.clone(),
724            basek: self.basek(),
725            k: self.k(),
726        }
727    }
728
729    pub fn decrypt<DataPt: AsMut<[u8]> + AsRef<[u8]>, DataSk: AsRef<[u8]>>(
730        &self,
731        module: &Module<FFT64>,
732        pt: &mut GLWEPlaintext<DataPt>,
733        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
734        scratch: &mut Scratch,
735    ) {
736        #[cfg(debug_assertions)]
737        {
738            assert_eq!(self.rank(), sk_dft.rank());
739            assert_eq!(self.n(), module.n());
740            assert_eq!(pt.n(), module.n());
741            assert_eq!(sk_dft.n(), module.n());
742        }
743
744        let cols: usize = self.rank() + 1;
745
746        let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
747        c0_big.zero();
748
749        {
750            (1..cols).for_each(|i| {
751                // ci_dft = DFT(a[i]) * DFT(s[i])
752                let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
753                module.vec_znx_dft(&mut ci_dft, 0, &self.data, i);
754                module.svp_apply_inplace(&mut ci_dft, 0, &sk_dft.data, i - 1);
755                let ci_big = module.vec_znx_idft_consume(ci_dft);
756
757                // c0_big += a[i] * s[i]
758                module.vec_znx_big_add_inplace(&mut c0_big, 0, &ci_big, 0);
759            });
760        }
761
762        // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
763        module.vec_znx_big_add_small_inplace(&mut c0_big, 0, &self.data, 0);
764
765        // pt = norm(BIG(m + e))
766        module.vec_znx_big_normalize(self.basek(), &mut pt.data, 0, &mut c0_big, 0, scratch_1);
767
768        pt.basek = self.basek();
769        pt.k = pt.k().min(self.k());
770    }
771}
772
773pub trait GLWECiphertextToRef {
774    fn to_ref(&self) -> GLWECiphertext<&[u8]>;
775}
776
777impl<D: AsRef<[u8]>> GLWECiphertextToRef for GLWECiphertext<D> {
778    fn to_ref(&self) -> GLWECiphertext<&[u8]> {
779        GLWECiphertext {
780            data: self.data.to_ref(),
781            basek: self.basek,
782            k: self.k,
783        }
784    }
785}
786
787pub trait GLWECiphertextToMut {
788    fn to_mut(&mut self) -> GLWECiphertext<&mut [u8]>;
789}
790
791impl<D: AsMut<[u8]> + AsRef<[u8]>> GLWECiphertextToMut for GLWECiphertext<D> {
792    fn to_mut(&mut self) -> GLWECiphertext<&mut [u8]> {
793        GLWECiphertext {
794            data: self.data.to_mut(),
795            basek: self.basek,
796            k: self.k,
797        }
798    }
799}
800
801impl<D> GLWEOps for GLWECiphertext<D>
802where
803    D: AsRef<[u8]> + AsMut<[u8]>,
804    GLWECiphertext<D>: GLWECiphertextToMut + Infos + SetMetaData,
805{
806}