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