Fix compressed encryptions & add GGSW compressed encryption (#67)

* Added decompress test

* updated encryption sampling & fixed bug in glwe -> lwe test

* Added GGSW compressed encryption

core/src/blind_rotation/test/cggi.rs

@@ -3,7 +3,7 @@ use backend::{
         api::{
             MatZnxAlloc, ModuleNew, ScalarZnxAlloc, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddNormal,
             VecZnxAddScalarInplace, VecZnxAlloc, VecZnxAllocBytes, VecZnxEncodeCoeffsi64, VecZnxFillUniform, VecZnxRotateInplace,
-            VecZnxSwithcDegree, ZnxView,
+            VecZnxSub, VecZnxSwithcDegree, ZnxView,
         },
         layouts::{Backend, Module, ScratchOwned},
         oep::{
@@ -57,7 +57,8 @@ pub(crate) trait CGGITestModuleFamily<B: Backend> = CCGIBlindRotationFamily<B>
     + VecZnxEncodeCoeffsi64
     + VecZnxRotateInplace
     + VecZnxSwithcDegree
-    + MatZnxAlloc;
+    + MatZnxAlloc
+    + VecZnxSub;
 pub(crate) trait CGGITestScratchFamily<B: Backend> = VecZnxDftAllocBytesImpl<B>
     + VecZnxBigAllocBytesImpl<B>
     + ScratchOwnedAllocImpl<B>

core/src/elem.rs

@@ -1,4 +1,7 @@
-use backend::hal::api::ZnxInfos;
+use backend::hal::{
+    api::{VecZnxCopy, VecZnxFillUniform, ZnxInfos},
+    layouts::{Backend, Module},
+};
 
 pub trait Infos {
     type Inner: ZnxInfos;
@@ -52,3 +55,9 @@ pub trait SetMetaData {
     fn set_basek(&mut self, basek: usize);
     fn set_k(&mut self, k: usize);
 }
+
+pub trait Decompress<B: Backend, C> {
+    fn decompress(&mut self, module: &Module<B>, other: &C)
+    where
+        Module<B>: VecZnxFillUniform + VecZnxCopy;
+}

core/src/gglwe/encryption.rs

@@ -192,6 +192,7 @@ impl<D: DataMut> GGLWECiphertextCompressed<D> {
         let basek: usize = self.basek();
         let k: usize = self.k();
         let rank_in: usize = self.rank_in();
+        let cols: usize = self.rank_out() + 1;
 
         let mut source_xa = Source::new(seed);
 
@@ -217,6 +218,7 @@ impl<D: DataMut> GGLWECiphertextCompressed<D> {
                     self.basek(),
                     self.k(),
                     &mut self.at_mut(row_i, col_i).data,
+                    cols,
                     true,
                     Some((&tmp_pt, 0)),
                     sk,

core/src/gglwe/layouts_compressed.rs

@@ -3,7 +3,7 @@ use backend::hal::{
     layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, WriterTo},
 };
 
-use crate::{AutomorphismKey, GGLWECiphertext, GLWECiphertextCompressed, GLWESwitchingKey, GLWETensorKey, Infos};
+use crate::{AutomorphismKey, Decompress, GGLWECiphertext, GLWECiphertextCompressed, GLWESwitchingKey, GLWETensorKey, Infos};
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 
 #[derive(PartialEq, Eq)]
@@ -169,8 +169,8 @@ impl<D: DataRef> WriterTo for GGLWECiphertextCompressed<D> {
     }
 }
 
-impl<D: DataMut> GGLWECiphertext<D> {
+impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, GGLWECiphertextCompressed<DR>> for GGLWECiphertext<D> {
-    pub fn decompress<DataOther: DataRef, B: Backend>(&mut self, module: &Module<B>, other: &GGLWECiphertextCompressed<DataOther>)
+    fn decompress(&mut self, module: &Module<B>, other: &GGLWECiphertextCompressed<DR>)
     where
         Module<B>: VecZnxFillUniform + VecZnxCopy,
     {

core/src/ggsw/encryption.rs

@@ -6,7 +6,10 @@ use backend::hal::{
 };
 use sampling::source::Source;
 
-use crate::{GGSWCiphertext, GLWECiphertext, GLWEEncryptSkFamily, GLWESecretExec, Infos, TakeGLWEPt};
+use crate::{
+    GGLWEEncryptSkFamily, GGSWCiphertext, GGSWCiphertextCompressed, GLWECiphertext, GLWEEncryptSkFamily, GLWESecretExec, Infos,
+    TakeGLWEPt, encrypt_sk_internal,
+};
 
 pub trait GGSWEncryptSkFamily<B: Backend> = GLWEEncryptSkFamily<B>;
 
@@ -77,3 +80,79 @@ impl<DataSelf: DataMut> GGSWCiphertext<DataSelf> {
         });
     }
 }
+
+impl GGSWCiphertextCompressed<Vec<u8>> {
+    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> usize
+    where
+        Module<B>: GGSWEncryptSkFamily<B> + VecZnxAllocBytes,
+    {
+        GGSWCiphertext::encrypt_sk_scratch_space(module, basek, k, rank)
+    }
+}
+
+impl<DataSelf: DataMut> GGSWCiphertextCompressed<DataSelf> {
+    pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
+        &mut self,
+        module: &Module<B>,
+        pt: &ScalarZnx<DataPt>,
+        sk: &GLWESecretExec<DataSk, B>,
+        seed_xa: [u8; 32],
+        source_xe: &mut Source,
+        sigma: f64,
+        scratch: &mut Scratch<B>,
+    ) where
+        Module<B>: GGSWEncryptSkFamily<B> + VecZnxAddScalarInplace,
+        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
+    {
+        #[cfg(debug_assertions)]
+        {
+            use backend::hal::api::ZnxInfos;
+
+            assert_eq!(self.rank(), sk.rank());
+            assert_eq!(self.n(), module.n());
+            assert_eq!(pt.n(), module.n());
+            assert_eq!(sk.n(), module.n());
+        }
+
+        let basek: usize = self.basek();
+        let k: usize = self.k();
+        let rank: usize = self.rank();
+        let cols: usize = rank + 1;
+        let digits: usize = self.digits();
+
+        let (mut tmp_pt, scratch_1) = scratch.take_glwe_pt(module, basek, k);
+
+        let mut source = Source::new(seed_xa);
+
+        (0..self.rows()).for_each(|row_i| {
+            tmp_pt.data.zero();
+
+            // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
+            module.vec_znx_add_scalar_inplace(&mut tmp_pt.data, 0, (digits - 1) + row_i * digits, pt, 0);
+            module.vec_znx_normalize_inplace(basek, &mut tmp_pt.data, 0, scratch_1);
+
+            (0..rank + 1).for_each(|col_j| {
+                // rlwe encrypt of vec_znx_pt into vec_znx_ct
+
+                let (seed, mut source_xa_tmp) = source.branch();
+
+                self.seed[row_i * cols + col_j] = seed;
+                
+                encrypt_sk_internal(
+                    module,
+                    self.basek(),
+                    self.k(),
+                    &mut self.at_mut(row_i, col_j).data,
+                    cols,
+                    true,
+                    Some((&tmp_pt, col_j)),
+                    sk,
+                    &mut source_xa_tmp,
+                    source_xe,
+                    sigma,
+                    scratch_1,
+                );
+            });
+        });
+    }
+}

core/src/ggsw/layout_compressed.rs

@@ -3,11 +3,16 @@ use backend::hal::{
     layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, WriterTo},
 };
 
-use crate::{GGLWECiphertextCompressed, GGSWCiphertext, Infos};
+use crate::{Decompress, GGSWCiphertext, GLWECiphertextCompressed, Infos};
 
 #[derive(PartialEq, Eq)]
 pub struct GGSWCiphertextCompressed<D: Data> {
-    pub(crate) data: GGLWECiphertextCompressed<D>,
+    pub(crate) data: MatZnx<D>,
+    pub(crate) basek: usize,
+    pub(crate) k: usize,
+    pub(crate) digits: usize,
+    pub(crate) rank: usize,
+    pub(crate) seed: Vec<[u8; 32]>,
 }
 
 impl GGSWCiphertextCompressed<Vec<u8>> {
@@ -15,8 +20,31 @@ impl GGSWCiphertextCompressed<Vec<u8>> {
     where
         Module<B>: MatZnxAlloc,
     {
-        GGSWCiphertextCompressed {
+        let size: usize = k.div_ceil(basek);
-            data: GGLWECiphertextCompressed::alloc(module, basek, k, rows, digits, rank, rank),
+        debug_assert!(digits > 0, "invalid ggsw: `digits` == 0");
+
+        debug_assert!(
+            size > digits,
+            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
+            size,
+            digits
+        );
+
+        assert!(
+            rows * digits <= size,
+            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
+            rows,
+            digits,
+            size
+        );
+
+        Self {
+            data: module.mat_znx_alloc(rows, rank + 1, 1, k.div_ceil(basek)),
+            basek,
+            k: k,
+            digits,
+            rank,
+            seed: vec![[0u8; 32]; rows * (rank + 1)],
         }
     }
 
@@ -24,7 +52,48 @@ impl GGSWCiphertextCompressed<Vec<u8>> {
     where
         Module<B>: MatZnxAllocBytes,
     {
-        GGLWECiphertextCompressed::bytes_of(module, basek, k, rows, digits, rank)
+        let size: usize = k.div_ceil(basek);
+        debug_assert!(
+            size > digits,
+            "invalid ggsw: ceil(k/basek): {} <= digits: {}",
+            size,
+            digits
+        );
+
+        assert!(
+            rows * digits <= size,
+            "invalid ggsw: rows: {} * digits:{} > ceil(k/basek): {}",
+            rows,
+            digits,
+            size
+        );
+
+        module.mat_znx_alloc_bytes(rows, rank + 1, 1, size)
+    }
+}
+
+impl<D: DataRef> GGSWCiphertextCompressed<D> {
+    pub fn at(&self, row: usize, col: usize) -> GLWECiphertextCompressed<&[u8]> {
+        GLWECiphertextCompressed {
+            data: self.data.at(row, col),
+            basek: self.basek,
+            k: self.k,
+            rank: self.rank(),
+            seed: self.seed[row * (self.rank() + 1) + col],
+        }
+    }
+}
+
+impl<D: DataMut> GGSWCiphertextCompressed<D> {
+    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertextCompressed<&mut [u8]> {
+        let rank: usize = self.rank();
+        GLWECiphertextCompressed {
+            data: self.data.at_mut(row, col),
+            basek: self.basek,
+            k: self.k,
+            rank: rank,
+            seed: self.seed[row * (rank + 1) + col],
+        }
     }
 }
 
@@ -32,25 +101,25 @@ impl<D: Data> Infos for GGSWCiphertextCompressed<D> {
     type Inner = MatZnx<D>;
 
     fn inner(&self) -> &Self::Inner {
-        self.data.inner()
+        &self.data
     }
 
     fn basek(&self) -> usize {
-        self.data.basek()
+        self.basek
     }
 
     fn k(&self) -> usize {
-        self.data.k()
+        self.k
     }
 }
 
 impl<D: Data> GGSWCiphertextCompressed<D> {
     pub fn rank(&self) -> usize {
-        self.data.rank()
+        self.rank
     }
 
     pub fn digits(&self) -> usize {
-        self.data.digits()
+        self.digits
     }
 }
 
@@ -66,15 +135,23 @@ impl<D: DataRef> WriterTo for GGSWCiphertextCompressed<D> {
     }
 }
 
-impl<D: DataMut> GGSWCiphertext<D> {
+impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, GGSWCiphertextCompressed<DR>> for GGSWCiphertext<D> {
-    pub fn decompress<DataOther: DataRef, B: Backend>(&mut self, module: &Module<B>, other: &GGSWCiphertextCompressed<DataOther>)
+    fn decompress(&mut self, module: &Module<B>, other: &GGSWCiphertextCompressed<DR>)
     where
         Module<B>: VecZnxFillUniform + VecZnxCopy,
     {
-        let rows = self.rows();
+        #[cfg(debug_assertions)]
+        {
+            assert_eq!(self.rank(), other.rank())
+        }
+
+        let rows: usize = self.rows();
+        let rank: usize = self.rank();
         (0..rows).for_each(|row_i| {
-            self.at_mut(row_i, 0)
+            (0..rank + 1).for_each(|col_j| {
-                .decompress(module, &other.data.at(row_i, 0));
+                self.at_mut(row_i, col_j)
+                    .decompress(module, &other.at(row_i, col_j));
+            });
         });
     }
 }

core/src/ggsw/test/cpu_spqlios/fft64.rs

@@ -4,8 +4,8 @@ use backend::{
 };
 
 use crate::ggsw::test::generic_tests::{
-    test_automorphism, test_automorphism_inplace, test_encrypt_sk, test_external_product, test_external_product_inplace,
+    test_automorphism, test_automorphism_inplace, test_encrypt_sk, test_encrypt_sk_compressed, test_external_product,
-    test_keyswitch, test_keyswitch_inplace,
+    test_external_product_inplace, test_keyswitch, test_keyswitch_inplace,
 };
 
 #[test]
@@ -23,6 +23,21 @@ fn encrypt_sk() {
     });
 }
 
+#[test]
+fn encrypt_sk_compressed() {
+    let log_n: usize = 8;
+    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
+    let basek: usize = 12;
+    let k_ct: usize = 54;
+    let digits: usize = k_ct / basek;
+    (1..4).for_each(|rank| {
+        (1..digits + 1).for_each(|di| {
+            println!("test encrypt_sk_compressed digits: {} rank: {}", di, rank);
+            test_encrypt_sk_compressed(&module, basek, k_ct, di, rank, 3.2);
+        });
+    });
+}
+
 #[test]
 fn keyswitch() {
     let log_n: usize = 8;

core/src/ggsw/test/generic_tests.rs

@@ -1,7 +1,7 @@
 use backend::hal::{
     api::{
         MatZnxAlloc, ScalarZnxAlloc, ScalarZnxAllocBytes, ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace,
-        VecZnxAlloc, VecZnxAllocBytes, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxRotateInplace, VecZnxStd,
+        VecZnxAlloc, VecZnxAllocBytes, VecZnxAutomorphism, VecZnxAutomorphismInplace, VecZnxCopy, VecZnxRotateInplace, VecZnxStd,
         VecZnxSubABInplace, VecZnxSwithcDegree, ZnxViewMut,
     },
     layouts::{Backend, Module, ScalarZnx, ScalarZnxToMut, ScratchOwned},
@@ -13,9 +13,10 @@ use backend::hal::{
 use sampling::source::Source;
 
 use crate::{
-    AutomorphismKey, AutomorphismKeyExec, GGLWEExecLayoutFamily, GGSWAssertNoiseFamily, GGSWCiphertext, GGSWCiphertextExec,
+    AutomorphismKey, AutomorphismKeyExec, Decompress, GGLWEExecLayoutFamily, GGSWAssertNoiseFamily, GGSWCiphertext,
-    GGSWEncryptSkFamily, GGSWKeySwitchFamily, GLWESecret, GLWESecretExec, GLWESecretFamily, GLWESwitchingKey,
+    GGSWCiphertextCompressed, GGSWCiphertextExec, GGSWEncryptSkFamily, GGSWKeySwitchFamily, GLWESecret, GLWESecretExec,
-    GLWESwitchingKeyEncryptSkFamily, GLWESwitchingKeyExec, GLWETensorKey, GLWETensorKeyEncryptSkFamily, GLWETensorKeyExec,
+    GLWESecretFamily, GLWESwitchingKey, GLWESwitchingKeyEncryptSkFamily, GLWESwitchingKeyExec, GLWETensorKey,
+    GLWETensorKeyEncryptSkFamily, GLWETensorKeyExec,
     noise::{noise_ggsw_keyswitch, noise_ggsw_product},
 };
 
@@ -29,7 +30,8 @@ pub(crate) trait TestModuleFamily<B: Backend> = GLWESecretFamily<B>
     + VecZnxAddScalarInplace
     + VecZnxSubABInplace
     + VecZnxStd
-    + ScalarZnxAllocBytes;
+    + ScalarZnxAllocBytes
+    + VecZnxCopy;
 pub(crate) trait TestScratchFamily<B: Backend> = TakeVecZnxDftImpl<B>
     + TakeVecZnxBigImpl<B>
     + TakeSvpPPolImpl<B>
@@ -83,6 +85,58 @@ where
     ct.assert_noise(module, &sk_exec, &pt_scalar, &noise_f);
 }
 
+pub(crate) fn test_encrypt_sk_compressed<B: Backend>(
+    module: &Module<B>,
+    basek: usize,
+    k: usize,
+    digits: usize,
+    rank: usize,
+    sigma: f64,
+) where
+    Module<B>: TestModuleFamily<B>,
+    B: TestScratchFamily<B>,
+{
+    let rows: usize = (k - digits * basek) / (digits * basek);
+
+    let mut ct_compressed: GGSWCiphertextCompressed<Vec<u8>> =
+        GGSWCiphertextCompressed::alloc(module, basek, k, rows, digits, rank);
+
+    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.scalar_znx_alloc(1);
+
+    let mut source_xs: Source = Source::new([0u8; 32]);
+    let mut source_xe: Source = Source::new([0u8; 32]);
+
+    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
+
+    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(GGSWCiphertextCompressed::encrypt_sk_scratch_space(
+        module, basek, k, rank,
+    ));
+
+    let mut sk: GLWESecret<Vec<u8>> = GLWESecret::alloc(module, rank);
+    sk.fill_ternary_prob(0.5, &mut source_xs);
+    let mut sk_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk);
+    sk_exec.prepare(module, &sk);
+
+    let seed_xa: [u8; 32] = [1u8; 32];
+
+    ct_compressed.encrypt_sk(
+        module,
+        &pt_scalar,
+        &sk_exec,
+        seed_xa,
+        &mut source_xe,
+        sigma,
+        scratch.borrow(),
+    );
+
+    let noise_f = |_col_i: usize| -(k as f64) + sigma.log2() + 0.5;
+
+    let mut ct: GGSWCiphertext<Vec<u8>> = GGSWCiphertext::alloc(module, basek, k, rows, digits, rank);
+    ct.decompress(module, &ct_compressed);
+
+    ct.assert_noise(module, &sk_exec, &pt_scalar, &noise_f);
+}
+
 pub(crate) fn test_keyswitch<B: Backend>(
     module: &Module<B>,
     basek: usize,

core/src/glwe/encryption.rs

@@ -1,10 +1,10 @@
 use backend::hal::{
     api::{
         ScalarZnxAllocBytes, ScratchAvailable, SvpApply, SvpApplyInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx,
-        TakeSvpPPol, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxBigAddNormal, VecZnxBigAddSmallInplace,
+        TakeSvpPPol, TakeVecZnx, TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxAllocBytes, VecZnxBigAddNormal,
-        VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume,
+        VecZnxBigAddSmallInplace, VecZnxBigAllocBytes, VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx,
-        VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSubABInplace, ZnxInfos,
+        VecZnxDftToVecZnxBigConsume, VecZnxFillUniform, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes,
-        ZnxZero,
+        VecZnxSub, VecZnxSubABInplace, ZnxInfos, ZnxZero,
     },
     layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, VecZnxBig},
 };
@@ -26,7 +26,9 @@ pub trait GLWEEncryptSkFamily<B: Backend> = VecZnxDftAllocBytes
     + VecZnxAddInplace
     + VecZnxNormalizeInplace<B>
     + VecZnxAddNormal
-    + VecZnxNormalize<B>;
+    + VecZnxNormalize<B>
+    + VecZnxSub
+    + VecZnxAllocBytes;
 
 pub trait GLWEEncryptPkFamily<B: Backend> = VecZnxDftAllocBytes
     + VecZnxBigAllocBytes
@@ -47,7 +49,7 @@ impl GLWECiphertext<Vec<u8>> {
     {
         let size: usize = k.div_ceil(basek);
         module.vec_znx_normalize_tmp_bytes(module.n())
-            + module.vec_znx_dft_alloc_bytes(1, size)
+            + 2 * module.vec_znx_alloc_bytes(1, size)
             + module.vec_znx_dft_alloc_bytes(1, size)
     }
     pub fn encrypt_pk_scratch_space<B: Backend>(module: &Module<B>, basek: usize, k: usize) -> usize
@@ -146,11 +148,13 @@ impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
         Module<B>: GLWEEncryptSkFamily<B>,
         Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
     {
+        let cols: usize = self.rank() + 1;
         encrypt_sk_internal(
             module,
             self.basek(),
             self.k(),
             &mut self.data,
+            cols,
             false,
             pt,
             sk,
@@ -345,11 +349,13 @@ impl<D: DataMut> GLWECiphertextCompressed<D> {
         Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx<B>,
     {
         let mut source_xa = Source::new(seed_xa);
+        let cols: usize = self.rank() + 1;
         encrypt_sk_internal(
             module,
             self.basek(),
             self.k(),
             &mut self.data,
+            cols,
             true,
             pt,
             sk,
@@ -367,6 +373,7 @@ pub(crate) fn encrypt_sk_internal<DataCt: DataMut, DataPt: DataRef, DataSk: Data
     basek: usize,
     k: usize,
     ct: &mut VecZnx<DataCt>,
+    cols: usize,
     compressed: bool,
     pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
     sk: &GLWESecretExec<DataSk, B>,
@@ -381,12 +388,6 @@ pub(crate) fn encrypt_sk_internal<DataCt: DataMut, DataPt: DataRef, DataSk: Data
     #[cfg(debug_assertions)]
     {
         if compressed {
-            if let Some((_, col)) = pt {
-                assert_eq!(
-                    col, 0,
-                    "invalid plaintext: cannot put pt in col>0 if compressed encryption"
-                )
-            }
             assert_eq!(
                 ct.cols(),
                 1,
@@ -397,14 +398,15 @@ pub(crate) fn encrypt_sk_internal<DataCt: DataMut, DataPt: DataRef, DataSk: Data
     }
 
     let size: usize = ct.size();
-    let cols: usize = ct.cols();
 
     let (mut c0, scratch_1) = scratch.take_vec_znx(module, 1, size);
     c0.zero();
 
     {
-        // c[i] = uniform
+        let (mut ci, scratch_2) = scratch_1.take_vec_znx(module, 1, size);
-        // c[0] -= c[i] * s[i],
+
+        // ct[i] = uniform
+        // ct[0] -= c[i] * s[i],
         (1..cols).for_each(|i| {
             let col_ct: usize;
             if compressed {
@@ -413,29 +415,34 @@ pub(crate) fn encrypt_sk_internal<DataCt: DataMut, DataPt: DataRef, DataSk: Data
                 col_ct = i;
             }
 
-            let (mut ci_dft, scratch_2) = scratch_1.take_vec_znx_dft(module, 1, size);
+            // ct[i] = uniform (+ pt)
-
-            // c[i] = uniform
             module.vec_znx_fill_uniform(basek, ct, col_ct, k, source_xa);
 
-            // c[i] = norm(IDFT(DFT(c[i]) * DFT(s[i])))
+            let (mut ci_dft, scratch_3) = scratch_2.take_vec_znx_dft(module, 1, size);
+
+            // ci = ct[i] - pt
+            // i.e. we act as we sample ct[i] already as uniform + pt
+            // and if there is a pt, then we subtract it before applying DFT
+            if let Some((pt, col)) = pt {
+                if i == col {
+                    module.vec_znx_sub(&mut ci, 0, ct, col_ct, &pt.data, 0);
+                    module.vec_znx_normalize_inplace(basek, &mut ci, 0, scratch_3);
+                    module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, 0, &ci, 0);
+                } else {
                     module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, 0, ct, col_ct);
+                }
+            } else {
+                module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, 0, ct, col_ct);
+            }
+
             module.svp_apply_inplace(&mut ci_dft, 0, &sk.data, i - 1);
             let ci_big: VecZnxBig<&mut [u8], B> = module.vec_znx_dft_to_vec_znx_big_consume(ci_dft);
 
             // use c[0] as buffer, which is overwritten later by the normalization step
-            module.vec_znx_big_normalize(basek, ct, 0, &ci_big, 0, scratch_2);
+            module.vec_znx_big_normalize(basek, &mut ci, 0, &ci_big, 0, scratch_3);
 
             // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
-            module.vec_znx_sub_ab_inplace(&mut c0, 0, ct, 0);
+            module.vec_znx_sub_ab_inplace(&mut c0, 0, &ci, 0);
-
-            // c[i] += m if col = i
-            // note: case cannot happen if compressed = true
-            if let Some((pt, col)) = pt {
-                if i == col {
-                    module.vec_znx_add_inplace(ct, i, &pt.data, 0);
-                }
-            }
         });
     }
 

core/src/glwe/keyswitch.rs

@@ -119,6 +119,27 @@ impl<DataSelf: DataRef> GLWECiphertext<DataSelf> {
                     rhs.digits(),
                     rhs.rank_in(),
                     rhs.rank_out(),
+                ),
+            "scratch.available()={} < GLWECiphertext::keyswitch_scratch_space(
+                    module,
+                    self.basek(),
+                    self.k(),
+                    lhs.k(),
+                    rhs.k(),
+                    rhs.digits(),
+                    rhs.rank_in(),
+                    rhs.rank_out(),
+                )={}",
+            scratch.available(),
+            GLWECiphertext::keyswitch_scratch_space(
+                module,
+                self.basek(),
+                self.k(),
+                lhs.k(),
+                rhs.k(),
+                rhs.digits(),
+                rhs.rank_in(),
+                rhs.rank_out(),
             )
         );
     }

core/src/glwe/layout.rs

@@ -6,7 +6,7 @@ use backend::hal::{
 };
 use sampling::source::Source;
 
-use crate::{GLWEOps, Infos, SetMetaData};
+use crate::{Decompress, GLWEOps, Infos, SetMetaData};
 
 #[derive(PartialEq, Eq, Clone)]
 pub struct GLWECiphertext<D: Data> {
@@ -266,10 +266,9 @@ impl<D: DataRef> WriterTo for GLWECiphertextCompressed<D> {
     }
 }
 
-impl<D: DataMut> GLWECiphertext<D> {
+impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, GLWECiphertextCompressed<DR>> for GLWECiphertext<D> {
-    pub fn decompress<DataOther, B: Backend>(&mut self, module: &Module<B>, other: &GLWECiphertextCompressed<DataOther>)
+    fn decompress(&mut self, module: &Module<B>, other: &GLWECiphertextCompressed<DR>)
     where
-        DataOther: DataRef,
         Module<B>: VecZnxFillUniform + VecZnxCopy,
     {
         #[cfg(debug_assertions)]
@@ -299,7 +298,9 @@ impl<D: DataMut> GLWECiphertext<D> {
             self.decompress_internal(module, other, &mut source);
         }
     }
+}
 
+impl<D: DataMut> GLWECiphertext<D> {
     pub(crate) fn decompress_internal<DataOther, B: Backend>(
         &mut self,
         module: &Module<B>,
@@ -309,13 +310,19 @@ impl<D: DataMut> GLWECiphertext<D> {
         DataOther: DataRef,
         Module<B>: VecZnxFillUniform + VecZnxCopy,
     {
+        #[cfg(debug_assertions)]
+        {
+            assert_eq!(self.rank(), other.rank())
+        }
+
         let k: usize = other.k;
         let basek: usize = other.basek;
-        let cols: usize = other.cols();
+        let cols: usize = other.rank() + 1;
         module.vec_znx_copy(&mut self.data, 0, &other.data, 0);
         (1..cols).for_each(|i| {
             module.vec_znx_fill_uniform(basek, &mut self.data, i, k, source);
         });
+
         self.basek = basek;
         self.k = k;
     }

core/src/glwe/tests/generic_encryption.rs

@@ -12,7 +12,7 @@ use backend::hal::{
 use sampling::source::Source;
 
 use crate::{
-    GLWECiphertext, GLWECiphertextCompressed, GLWEDecryptFamily, GLWEEncryptPkFamily, GLWEEncryptSkFamily, GLWEOps,
+    Decompress, GLWECiphertext, GLWECiphertextCompressed, GLWEDecryptFamily, GLWEEncryptPkFamily, GLWEEncryptSkFamily, GLWEOps,
     GLWEPlaintext, GLWEPublicKey, GLWEPublicKeyExec, GLWESecret, GLWESecretExec, GLWESecretFamily, Infos,
 };
 
@@ -125,7 +125,12 @@ pub(crate) fn test_encrypt_sk_compressed<B: Backend>(
     let noise_have: f64 = module.vec_znx_std(basek, &pt_want.data, 0) * (ct.k() as f64).exp2();
     let noise_want: f64 = sigma;
 
-    assert!(noise_have <= noise_want + 0.2);
+    assert!(
+        noise_have <= noise_want + 0.2,
+        "{} <= {}",
+        noise_have,
+        noise_want + 0.2
+    );
 }
 
 pub(crate) fn test_encrypt_zero_sk<B: Backend>(module: &Module<B>, basek: usize, k_ct: usize, sigma: f64, rank: usize)

core/src/lwe/test_fft64/conversion.rs

@@ -163,7 +163,7 @@ where
 
     let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
         LWEToGLWESwitchingKey::encrypt_sk_scratch_space(module, basek, k_ksk, rank)
-            | GLWECiphertext::from_lwe_scratch_space(module, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
+            | LWECiphertext::from_glwe_scratch_space(module, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
             | GLWECiphertext::decrypt_scratch_space(module, basek, k_glwe_ct),
     );