arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/core/src/ggsw_ciphertext.rs
core/src/ggsw_ciphertext.rsRaw
  1use backend::{
  2    Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftScratch, Module, ScalarZnx, Scratch, VecZnxAlloc,
  3    VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, VecZnxToMut, ZnxInfos,
  4    ZnxZero,
  5};
  6use sampling::source::Source;
  7
  8use crate::{
  9    automorphism::AutomorphismKey,
 10    elem::{GetRow, Infos, SetRow},
 11    glwe_ciphertext::GLWECiphertext,
 12    glwe_ciphertext_fourier::GLWECiphertextFourier,
 13    glwe_plaintext::GLWEPlaintext,
 14    keys::SecretKeyFourier,
 15    keyswitch_key::GLWESwitchingKey,
 16    tensor_key::TensorKey,
 17    utils::derive_size,
 18};
 19
 20pub struct GGSWCiphertext<C, B: Backend> {
 21    pub data: MatZnxDft<C, B>,
 22    pub basek: usize,
 23    pub k: usize,
 24}
 25
 26impl<B: Backend> GGSWCiphertext<Vec<u8>, B> {
 27    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
 28        Self {
 29            data: module.new_mat_znx_dft(rows, rank + 1, rank + 1, derive_size(basek, k)),
 30            basek: basek,
 31            k: k,
 32        }
 33    }
 34
 35    pub fn bytes_of(module: &Module<FFT64>, basek: usize, k: usize, rows: usize, rank: usize) -> usize {
 36        module.bytes_of_mat_znx_dft(rows, rank + 1, rank + 1, derive_size(basek, k))
 37    }
 38}
 39
 40impl<T, B: Backend> Infos for GGSWCiphertext<T, B> {
 41    type Inner = MatZnxDft<T, B>;
 42
 43    fn inner(&self) -> &Self::Inner {
 44        &self.data
 45    }
 46
 47    fn basek(&self) -> usize {
 48        self.basek
 49    }
 50
 51    fn k(&self) -> usize {
 52        self.k
 53    }
 54}
 55
 56impl<T, B: Backend> GGSWCiphertext<T, B> {
 57    pub fn rank(&self) -> usize {
 58        self.data.cols_out() - 1
 59    }
 60}
 61
 62impl GGSWCiphertext<Vec<u8>, FFT64> {
 63    pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
 64        GLWECiphertext::encrypt_sk_scratch_space(module, size)
 65            + module.bytes_of_vec_znx(rank + 1, size)
 66            + module.bytes_of_vec_znx(1, size)
 67            + module.bytes_of_vec_znx_dft(rank + 1, size)
 68    }
 69
 70    pub(crate) fn expand_row_scratch_space(
 71        module: &Module<FFT64>,
 72        self_size: usize,
 73        tensor_key_size: usize,
 74        rank: usize,
 75    ) -> usize {
 76        let tmp_dft_i: usize = module.bytes_of_vec_znx_dft(rank + 1, tensor_key_size);
 77        let tmp_dft_col_data: usize = module.bytes_of_vec_znx_dft(1, self_size);
 78        let vmp: usize =
 79            tmp_dft_col_data + module.vmp_apply_tmp_bytes(self_size, self_size, self_size, rank, rank, tensor_key_size);
 80        let tmp_idft: usize = module.bytes_of_vec_znx_big(1, tensor_key_size);
 81        let norm: usize = module.vec_znx_big_normalize_tmp_bytes();
 82        tmp_dft_i + ((tmp_dft_col_data + vmp) | (tmp_idft + norm))
 83    }
 84
 85    pub(crate) fn keyswitch_internal_col0_scratch_space(
 86        module: &Module<FFT64>,
 87        out_size: usize,
 88        in_size: usize,
 89        ksk_size: usize,
 90        rank: usize,
 91    ) -> usize {
 92        GLWECiphertext::keyswitch_from_fourier_scratch_space(module, out_size, rank, in_size, rank, ksk_size)
 93            + module.bytes_of_vec_znx_dft(rank + 1, in_size)
 94    }
 95
 96    pub fn keyswitch_scratch_space(
 97        module: &Module<FFT64>,
 98        out_size: usize,
 99        in_size: usize,
100        ksk_size: usize,
101        tensor_key_size: usize,
102        rank: usize,
103    ) -> usize {
104        let res_znx: usize = module.bytes_of_vec_znx(rank + 1, out_size);
105        let ci_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
106        let ks: usize = GGSWCiphertext::keyswitch_internal_col0_scratch_space(module, out_size, in_size, ksk_size, rank);
107        let expand_rows: usize = GGSWCiphertext::expand_row_scratch_space(module, out_size, tensor_key_size, rank);
108        let res_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
109        res_znx + ci_dft + (ks | expand_rows | res_dft)
110    }
111
112    pub fn keyswitch_inplace_scratch_space(
113        module: &Module<FFT64>,
114        out_size: usize,
115        ksk_size: usize,
116        tensor_key_size: usize,
117        rank: usize,
118    ) -> usize {
119        GGSWCiphertext::keyswitch_scratch_space(module, out_size, out_size, ksk_size, tensor_key_size, rank)
120    }
121
122    pub fn automorphism_scratch_space(
123        module: &Module<FFT64>,
124        out_size: usize,
125        in_size: usize,
126        auto_key_size: usize,
127        tensor_key_size: usize,
128        rank: usize,
129    ) -> usize {
130        let cols: usize = rank + 1;
131        let res: usize = module.bytes_of_vec_znx(cols, out_size);
132        let res_dft: usize = module.bytes_of_vec_znx_dft(cols, out_size);
133        let ci_dft: usize = module.bytes_of_vec_znx_dft(cols, out_size);
134        let ks_internal: usize =
135            GGSWCiphertext::keyswitch_internal_col0_scratch_space(module, out_size, in_size, auto_key_size, rank);
136        let expand: usize = GGSWCiphertext::expand_row_scratch_space(module, out_size, tensor_key_size, rank);
137        res + ci_dft + (ks_internal | expand | res_dft)
138    }
139
140    pub fn automorphism_inplace_scratch_space(
141        module: &Module<FFT64>,
142        out_size: usize,
143        auto_key_size: usize,
144        tensor_key_size: usize,
145        rank: usize,
146    ) -> usize {
147        GGSWCiphertext::automorphism_scratch_space(
148            module,
149            out_size,
150            out_size,
151            auto_key_size,
152            tensor_key_size,
153            rank,
154        )
155    }
156
157    pub fn external_product_scratch_space(
158        module: &Module<FFT64>,
159        out_size: usize,
160        in_size: usize,
161        ggsw_size: usize,
162        rank: usize,
163    ) -> usize {
164        let tmp_in: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
165        let tmp_out: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
166        let ggsw: usize = GLWECiphertextFourier::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank);
167        tmp_in + tmp_out + ggsw
168    }
169
170    pub fn external_product_inplace_scratch_space(
171        module: &Module<FFT64>,
172        out_size: usize,
173        ggsw_size: usize,
174        rank: usize,
175    ) -> usize {
176        let tmp: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
177        let ggsw: usize = GLWECiphertextFourier::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank);
178        tmp + ggsw
179    }
180}
181
182impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> GGSWCiphertext<DataSelf, FFT64> {
183    pub fn encrypt_sk<DataPt: AsRef<[u8]>, DataSk: AsRef<[u8]>>(
184        &mut self,
185        module: &Module<FFT64>,
186        pt: &ScalarZnx<DataPt>,
187        sk_dft: &SecretKeyFourier<DataSk, FFT64>,
188        source_xa: &mut Source,
189        source_xe: &mut Source,
190        sigma: f64,
191        scratch: &mut Scratch,
192    ) {
193        #[cfg(debug_assertions)]
194        {
195            assert_eq!(self.rank(), sk_dft.rank());
196            assert_eq!(self.n(), module.n());
197            assert_eq!(pt.n(), module.n());
198            assert_eq!(sk_dft.n(), module.n());
199        }
200
201        let size: usize = self.size();
202        let basek: usize = self.basek();
203        let k: usize = self.k();
204        let cols: usize = self.rank() + 1;
205
206        let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
207        let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, cols, size);
208
209        let mut vec_znx_pt: GLWEPlaintext<&mut [u8]> = GLWEPlaintext {
210            data: tmp_znx_pt,
211            basek: basek,
212            k: k,
213        };
214
215        let mut vec_znx_ct: GLWECiphertext<&mut [u8]> = GLWECiphertext {
216            data: tmp_znx_ct,
217            basek: basek,
218            k,
219        };
220
221        (0..self.rows()).for_each(|row_i| {
222            vec_znx_pt.data.zero();
223
224            // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
225            module.vec_znx_add_scalar_inplace(&mut vec_znx_pt.data, 0, row_i, pt, 0);
226            module.vec_znx_normalize_inplace(basek, &mut vec_znx_pt.data, 0, scrach_2);
227
228            (0..cols).for_each(|col_j| {
229                // rlwe encrypt of vec_znx_pt into vec_znx_ct
230
231                vec_znx_ct.encrypt_sk_private(
232                    module,
233                    Some((&vec_znx_pt, col_j)),
234                    sk_dft,
235                    source_xa,
236                    source_xe,
237                    sigma,
238                    scrach_2,
239                );
240
241                // Switch vec_znx_ct into DFT domain
242                {
243                    let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, cols, size);
244
245                    (0..cols).for_each(|i| {
246                        module.vec_znx_dft(&mut vec_znx_dft_ct, i, &vec_znx_ct.data, i);
247                    });
248
249                    module.vmp_prepare_row(&mut self.data, row_i, col_j, &vec_znx_dft_ct);
250                }
251            });
252        });
253    }
254
255    pub(crate) fn expand_row<R, DataCi: AsRef<[u8]>, DataTsk: AsRef<[u8]>>(
256        &mut self,
257        module: &Module<FFT64>,
258        col_j: usize,
259        res: &mut R,
260        ci_dft: &VecZnxDft<DataCi, FFT64>,
261        tsk: &TensorKey<DataTsk, FFT64>,
262        scratch: &mut Scratch,
263    ) where
264        R: VecZnxToMut,
265    {
266        let cols: usize = self.rank() + 1;
267
268        assert!(scratch.available() >= GGSWCiphertext::expand_row_scratch_space(module, self.size(), tsk.size(), tsk.rank()));
269
270        // Example for rank 3:
271        //
272        // Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
273        // actually composed of that many rows and we focus on a specific row here
274        // implicitely given ci_dft.
275        //
276        // # Input
277        //
278        // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0    , a1    , a2    )
279        // col 1: (0, 0, 0, 0)
280        // col 2: (0, 0, 0, 0)
281        // col 3: (0, 0, 0, 0)
282        //
283        // # Output
284        //
285        // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0       , a1       , a2       )
286        // col 1: (-(b0s0 + b1s1 + b2s2)       , b0 + M[i], b1       , b2       )
287        // col 2: (-(c0s0 + c1s1 + c2s2)       , c0       , c1 + M[i], c2       )
288        // col 3: (-(d0s0 + d1s1 + d2s2)       , d0       , d1       , d2 + M[i])
289
290        let (mut tmp_dft_i, scratch1) = scratch.tmp_vec_znx_dft(module, cols, tsk.size());
291        {
292            let (mut tmp_dft_col_data, scratch2) = scratch1.tmp_vec_znx_dft(module, 1, self.size());
293
294            // Performs a key-switch for each combination of s[i]*s[j], i.e. for a0, a1, a2
295            //
296            // # Example for col=1
297            //
298            // a0 * (-(f0s0 + f1s1 + f1s2) + s0^2, f0, f1, f2) = (-(a0f0s0 + a0f1s1 + a0f1s2) + a0s0^2, a0f0, a0f1, a0f2)
299            // +
300            // a1 * (-(g0s0 + g1s1 + g1s2) + s0s1, g0, g1, g2) = (-(a1g0s0 + a1g1s1 + a1g1s2) + a1s0s1, a1g0, a1g1, a1g2)
301            // +
302            // a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
303            // =
304            // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0, x1, x2)
305            (1..cols).for_each(|col_i| {
306                // Extracts a[i] and multipies with Enc(s[i]s[j])
307                tmp_dft_col_data.extract_column(0, ci_dft, col_i);
308
309                if col_i == 1 {
310                    module.vmp_apply(
311                        &mut tmp_dft_i,
312                        &tmp_dft_col_data,
313                        &tsk.at(col_i - 1, col_j - 1).0.data, // Selects Enc(s[i]s[j])
314                        scratch2,
315                    );
316                } else {
317                    module.vmp_apply_add(
318                        &mut tmp_dft_i,
319                        &tmp_dft_col_data,
320                        &tsk.at(col_i - 1, col_j - 1).0.data, // Selects Enc(s[i]s[j])
321                        scratch2,
322                    );
323                }
324            });
325        }
326
327        // Adds -(sum a[i] * s[i]) + m)  on the i-th column of tmp_idft_i
328        //
329        // (-(x0s0 + x1s1 + x2s2) + a0s0s0 + a1s0s1 + a2s0s2, x0, x1, x2)
330        // +
331        // (0, -(a0s0 + a1s1 + a2s2) + M[i], 0, 0)
332        // =
333        // (-(x0s0 + x1s1 + x2s2) + s0(a0s0 + a1s1 + a2s2), x0 -(a0s0 + a1s1 + a2s2) + M[i], x1, x2)
334        // =
335        // (-(x0s0 + x1s1 + x2s2), x0 + M[i], x1, x2)
336        module.vec_znx_dft_add_inplace(&mut tmp_dft_i, col_j, ci_dft, 0);
337        let (mut tmp_idft, scratch2) = scratch1.tmp_vec_znx_big(module, 1, tsk.size());
338        (0..cols).for_each(|i| {
339            module.vec_znx_idft_tmp_a(&mut tmp_idft, 0, &mut tmp_dft_i, i);
340            module.vec_znx_big_normalize(self.basek(), res, i, &tmp_idft, 0, scratch2);
341        });
342    }
343
344    pub fn keyswitch<DataLhs: AsRef<[u8]>, DataKsk: AsRef<[u8]>, DataTsk: AsRef<[u8]>>(
345        &mut self,
346        module: &Module<FFT64>,
347        lhs: &GGSWCiphertext<DataLhs, FFT64>,
348        ksk: &GLWESwitchingKey<DataKsk, FFT64>,
349        tsk: &TensorKey<DataTsk, FFT64>,
350        scratch: &mut Scratch,
351    ) {
352        let cols: usize = self.rank() + 1;
353
354        let (res_data, scratch1) = scratch.tmp_vec_znx(&module, cols, self.size());
355        let mut res: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
356            data: res_data,
357            basek: self.basek(),
358            k: self.k(),
359        };
360
361        let (mut ci_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, lhs.size());
362
363        // Keyswitch the j-th row of the col 0
364        (0..lhs.rows()).for_each(|row_i| {
365            // Key-switch column 0, i.e.
366            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
367            lhs.keyswitch_internal_col0(module, row_i, &mut res, ksk, scratch2);
368
369            // Isolates DFT(a[i])
370            (0..cols).for_each(|col_i| {
371                module.vec_znx_dft(&mut ci_dft, col_i, &res.data, col_i);
372            });
373
374            module.vmp_prepare_row(&mut self.data, row_i, 0, &ci_dft);
375
376            // Generates
377            //
378            // col 1: (-(b0s0' + b1s1' + b2s2')    , b0 + M[i], b1       , b2       )
379            // col 2: (-(c0s0' + c1s1' + c2s2')    , c0       , c1 + M[i], c2       )
380            // col 3: (-(d0s0' + d1s1' + d2s2')    , d0       , d1       , d2 + M[i])
381            (1..cols).for_each(|col_j| {
382                self.expand_row(module, col_j, &mut res.data, &ci_dft, tsk, scratch2);
383
384                let (mut res_dft, _) = scratch2.tmp_vec_znx_dft(module, cols, self.size());
385                (0..cols).for_each(|i| {
386                    module.vec_znx_dft(&mut res_dft, i, &res.data, i);
387                });
388
389                module.vmp_prepare_row(&mut self.data, row_i, col_j, &res_dft);
390            });
391        })
392    }
393
394    pub fn keyswitch_inplace<DataKsk: AsRef<[u8]>, DataTsk: AsRef<[u8]>>(
395        &mut self,
396        module: &Module<FFT64>,
397        ksk: &GLWESwitchingKey<DataKsk, FFT64>,
398        tsk: &TensorKey<DataTsk, FFT64>,
399        scratch: &mut Scratch,
400    ) {
401        unsafe {
402            let self_ptr: *mut GGSWCiphertext<DataSelf, FFT64> = self as *mut GGSWCiphertext<DataSelf, FFT64>;
403            self.keyswitch(module, &*self_ptr, ksk, tsk, scratch);
404        }
405    }
406
407    pub fn automorphism<DataLhs: AsRef<[u8]>, DataAk: AsRef<[u8]>, DataTsk: AsRef<[u8]>>(
408        &mut self,
409        module: &Module<FFT64>,
410        lhs: &GGSWCiphertext<DataLhs, FFT64>,
411        auto_key: &AutomorphismKey<DataAk, FFT64>,
412        tensor_key: &TensorKey<DataTsk, FFT64>,
413        scratch: &mut Scratch,
414    ) {
415        #[cfg(debug_assertions)]
416        {
417            assert_eq!(
418                self.rank(),
419                lhs.rank(),
420                "ggsw_out rank: {} != ggsw_in rank: {}",
421                self.rank(),
422                lhs.rank()
423            );
424            assert_eq!(
425                self.rank(),
426                auto_key.rank(),
427                "ggsw_in rank: {} != auto_key rank: {}",
428                self.rank(),
429                auto_key.rank()
430            );
431            assert_eq!(
432                self.rank(),
433                tensor_key.rank(),
434                "ggsw_in rank: {} != tensor_key rank: {}",
435                self.rank(),
436                tensor_key.rank()
437            );
438            assert!(
439                scratch.available()
440                    >= GGSWCiphertext::automorphism_scratch_space(
441                        module,
442                        self.size(),
443                        lhs.size(),
444                        auto_key.size(),
445                        tensor_key.size(),
446                        self.rank()
447                    )
448            )
449        };
450
451        let cols: usize = self.rank() + 1;
452
453        let (res_data, scratch1) = scratch.tmp_vec_znx(&module, cols, self.size());
454        let mut res: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
455            data: res_data,
456            basek: self.basek(),
457            k: self.k(),
458        };
459
460        let (mut ci_dft, scratch2) = scratch1.tmp_vec_znx_dft(module, cols, self.size());
461
462        // Keyswitch the j-th row of the col 0
463        (0..lhs.rows()).for_each(|row_i| {
464            // Key-switch column 0, i.e.
465            // col 0: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2)
466            lhs.keyswitch_internal_col0(module, row_i, &mut res, &auto_key.key, scratch2);
467
468            // Isolates DFT(AUTO(a[i]))
469            (0..cols).for_each(|col_i| {
470                // (-(a0pi^-1(s0) + a1pi^-1(s1) + a2pi^-1(s2)) + M[i], a0, a1, a2) -> (-(a0s0 + a1s1 + a2s2) + pi(M[i]), a0, a1, a2)
471                module.vec_znx_automorphism_inplace(auto_key.p(), &mut res.data, col_i);
472                module.vec_znx_dft(&mut ci_dft, col_i, &res.data, col_i);
473            });
474
475            module.vmp_prepare_row(&mut self.data, row_i, 0, &ci_dft);
476
477            // Generates
478            //
479            // col 1: (-(b0s0 + b1s1 + b2s2)    , b0 + pi(M[i]), b1           , b2           )
480            // col 2: (-(c0s0 + c1s1 + c2s2)    , c0           , c1 + pi(M[i]), c2           )
481            // col 3: (-(d0s0 + d1s1 + d2s2)    , d0           , d1           , d2 + pi(M[i]))
482            (1..cols).for_each(|col_j| {
483                self.expand_row(module, col_j, &mut res.data, &ci_dft, tensor_key, scratch2);
484
485                let (mut res_dft, _) = scratch2.tmp_vec_znx_dft(module, cols, self.size());
486                (0..cols).for_each(|i| {
487                    module.vec_znx_dft(&mut res_dft, i, &res.data, i);
488                });
489
490                module.vmp_prepare_row(&mut self.data, row_i, col_j, &res_dft);
491            });
492        })
493    }
494
495    pub fn automorphism_inplace<DataKsk: AsRef<[u8]>, DataTsk: AsRef<[u8]>>(
496        &mut self,
497        module: &Module<FFT64>,
498        auto_key: &AutomorphismKey<DataKsk, FFT64>,
499        tensor_key: &TensorKey<DataTsk, FFT64>,
500        scratch: &mut Scratch,
501    ) {
502        unsafe {
503            let self_ptr: *mut GGSWCiphertext<DataSelf, FFT64> = self as *mut GGSWCiphertext<DataSelf, FFT64>;
504            self.automorphism(module, &*self_ptr, auto_key, tensor_key, scratch);
505        }
506    }
507
508    pub fn external_product<DataLhs: AsRef<[u8]>, DataRhs: AsRef<[u8]>>(
509        &mut self,
510        module: &Module<FFT64>,
511        lhs: &GGSWCiphertext<DataLhs, FFT64>,
512        rhs: &GGSWCiphertext<DataRhs, FFT64>,
513        scratch: &mut Scratch,
514    ) {
515        #[cfg(debug_assertions)]
516        {
517            assert_eq!(
518                self.rank(),
519                lhs.rank(),
520                "ggsw_out rank: {} != ggsw_in rank: {}",
521                self.rank(),
522                lhs.rank()
523            );
524            assert_eq!(
525                self.rank(),
526                rhs.rank(),
527                "ggsw_in rank: {} != ggsw_apply rank: {}",
528                self.rank(),
529                rhs.rank()
530            );
531        }
532
533        let (tmp_in_data, scratch1) = scratch.tmp_vec_znx_dft(module, lhs.rank() + 1, lhs.size());
534
535        let mut tmp_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
536            data: tmp_in_data,
537            basek: lhs.basek(),
538            k: lhs.k(),
539        };
540
541        let (tmp_out_data, scratch2) = scratch1.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
542
543        let mut tmp_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
544            data: tmp_out_data,
545            basek: self.basek(),
546            k: self.k(),
547        };
548
549        (0..self.rank() + 1).for_each(|col_i| {
550            (0..self.rows()).for_each(|row_j| {
551                lhs.get_row(module, row_j, col_i, &mut tmp_in);
552                tmp_out.external_product(module, &tmp_in, rhs, scratch2);
553                self.set_row(module, row_j, col_i, &tmp_out);
554            });
555        });
556
557        tmp_out.data.zero();
558
559        (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
560            (0..self.rank() + 1).for_each(|col_j| {
561                self.set_row(module, row_i, col_j, &tmp_out);
562            });
563        });
564    }
565
566    pub fn external_product_inplace<DataRhs: AsRef<[u8]>>(
567        &mut self,
568        module: &Module<FFT64>,
569        rhs: &GGSWCiphertext<DataRhs, FFT64>,
570        scratch: &mut Scratch,
571    ) {
572        #[cfg(debug_assertions)]
573        {
574            assert_eq!(
575                self.rank(),
576                rhs.rank(),
577                "ggsw_out rank: {} != ggsw_apply: {}",
578                self.rank(),
579                rhs.rank()
580            );
581        }
582
583        let (tmp_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
584
585        let mut tmp: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
586            data: tmp_data,
587            basek: self.basek(),
588            k: self.k(),
589        };
590
591        (0..self.rank() + 1).for_each(|col_i| {
592            (0..self.rows()).for_each(|row_j| {
593                self.get_row(module, row_j, col_i, &mut tmp);
594                tmp.external_product_inplace(module, rhs, scratch1);
595                self.set_row(module, row_j, col_i, &tmp);
596            });
597        });
598    }
599}
600
601impl<DataSelf: AsRef<[u8]>> GGSWCiphertext<DataSelf, FFT64> {
602    pub(crate) fn keyswitch_internal_col0<DataRes: AsMut<[u8]> + AsRef<[u8]>, DataKsk: AsRef<[u8]>>(
603        &self,
604        module: &Module<FFT64>,
605        row_i: usize,
606        res: &mut GLWECiphertext<DataRes>,
607        ksk: &GLWESwitchingKey<DataKsk, FFT64>,
608        scratch: &mut Scratch,
609    ) {
610        #[cfg(debug_assertions)]
611        {
612            assert_eq!(self.rank(), ksk.rank());
613            assert_eq!(res.rank(), ksk.rank());
614            assert!(
615                scratch.available()
616                    >= GGSWCiphertext::keyswitch_internal_col0_scratch_space(
617                        module,
618                        res.size(),
619                        self.size(),
620                        ksk.size(),
621                        ksk.rank()
622                    )
623            )
624        }
625
626        let (tmp_dft_in_data, scratch2) = scratch.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
627        let mut tmp_dft_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
628            data: tmp_dft_in_data,
629            basek: self.basek(),
630            k: self.k(),
631        };
632        self.get_row(module, row_i, 0, &mut tmp_dft_in);
633        res.keyswitch_from_fourier(module, &tmp_dft_in, ksk, scratch2);
634    }
635}
636
637impl<DataSelf: AsRef<[u8]>> GetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64> {
638    fn get_row<R: AsMut<[u8]> + AsRef<[u8]>>(
639        &self,
640        module: &Module<FFT64>,
641        row_i: usize,
642        col_j: usize,
643        res: &mut GLWECiphertextFourier<R, FFT64>,
644    ) {
645        module.vmp_extract_row(&mut res.data, &self.data, row_i, col_j);
646    }
647}
648
649impl<DataSelf: AsMut<[u8]> + AsRef<[u8]>> SetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64> {
650    fn set_row<R: AsRef<[u8]>>(
651        &mut self,
652        module: &Module<FFT64>,
653        row_i: usize,
654        col_j: usize,
655        a: &GLWECiphertextFourier<R, FFT64>,
656    ) {
657        module.vmp_prepare_row(&mut self.data, row_i, col_j, &a.data);
658    }
659}