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