arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/core/src/glwe_ciphertext.rs
core/src/glwe_ciphertext.rsRaw
  1use base2k::{
  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,
  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 new<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 a0_big: usize = module.bytes_of_vec_znx_big(1, in_size) + module.vec_znx_idft_tmp_bytes();
145
146        let norm: usize = module.vec_znx_big_normalize_tmp_bytes();
147
148        res_dft + (vmp | a0_big | norm)
149    }
150
151    pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, out_rank: usize, ksk_size: usize) -> usize {
152        GLWECiphertext::keyswitch_scratch_space(module, out_size, out_rank, out_size, out_rank, ksk_size)
153    }
154
155    pub fn automorphism_scratch_space(
156        module: &Module<FFT64>,
157        out_size: usize,
158        out_rank: usize,
159        in_size: usize,
160        autokey_size: usize,
161    ) -> usize {
162        GLWECiphertext::keyswitch_scratch_space(module, out_size, out_rank, in_size, out_rank, autokey_size)
163    }
164
165    pub fn automorphism_inplace_scratch_space(
166        module: &Module<FFT64>,
167        out_size: usize,
168        out_rank: usize,
169        autokey_size: usize,
170    ) -> usize {
171        GLWECiphertext::keyswitch_scratch_space(module, out_size, out_rank, out_size, out_rank, autokey_size)
172    }
173
174    pub fn external_product_scratch_space(
175        module: &Module<FFT64>,
176        out_size: usize,
177        out_rank: usize,
178        in_size: usize,
179        ggsw_size: usize,
180    ) -> usize {
181        let res_dft: usize = module.bytes_of_vec_znx_dft(out_rank + 1, ggsw_size);
182        let vmp: usize = module.bytes_of_vec_znx_dft(out_rank + 1, in_size)
183            + module.vmp_apply_tmp_bytes(
184                out_size,
185                in_size,
186                in_size,      // rows
187                out_rank + 1, // cols in
188                out_rank + 1, // cols out
189                ggsw_size,
190            );
191        let normalize: usize = module.vec_znx_big_normalize_tmp_bytes();
192
193        res_dft + (vmp | normalize)
194    }
195
196    pub fn external_product_inplace_scratch_space(
197        module: &Module<FFT64>,
198        out_size: usize,
199        out_rank: usize,
200        ggsw_size: usize,
201    ) -> usize {
202        GLWECiphertext::external_product_scratch_space(module, out_size, out_rank, out_size, ggsw_size)
203    }
204}
205
206impl<DataSelf> GLWECiphertext<DataSelf>
207where
208    VecZnx<DataSelf>: VecZnxToMut + VecZnxToRef,
209{
210    pub fn encrypt_sk<DataPt, DataSk>(
211        &mut self,
212        module: &Module<FFT64>,
213        pt: &GLWEPlaintext<DataPt>,
214        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
215        source_xa: &mut Source,
216        source_xe: &mut Source,
217        sigma: f64,
218        scratch: &mut Scratch,
219    ) where
220        VecZnx<DataPt>: VecZnxToRef,
221        ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
222    {
223        self.encrypt_sk_private(
224            module,
225            Some((pt, 0)),
226            sk_dft,
227            source_xa,
228            source_xe,
229            sigma,
230            scratch,
231        );
232    }
233
234    pub fn encrypt_zero_sk<DataSk>(
235        &mut self,
236        module: &Module<FFT64>,
237        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
238        source_xa: &mut Source,
239        source_xe: &mut Source,
240        sigma: f64,
241        scratch: &mut Scratch,
242    ) where
243        ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
244    {
245        self.encrypt_sk_private(module, None, sk_dft, source_xa, source_xe, sigma, scratch);
246    }
247
248    pub fn encrypt_pk<DataPt, DataPk>(
249        &mut self,
250        module: &Module<FFT64>,
251        pt: &GLWEPlaintext<DataPt>,
252        pk: &GLWEPublicKey<DataPk, FFT64>,
253        source_xu: &mut Source,
254        source_xe: &mut Source,
255        sigma: f64,
256        scratch: &mut Scratch,
257    ) where
258        VecZnx<DataPt>: VecZnxToRef,
259        VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
260    {
261        self.encrypt_pk_private(
262            module,
263            Some((pt, 0)),
264            pk,
265            source_xu,
266            source_xe,
267            sigma,
268            scratch,
269        );
270    }
271
272    pub fn encrypt_zero_pk<DataPk>(
273        &mut self,
274        module: &Module<FFT64>,
275        pk: &GLWEPublicKey<DataPk, FFT64>,
276        source_xu: &mut Source,
277        source_xe: &mut Source,
278        sigma: f64,
279        scratch: &mut Scratch,
280    ) where
281        VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
282    {
283        self.encrypt_pk_private(module, None, pk, source_xu, source_xe, sigma, scratch);
284    }
285
286    pub fn automorphism<DataLhs, DataRhs>(
287        &mut self,
288        module: &Module<FFT64>,
289        lhs: &GLWECiphertext<DataLhs>,
290        rhs: &AutomorphismKey<DataRhs, FFT64>,
291        scratch: &mut Scratch,
292    ) where
293        VecZnx<DataLhs>: VecZnxToRef,
294        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
295    {
296        self.keyswitch(module, lhs, &rhs.key, scratch);
297        (0..self.rank() + 1).for_each(|i| {
298            module.vec_znx_automorphism_inplace(rhs.p(), self, i);
299        })
300    }
301
302    pub fn automorphism_inplace<DataRhs>(
303        &mut self,
304        module: &Module<FFT64>,
305        rhs: &AutomorphismKey<DataRhs, FFT64>,
306        scratch: &mut Scratch,
307    ) where
308        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
309    {
310        self.keyswitch_inplace(module, &rhs.key, scratch);
311        (0..self.rank() + 1).for_each(|i| {
312            module.vec_znx_automorphism_inplace(rhs.p(), self, i);
313        })
314    }
315
316    pub(crate) fn keyswitch_from_fourier<DataLhs, DataRhs>(
317        &mut self,
318        module: &Module<FFT64>,
319        lhs: &GLWECiphertextFourier<DataLhs, FFT64>,
320        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
321        scratch: &mut Scratch,
322    ) where
323        VecZnxDft<DataLhs, FFT64>: VecZnxDftToRef<FFT64>,
324        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
325    {
326        let basek: usize = self.basek();
327
328        #[cfg(debug_assertions)]
329        {
330            assert_eq!(lhs.rank(), rhs.rank_in());
331            assert_eq!(self.rank(), rhs.rank_out());
332            assert_eq!(self.basek(), basek);
333            assert_eq!(lhs.basek(), basek);
334            assert_eq!(rhs.n(), module.n());
335            assert_eq!(self.n(), module.n());
336            assert_eq!(lhs.n(), module.n());
337            assert!(
338                scratch.available()
339                    >= GLWECiphertext::keyswitch_from_fourier_scratch_space(
340                        module,
341                        self.size(),
342                        self.rank(),
343                        lhs.size(),
344                        lhs.rank(),
345                        rhs.size(),
346                    )
347            );
348        }
349
350        let cols_in: usize = rhs.rank_in();
351        let cols_out: usize = rhs.rank_out() + 1;
352
353        // Buffer of the result of VMP in DFT
354        let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
355
356        {
357            // Applies VMP
358            let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, lhs.size());
359            (0..cols_in).for_each(|col_i| {
360                module.vec_znx_dft_copy(&mut ai_dft, col_i, lhs, col_i + 1);
361            });
362            module.vmp_apply(&mut res_dft, &ai_dft, rhs, scratch2);
363        }
364
365        // Switches result of VMP outside of DFT
366        let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume::<&mut [u8]>(res_dft);
367
368        {
369            // Switches lhs 0-th outside of DFT domain and adds on
370            let (mut a0_big, scratch2) = scratch1.tmp_vec_znx_big(module, 1, lhs.size());
371            module.vec_znx_idft(&mut a0_big, 0, lhs, 0, scratch2);
372            module.vec_znx_big_add_inplace(&mut res_big, 0, &a0_big, 0);
373        }
374
375        (0..cols_out).for_each(|i| {
376            module.vec_znx_big_normalize(basek, self, i, &res_big, i, scratch1);
377        });
378    }
379
380    pub fn keyswitch<DataLhs, DataRhs>(
381        &mut self,
382        module: &Module<FFT64>,
383        lhs: &GLWECiphertext<DataLhs>,
384        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
385        scratch: &mut Scratch,
386    ) where
387        VecZnx<DataLhs>: VecZnxToRef,
388        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
389    {
390        let basek: usize = self.basek();
391
392        #[cfg(debug_assertions)]
393        {
394            assert_eq!(lhs.rank(), rhs.rank_in());
395            assert_eq!(self.rank(), rhs.rank_out());
396            assert_eq!(self.basek(), basek);
397            assert_eq!(lhs.basek(), basek);
398            assert_eq!(rhs.n(), module.n());
399            assert_eq!(self.n(), module.n());
400            assert_eq!(lhs.n(), module.n());
401            assert!(
402                scratch.available()
403                    >= GLWECiphertext::keyswitch_scratch_space(
404                        module,
405                        self.size(),
406                        self.rank(),
407                        lhs.size(),
408                        lhs.rank(),
409                        rhs.size(),
410                    )
411            );
412        }
413
414        let cols_in: usize = rhs.rank_in();
415        let cols_out: usize = rhs.rank_out() + 1;
416
417        let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols_out, rhs.size()); // Todo optimise
418
419        {
420            let (mut ai_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols_in, lhs.size());
421            (0..cols_in).for_each(|col_i| {
422                module.vec_znx_dft(&mut ai_dft, col_i, lhs, col_i + 1);
423            });
424            module.vmp_apply(&mut res_dft, &ai_dft, rhs, scratch2);
425        }
426
427        let mut res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
428
429        module.vec_znx_big_add_small_inplace(&mut res_big, 0, lhs, 0);
430
431        (0..cols_out).for_each(|i| {
432            module.vec_znx_big_normalize(basek, self, i, &res_big, i, scratch1);
433        });
434    }
435
436    pub fn keyswitch_inplace<DataRhs>(
437        &mut self,
438        module: &Module<FFT64>,
439        rhs: &GLWESwitchingKey<DataRhs, FFT64>,
440        scratch: &mut Scratch,
441    ) where
442        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
443    {
444        unsafe {
445            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
446            self.keyswitch(&module, &*self_ptr, rhs, scratch);
447        }
448    }
449
450    pub fn external_product<DataLhs, DataRhs>(
451        &mut self,
452        module: &Module<FFT64>,
453        lhs: &GLWECiphertext<DataLhs>,
454        rhs: &GGSWCiphertext<DataRhs, FFT64>,
455        scratch: &mut Scratch,
456    ) where
457        VecZnx<DataLhs>: VecZnxToRef,
458        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
459    {
460        let basek: usize = self.basek();
461
462        #[cfg(debug_assertions)]
463        {
464            assert_eq!(rhs.rank(), lhs.rank());
465            assert_eq!(rhs.rank(), self.rank());
466            assert_eq!(self.basek(), basek);
467            assert_eq!(lhs.basek(), basek);
468            assert_eq!(rhs.n(), module.n());
469            assert_eq!(self.n(), module.n());
470            assert_eq!(lhs.n(), module.n());
471        }
472
473        let cols: usize = rhs.rank() + 1;
474
475        let (mut res_dft, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size()); // Todo optimise
476
477        {
478            let (mut a_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, lhs.size());
479            (0..cols).for_each(|col_i| {
480                module.vec_znx_dft(&mut a_dft, col_i, lhs, col_i);
481            });
482            module.vmp_apply(&mut res_dft, &a_dft, rhs, scratch2);
483        }
484
485        let res_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(res_dft);
486
487        (0..cols).for_each(|i| {
488            module.vec_znx_big_normalize(basek, self, i, &res_big, i, scratch1);
489        });
490    }
491
492    pub fn external_product_inplace<DataRhs>(
493        &mut self,
494        module: &Module<FFT64>,
495        rhs: &GGSWCiphertext<DataRhs, FFT64>,
496        scratch: &mut Scratch,
497    ) where
498        MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
499    {
500        unsafe {
501            let self_ptr: *mut GLWECiphertext<DataSelf> = self as *mut GLWECiphertext<DataSelf>;
502            self.external_product(&module, &*self_ptr, rhs, scratch);
503        }
504    }
505
506    pub(crate) fn encrypt_sk_private<DataPt, DataSk>(
507        &mut self,
508        module: &Module<FFT64>,
509        pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
510        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
511        source_xa: &mut Source,
512        source_xe: &mut Source,
513        sigma: f64,
514        scratch: &mut Scratch,
515    ) where
516        VecZnx<DataPt>: VecZnxToRef,
517        ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
518    {
519        #[cfg(debug_assertions)]
520        {
521            assert_eq!(self.rank(), sk_dft.rank());
522            assert_eq!(sk_dft.n(), module.n());
523            assert_eq!(self.n(), module.n());
524            if let Some((pt, col)) = pt {
525                assert_eq!(pt.n(), module.n());
526                assert!(col < self.rank() + 1);
527            }
528        }
529
530        let log_base2k: usize = self.basek();
531        let log_k: usize = self.k();
532        let size: usize = self.size();
533        let cols: usize = self.rank() + 1;
534
535        let (mut c0_big, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
536        c0_big.zero();
537
538        {
539            // c[i] = uniform
540            // c[0] -= c[i] * s[i],
541            (1..cols).for_each(|i| {
542                let (mut ci_dft, scratch_2) = scratch_1.tmp_vec_znx_dft(module, 1, size);
543
544                // c[i] = uniform
545                self.data.fill_uniform(log_base2k, i, size, source_xa);
546
547                // c[i] = norm(IDFT(DFT(c[i]) * DFT(s[i])))
548                module.vec_znx_dft(&mut ci_dft, 0, self, i);
549                module.svp_apply_inplace(&mut ci_dft, 0, sk_dft, i - 1);
550                let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
551
552                // use c[0] as buffer, which is overwritten later by the normalization step
553                module.vec_znx_big_normalize(log_base2k, self, 0, &ci_big, 0, scratch_2);
554
555                // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
556                module.vec_znx_sub_ab_inplace(&mut c0_big, 0, self, 0);
557
558                // c[i] += m if col = i
559                if let Some((pt, col)) = pt {
560                    if i == col {
561                        module.vec_znx_add_inplace(self, i, pt, 0);
562                        module.vec_znx_normalize_inplace(log_base2k, self, i, scratch_2);
563                    }
564                }
565            });
566        }
567
568        // c[0] += e
569        c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, sigma * SIX_SIGMA);
570
571        // c[0] += m if col = 0
572        if let Some((pt, col)) = pt {
573            if col == 0 {
574                module.vec_znx_add_inplace(&mut c0_big, 0, pt, 0);
575            }
576        }
577
578        // c[0] = norm(c[0])
579        module.vec_znx_normalize(log_base2k, self, 0, &c0_big, 0, scratch_1);
580    }
581
582    pub(crate) fn encrypt_pk_private<DataPt, DataPk>(
583        &mut self,
584        module: &Module<FFT64>,
585        pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
586        pk: &GLWEPublicKey<DataPk, FFT64>,
587        source_xu: &mut Source,
588        source_xe: &mut Source,
589        sigma: f64,
590        scratch: &mut Scratch,
591    ) where
592        VecZnx<DataPt>: VecZnxToRef,
593        VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
594    {
595        #[cfg(debug_assertions)]
596        {
597            assert_eq!(self.basek(), pk.basek());
598            assert_eq!(self.n(), module.n());
599            assert_eq!(pk.n(), module.n());
600            assert_eq!(self.rank(), pk.rank());
601            if let Some((pt, _)) = pt {
602                assert_eq!(pt.basek(), pk.basek());
603                assert_eq!(pt.n(), module.n());
604            }
605        }
606
607        let log_base2k: usize = pk.basek();
608        let size_pk: usize = pk.size();
609        let cols: usize = self.rank() + 1;
610
611        // Generates u according to the underlying secret distribution.
612        let (mut u_dft, scratch_1) = scratch.tmp_scalar_znx_dft(module, 1);
613
614        {
615            let (mut u, _) = scratch_1.tmp_scalar_znx(module, 1);
616            match pk.dist {
617                SecretDistribution::NONE => panic!(
618                    "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through \
619                     Self::generate"
620                ),
621                SecretDistribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
622                SecretDistribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
623                SecretDistribution::ZERO => {}
624            }
625
626            module.svp_prepare(&mut u_dft, 0, &u, 0);
627        }
628
629        // ct[i] = pk[i] * u + ei (+ m if col = i)
630        (0..cols).for_each(|i| {
631            let (mut ci_dft, scratch_2) = scratch_1.tmp_vec_znx_dft(module, 1, size_pk);
632            // ci_dft = DFT(u) * DFT(pk[i])
633            module.svp_apply(&mut ci_dft, 0, &u_dft, 0, pk, i);
634
635            // ci_big = u * p[i]
636            let mut ci_big = module.vec_znx_idft_consume(ci_dft);
637
638            // ci_big = u * pk[i] + e
639            ci_big.add_normal(log_base2k, 0, pk.k(), source_xe, sigma, sigma * SIX_SIGMA);
640
641            // ci_big = u * pk[i] + e + m (if col = i)
642            if let Some((pt, col)) = pt {
643                if col == i {
644                    module.vec_znx_big_add_small_inplace(&mut ci_big, 0, pt, 0);
645                }
646            }
647
648            // ct[i] = norm(ci_big)
649            module.vec_znx_big_normalize(log_base2k, self, i, &ci_big, 0, scratch_2);
650        });
651    }
652}
653
654impl<DataSelf> GLWECiphertext<DataSelf>
655where
656    VecZnx<DataSelf>: VecZnxToRef,
657{
658    pub fn decrypt<DataPt, DataSk>(
659        &self,
660        module: &Module<FFT64>,
661        pt: &mut GLWEPlaintext<DataPt>,
662        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
663        scratch: &mut Scratch,
664    ) where
665        VecZnx<DataPt>: VecZnxToMut,
666        ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
667    {
668        #[cfg(debug_assertions)]
669        {
670            assert_eq!(self.rank(), sk_dft.rank());
671            assert_eq!(self.n(), module.n());
672            assert_eq!(pt.n(), module.n());
673            assert_eq!(sk_dft.n(), module.n());
674        }
675
676        let cols: usize = self.rank() + 1;
677
678        let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
679        c0_big.zero();
680
681        {
682            (1..cols).for_each(|i| {
683                // ci_dft = DFT(a[i]) * DFT(s[i])
684                let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
685                module.vec_znx_dft(&mut ci_dft, 0, self, i);
686                module.svp_apply_inplace(&mut ci_dft, 0, sk_dft, i - 1);
687                let ci_big = module.vec_znx_idft_consume(ci_dft);
688
689                // c0_big += a[i] * s[i]
690                module.vec_znx_big_add_inplace(&mut c0_big, 0, &ci_big, 0);
691            });
692        }
693
694        // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
695        module.vec_znx_big_add_small_inplace(&mut c0_big, 0, self, 0);
696
697        // pt = norm(BIG(m + e))
698        module.vec_znx_big_normalize(self.basek(), pt, 0, &mut c0_big, 0, scratch_1);
699
700        pt.basek = self.basek();
701        pt.k = pt.k().min(self.k());
702    }
703}