Merge branch 'dev_trait' into dev_trait_practice

poulpy-core/src/conversion/glwe_to_lwe.rs

@@ -1,100 +1,125 @@
 use poulpy_hal::{
-    api::{
+    api::ModuleN,
-        VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply, VecZnxDftBytesOf,
+    layouts::{Backend, DataMut, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
-        VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft, VmpApplyDftToDftAdd,
-        VmpApplyDftToDftTmpBytes,
-    },
-    layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
 };
 
-use crate::layouts::{GGLWEInfos, GLWE, GLWEInfos, GLWELayout, LWE, LWEInfos, Rank, prepared::GLWEToLWESwitchingKeyPrepared};
+use crate::{
+    GLWEKeyswitch, ScratchTakeCore,
+    layouts::{
+        GGLWEInfos, GLWE, GLWEAlloc, GLWEInfos, GLWELayout, GLWEToRef, LWE, LWEInfos, LWEToMut, Rank,
+        prepared::{LWEToGLWESwitchingKeyPrepared, LWEToGLWESwitchingKeyPreparedToRef},
+    },
+};
 
-impl LWE<Vec<u8>> {
+pub trait LWESampleExtract
-    pub fn from_glwe_tmp_bytes<B: Backend, OUT, IN, KEY>(
+where
-        module: &Module<B>,
+    Self: ModuleN,
-        lwe_infos: &OUT,
+{
-        glwe_infos: &IN,
+    fn lwe_sample_extract<R, A>(&self, res: &mut R, a: &A)
-        key_infos: &KEY,
-    ) -> usize
     where
-        OUT: LWEInfos,
+        R: LWEToMut,
-        IN: GLWEInfos,
+        A: GLWEToRef,
-        KEY: GGLWEInfos,
-        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
     {
-        let glwe_layout: GLWELayout = GLWELayout {
+        let res: &mut LWE<&mut [u8]> = &mut res.to_mut();
-            n: module.n().into(),
+        let a: &GLWE<&[u8]> = &a.to_ref();
+
+        assert!(res.n() <= a.n());
+        assert_eq!(a.n(), self.n() as u32);
+        assert!(res.base2k() == a.base2k());
+
+        let min_size: usize = res.size().min(a.size());
+        let n: usize = res.n().into();
+
+        res.data.zero();
+        (0..min_size).for_each(|i| {
+            let data_lwe: &mut [i64] = res.data.at_mut(0, i);
+            data_lwe[0] = a.data.at(0, i)[0];
+            data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]);
+        });
+    }
+}
+
+impl<BE: Backend> LWESampleExtract for Module<BE> where Self: ModuleN {}
+impl<BE: Backend> LWEFromGLWE<BE> for Module<BE> where Self: GLWEKeyswitch<BE> + GLWEAlloc + LWESampleExtract {}
+
+pub trait LWEFromGLWE<BE: Backend>
+where
+    Self: GLWEKeyswitch<BE> + GLWEAlloc + LWESampleExtract,
+{
+    fn lwe_from_glwe_tmp_bytes<R, A, K>(&self, lwe_infos: &R, glwe_infos: &A, key_infos: &K) -> usize
+    where
+        R: LWEInfos,
+        A: GLWEInfos,
+        K: GGLWEInfos,
+    {
+        let res_infos: GLWELayout = GLWELayout {
+            n: self.n().into(),
             base2k: lwe_infos.base2k(),
             k: lwe_infos.k(),
             rank: Rank(1),
         };
 
-        GLWE::bytes_of(
+        self.bytes_of_glwe(lwe_infos.base2k(), lwe_infos.k(), 1u32.into())
-            module.n().into(),
+            + self.glwe_keyswitch_tmp_bytes(&res_infos, glwe_infos, key_infos)
-            lwe_infos.base2k(),
-            lwe_infos.k(),
-            1u32.into(),
-        ) + GLWE::keyswitch_tmp_bytes(module, &glwe_layout, glwe_infos, key_infos)
     }
-}
 
-impl<DLwe: DataMut> LWE<DLwe> {
+    fn lwe_from_glwe<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
-    pub fn sample_extract<DGlwe: DataRef>(&mut self, a: &GLWE<DGlwe>) {
+    where
-        #[cfg(debug_assertions)]
+        R: LWEToMut,
+        A: GLWEToRef,
+        K: LWEToGLWESwitchingKeyPreparedToRef<BE> + GGLWEInfos,
+        Scratch<BE>: ScratchTakeCore<BE>,
     {
-            assert!(self.n() <= a.n());
+        let res: &mut LWE<&mut [u8]> = &mut res.to_mut();
-            assert!(self.base2k() == a.base2k());
+        let a: &GLWE<&[u8]> = &a.to_ref();
-        }
+        let key: &LWEToGLWESwitchingKeyPrepared<&[u8], BE> = &key.to_ref();
 
-        let min_size: usize = self.size().min(a.size());
+        assert_eq!(a.n(), self.n() as u32);
-        let n: usize = self.n().into();
+        assert_eq!(key.n(), self.n() as u32);
-
+        assert!(res.n() <= self.n() as u32);
-        self.data.zero();
-        (0..min_size).for_each(|i| {
-            let data_lwe: &mut [i64] = self.data.at_mut(0, i);
-            data_lwe[0] = a.data.at(0, i)[0];
-            data_lwe[1..].copy_from_slice(&a.data.at(1, i)[..n]);
-        });
-    }
-
-    pub fn from_glwe<DGlwe, DKs, B: Backend>(
-        &mut self,
-        module: &Module<B>,
-        a: &GLWE<DGlwe>,
-        ks: &GLWEToLWESwitchingKeyPrepared<DKs, B>,
-        scratch: &mut Scratch<B>,
-    ) where
-        DGlwe: DataRef,
-        DKs: DataRef,
-        Module<B>: VecZnxDftBytesOf
-            + VmpApplyDftToDftTmpBytes
-            + VecZnxBigNormalizeTmpBytes
-            + VmpApplyDftToDft<B>
-            + VmpApplyDftToDftAdd<B>
-            + VecZnxDftApply<B>
-            + VecZnxIdftApplyConsume<B>
-            + VecZnxBigAddSmallInplace<B>
-            + VecZnxBigNormalize<B>
-            + VecZnxNormalize<B>
-            + VecZnxNormalizeTmpBytes,
-        Scratch<B>:,
-    {
-        #[cfg(debug_assertions)]
-        {
-            assert_eq!(a.n(), module.n() as u32);
-            assert_eq!(ks.n(), module.n() as u32);
-            assert!(self.n() <= module.n() as u32);
-        }
 
         let glwe_layout: GLWELayout = GLWELayout {
-            n: module.n().into(),
+            n: self.n().into(),
-            base2k: self.base2k(),
+            base2k: res.base2k(),
-            k: self.k(),
+            k: res.k(),
             rank: Rank(1),
         };
 
-        let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(&glwe_layout);
+        let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(self, &glwe_layout);
-        tmp_glwe.keyswitch(module, a, &ks.0, scratch_1);
+        self.glwe_keyswitch(&mut tmp_glwe, a, &key.0, scratch_1);
-        self.sample_extract(&tmp_glwe);
+        self.lwe_sample_extract(res, &tmp_glwe);
+    }
+}
+
+impl LWE<Vec<u8>> {
+    pub fn from_glwe_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, lwe_infos: &R, glwe_infos: &A, key_infos: &K) -> usize
+    where
+        R: LWEInfos,
+        A: GLWEInfos,
+        K: GGLWEInfos,
+        M: LWEFromGLWE<BE>,
+    {
+        module.lwe_from_glwe_tmp_bytes(lwe_infos, glwe_infos, key_infos)
+    }
+}
+
+impl<D: DataMut> LWE<D> {
+    pub fn sample_extract<A, M>(&mut self, module: &M, a: &A)
+    where
+        A: GLWEToRef,
+        M: LWESampleExtract,
+    {
+        module.lwe_sample_extract(self, a);
+    }
+
+    pub fn from_glwe<R, A, K, M, BE: Backend>(&self, module: &M, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
+    where
+        R: LWEToMut,
+        A: GLWEToRef,
+        K: LWEToGLWESwitchingKeyPreparedToRef<BE> + GGLWEInfos,
+        M: LWEFromGLWE<BE>,
+        Scratch<BE>: ScratchTakeCore<BE>,
+    {
+        module.lwe_from_glwe(res, a, key, scratch);
     }
 }

poulpy-core/src/conversion/lwe_to_glwe.rs

@@ -1,79 +1,67 @@
 use poulpy_hal::{
-    api::{
+    api::ScratchTakeBasic,
-        ScratchAvailable, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes, VecZnxDftApply,
+    layouts::{Backend, DataMut, Module, Scratch, VecZnx, ZnxView, ZnxViewMut, ZnxZero},
-        VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
-        VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
-    },
-    layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx, ZnxView, ZnxViewMut, ZnxZero},
 };
 
-use crate::layouts::{GGLWEInfos, GLWE, GLWEInfos, GLWELayout, LWE, LWEInfos, prepared::LWEToGLWESwitchingKeyPrepared};
+use crate::{
+    GLWEKeyswitch, ScratchTakeCore,
+    layouts::{
+        GGLWEInfos, GLWE, GLWEAlloc, GLWEInfos, GLWELayout, GLWEToMut, LWE, LWEInfos, LWEToRef,
+        prepared::{LWEToGLWESwitchingKeyPrepared, LWEToGLWESwitchingKeyPreparedToRef},
+    },
+};
 
-impl GLWE<Vec<u8>> {
+impl<BE: Backend> GLWEFromLWE<BE> for Module<BE> where Self: GLWEKeyswitch<BE> + GLWEAlloc {}
-    pub fn from_lwe_tmp_bytes<B: Backend, OUT, IN, KEY>(
+
-        module: &Module<B>,
+pub trait GLWEFromLWE<BE: Backend>
-        glwe_infos: &OUT,
+where
-        lwe_infos: &IN,
+    Self: GLWEKeyswitch<BE> + GLWEAlloc,
-        key_infos: &KEY,
+{
-    ) -> usize
+    fn glwe_from_lwe_tmp_bytes<R, A, K>(&self, glwe_infos: &R, lwe_infos: &A, key_infos: &K) -> usize
     where
-        OUT: GLWEInfos,
+        R: GLWEInfos,
-        IN: LWEInfos,
+        A: LWEInfos,
-        KEY: GGLWEInfos,
+        K: GGLWEInfos,
-        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
     {
-        let ct: usize = GLWE::bytes_of(
+        let ct: usize = self.bytes_of_glwe(
-            module.n().into(),
             key_infos.base2k(),
             lwe_infos.k().max(glwe_infos.k()),
             1u32.into(),
         );
-        let ks: usize = GLWE::keyswitch_inplace_tmp_bytes(module, glwe_infos, key_infos);
+
+        let ks: usize = self.glwe_keyswitch_tmp_bytes(glwe_infos, glwe_infos, key_infos);
         if lwe_infos.base2k() == key_infos.base2k() {
             ct + ks
         } else {
-            let a_conv = VecZnx::bytes_of(module.n(), 1, lwe_infos.size()) + module.vec_znx_normalize_tmp_bytes();
+            let a_conv = VecZnx::bytes_of(self.n(), 1, lwe_infos.size()) + self.vec_znx_normalize_tmp_bytes();
             ct + a_conv + ks
         }
     }
-}
 
-impl<D: DataMut> GLWE<D> {
+    fn glwe_from_lwe<R, A, K>(&self, res: &mut R, lwe: &A, ksk: &K, scratch: &mut Scratch<BE>)
-    pub fn from_lwe<DLwe, DKsk, B: Backend>(
+    where
-        &mut self,
+        R: GLWEToMut,
-        module: &Module<B>,
+        A: LWEToRef,
-        lwe: &LWE<DLwe>,
+        K: LWEToGLWESwitchingKeyPreparedToRef<BE>,
-        ksk: &LWEToGLWESwitchingKeyPrepared<DKsk, B>,
+        Scratch<BE>: ScratchTakeCore<BE>,
-        scratch: &mut Scratch<B>,
-    ) where
-        DLwe: DataRef,
-        DKsk: DataRef,
-        Module<B>: VecZnxDftBytesOf
-            + VmpApplyDftToDftTmpBytes
-            + VecZnxBigNormalizeTmpBytes
-            + VmpApplyDftToDft<B>
-            + VmpApplyDftToDftAdd<B>
-            + VecZnxDftApply<B>
-            + VecZnxIdftApplyConsume<B>
-            + VecZnxBigAddSmallInplace<B>
-            + VecZnxBigNormalize<B>
-            + VecZnxNormalize<B>
-            + VecZnxNormalizeTmpBytes,
-        Scratch<B>: ScratchAvailable,
     {
-        #[cfg(debug_assertions)]
+        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
-        {
+        let lwe: &LWE<&[u8]> = &lwe.to_ref();
-            assert_eq!(self.n(), module.n() as u32);
+        let ksk: &LWEToGLWESwitchingKeyPrepared<&[u8], BE> = &ksk.to_ref();
-            assert_eq!(ksk.n(), module.n() as u32);
-            assert!(lwe.n() <= module.n() as u32);
-        }
 
-        let (mut glwe, scratch_1) = scratch.take_glwe_ct(&GLWELayout {
+        assert_eq!(res.n(), self.n() as u32);
+        assert_eq!(ksk.n(), self.n() as u32);
+        assert!(lwe.n() <= self.n() as u32);
+
+        let (mut glwe, scratch_1) = scratch.take_glwe_ct(
+            self,
+            &GLWELayout {
                 n: ksk.n(),
                 base2k: ksk.base2k(),
                 k: lwe.k(),
                 rank: 1u32.into(),
-        });
+            },
+        );
         glwe.data.zero();
 
         let n_lwe: usize = lwe.n().into();
@@ -85,14 +73,14 @@ impl<D: DataMut> GLWE<D> {
                 glwe.data.at_mut(1, i)[..n_lwe].copy_from_slice(&data_lwe[1..]);
             }
         } else {
-            let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(module.n(), 1, lwe.size());
+            let (mut a_conv, scratch_2) = scratch_1.take_vec_znx(self, 1, lwe.size());
             a_conv.zero();
             for j in 0..lwe.size() {
                 let data_lwe: &[i64] = lwe.data.at(0, j);
                 a_conv.at_mut(0, j)[0] = data_lwe[0]
             }
 
-            module.vec_znx_normalize(
+            self.vec_znx_normalize(
                 ksk.base2k().into(),
                 &mut glwe.data,
                 0,
@@ -108,7 +96,7 @@ impl<D: DataMut> GLWE<D> {
                 a_conv.at_mut(0, j)[..n_lwe].copy_from_slice(&data_lwe[1..]);
             }
 
-            module.vec_znx_normalize(
+            self.vec_znx_normalize(
                 ksk.base2k().into(),
                 &mut glwe.data,
                 1,
@@ -119,6 +107,30 @@ impl<D: DataMut> GLWE<D> {
             );
         }
 
-        self.keyswitch(module, &glwe, &ksk.0, scratch_1);
+        self.glwe_keyswitch(res, &glwe, &ksk.0, scratch_1);
+    }
+}
+
+impl GLWE<Vec<u8>> {
+    pub fn from_lwe_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, glwe_infos: &R, lwe_infos: &A, key_infos: &K) -> usize
+    where
+        R: GLWEInfos,
+        A: LWEInfos,
+        K: GGLWEInfos,
+        M: GLWEFromLWE<BE>,
+    {
+        module.glwe_from_lwe_tmp_bytes(glwe_infos, lwe_infos, key_infos)
+    }
+}
+
+impl<D: DataMut> GLWE<D> {
+    pub fn from_lwe<A, K, M, BE: Backend>(&mut self, module: &M, lwe: &A, ksk: &K, scratch: &mut Scratch<BE>)
+    where
+        M: GLWEFromLWE<BE>,
+        A: LWEToRef,
+        K: LWEToGLWESwitchingKeyPreparedToRef<BE>,
+        Scratch<BE>: ScratchTakeCore<BE>,
+    {
+        module.glwe_from_lwe(self, lwe, ksk, scratch);
     }
 }

poulpy-core/src/conversion/mod.rs

@@ -3,3 +3,5 @@ mod glwe_to_lwe;
 mod lwe_to_glwe;
 
 pub use gglwe_to_ggsw::*;
+pub use glwe_to_lwe::*;
+pub use lwe_to_glwe::*;

poulpy-core/src/glwe_packing.rs

@@ -1,12 +1,13 @@
 use std::collections::HashMap;
 
 use poulpy_hal::{
-    api::{ModuleLogN, VecZnxCopy, VecZnxRotateInplace},
+    api::ModuleLogN,
-    layouts::{Backend, DataMut, DataRef, GaloisElement, Module, Scratch},
+    layouts::{Backend, GaloisElement, Module, Scratch},
 };
 
 use crate::{
-    GLWEAdd, GLWEAutomorphism, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore,
+    GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore,
+    glwe_trace::GLWETrace,
     layouts::{
         GGLWEInfos, GLWE, GLWEAlloc, GLWEInfos, GLWEToMut, GLWEToRef, LWEInfos,
         prepared::{AutomorphismKeyPreparedToRef, GetAutomorphismGaloisElement},
@@ -40,10 +41,10 @@ impl Accumulator {
     /// * `base2k`: base 2 logarithm of the GLWE ciphertext in memory digit representation.
     /// * `k`: base 2 precision of the GLWE ciphertext precision over the Torus.
     /// * `rank`: rank of the GLWE ciphertext.
-    pub fn alloc<A, M>(module: &M, infos: &A) -> Self
+    pub fn alloc<A, M, BE: Backend>(module: &M, infos: &A) -> Self
     where
         A: GLWEInfos,
-        M: GLWEAlloc,
+        M: GLWEPacking<BE>,
     {
         Self {
             data: GLWE::alloc_from_infos(module, infos),
@@ -64,15 +65,15 @@ impl GLWEPacker {
     ///   and N GLWE ciphertext can be packed. With `log_batch=2` all coefficients
     ///   which are multiples of X^{N/4} are packed. Meaning that N/4 ciphertexts
     ///   can be packed.
-    pub fn new<A, M>(module: &M, infos: &A, log_batch: usize) -> Self
+    pub fn alloc<A, M, BE: Backend>(module: &M, infos: &A, log_batch: usize) -> Self
     where
         A: GLWEInfos,
-        M: GLWEAlloc,
+        M: GLWEPacking<BE>,
     {
         let mut accumulators: Vec<Accumulator> = Vec::<Accumulator>::new();
         let log_n: usize = infos.n().log2();
         (0..log_n - log_batch).for_each(|_| accumulators.push(Accumulator::alloc(module, infos)));
-        Self {
+        GLWEPacker {
             accumulators,
             log_batch,
             counter: 0,
@@ -93,13 +94,19 @@ impl GLWEPacker {
     where
         R: GLWEInfos,
         K: GGLWEInfos,
-        M: GLWEAlloc + GLWEAutomorphism<BE>,
+        M: GLWEPacking<BE>,
     {
-        pack_core_tmp_bytes(module, res_infos, key_infos)
+        module.bytes_of_glwe_from_infos(res_infos)
+            + module
+                .glwe_rsh_tmp_byte()
+                .max(module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos))
     }
 
-    pub fn galois_elements<B: Backend>(module: &Module<B>) -> Vec<i64> {
+    pub fn galois_elements<M, BE: Backend>(module: &M) -> Vec<i64>
-        GLWE::trace_galois_elements(module)
+    where
+        M: GLWETrace<BE>,
+    {
+        module.glwe_trace_galois_elements()
     }
 
     /// Adds a GLWE ciphertext to the [GLWEPacker].
@@ -111,11 +118,11 @@ impl GLWEPacker {
     /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext.
     /// * `auto_keys`: a [HashMap] containing the [AutomorphismKeyExec]s.
     /// * `scratch`: scratch space of size at least [Self::tmp_bytes].
-    pub fn add<A, K, M, BE: Backend>(&mut self, module: &M, a: Option<&A>, auto_keys: &HashMap<i64, K>, scratch: &mut Scratch<B>)
+    pub fn add<A, K, M, BE: Backend>(&mut self, module: &M, a: Option<&A>, auto_keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
     where
-        A: GLWEToRef,
+        A: GLWEToRef + GLWEInfos,
-        K: AutomorphismKeyPreparedToRef<BE>,
+        K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
-        M: GLWEAutomorphism<BE>,
+        M: GLWEPacking<BE>,
         Scratch<BE>: ScratchTakeCore<BE>,
     {
         assert!(
@@ -136,14 +143,15 @@ impl GLWEPacker {
     }
 
     /// Flush result to`res`.
-    pub fn flush<Data: DataMut, B: Backend>(&mut self, module: &Module<B>, res: &mut GLWE<Data>)
+    pub fn flush<R, M, BE: Backend>(&mut self, module: &M, res: &mut R)
     where
-        Module<B>: VecZnxCopy,
+        R: GLWEToMut,
+        M: GLWEPacking<BE>,
     {
         assert!(self.counter as u32 == self.accumulators[0].data.n());
         // Copy result GLWE into res GLWE
-        res.copy(
+        module.glwe_copy(
-            module,
+            res,
             &self.accumulators[module.log_n() - self.log_batch - 1].data,
         );
 
@@ -151,13 +159,76 @@ impl GLWEPacker {
     }
 }
 
-fn pack_core_tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
+impl<BE: Backend> GLWEPacking<BE> for Module<BE> where
-where
+    Self: GLWEAutomorphism<BE>
-    R: GLWEInfos,
+        + GaloisElement
-    K: GGLWEInfos,
+        + ModuleLogN
-    M: GLWEAlloc + GLWEAutomorphism<BE>,
+        + GLWERotate<BE>
+        + GLWESub
+        + GLWEShift<BE>
+        + GLWEAdd
+        + GLWENormalize<BE>
+        + GLWECopy
+        + GLWEAlloc
 {
-    combine_tmp_bytes(module, res_infos, key_infos)
+}
+
+pub trait GLWEPacking<BE: Backend>
+where
+    Self: GLWEAutomorphism<BE>
+        + GaloisElement
+        + ModuleLogN
+        + GLWERotate<BE>
+        + GLWESub
+        + GLWEShift<BE>
+        + GLWEAdd
+        + GLWENormalize<BE>
+        + GLWECopy
+        + GLWEAlloc,
+{
+    /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
+    /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
+    fn glwe_pack<R, K>(
+        &self,
+        cts: &mut HashMap<usize, &mut R>,
+        log_gap_out: usize,
+        keys: &HashMap<i64, K>,
+        scratch: &mut Scratch<BE>,
+    ) where
+        R: GLWEToMut + GLWEToRef + GLWEInfos,
+        K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
+        Scratch<BE>: ScratchTakeCore<BE>,
+    {
+        #[cfg(debug_assertions)]
+        {
+            assert!(*cts.keys().max().unwrap() < self.n())
+        }
+
+        let log_n: usize = self.log_n();
+
+        for i in 0..(log_n - log_gap_out) {
+            let t: usize = (1 << log_n).min(1 << (log_n - 1 - i));
+
+            let key: &K = if i == 0 {
+                keys.get(&-1).unwrap()
+            } else {
+                keys.get(&self.galois_element(1 << (i - 1))).unwrap()
+            };
+
+            for j in 0..t {
+                let mut a: Option<&mut R> = cts.remove(&j);
+                let mut b: Option<&mut R> = cts.remove(&(j + t));
+
+                pack_internal(self, &mut a, &mut b, i, key, scratch);
+
+                if let Some(a) = a {
+                    cts.insert(j, a);
+                } else if let Some(b) = b {
+                    cts.insert(j, b);
+                }
+            }
+        }
+    }
 }
 
 fn pack_core<A, K, M, BE: Backend>(
@@ -169,8 +240,16 @@ fn pack_core<A, K, M, BE: Backend>(
     scratch: &mut Scratch<BE>,
 ) where
     A: GLWEToRef + GLWEInfos,
-    K: AutomorphismKeyPreparedToRef<BE>,
+    K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
-    M: GLWEAutomorphism<BE> + ModuleLogN + VecZnxCopy,
+    M: ModuleLogN
+        + GLWEAutomorphism<BE>
+        + GaloisElement
+        + GLWERotate<BE>
+        + GLWESub
+        + GLWEShift<BE>
+        + GLWEAdd
+        + GLWENormalize<BE>
+        + GLWECopy,
     Scratch<BE>: ScratchTakeCore<BE>,
 {
     let log_n: usize = module.log_n();
@@ -188,7 +267,7 @@ fn pack_core<A, K, M, BE: Backend>(
 
         // No previous value -> copies and sets flags accordingly
         if let Some(a_ref) = a {
-            acc_mut_ref.data.copy(module, a_ref);
+            module.glwe_copy(&mut acc_mut_ref.data, a_ref);
             acc_mut_ref.value = true
         } else {
             acc_mut_ref.value = false
@@ -222,16 +301,6 @@ fn pack_core<A, K, M, BE: Backend>(
     }
 }
 
-fn combine_tmp_bytes<R, K, M, BE: Backend>(module: &M, res_infos: &R, key_infos: &K) -> usize
-where
-    R: GLWEInfos,
-    K: GGLWEInfos,
-    M: GLWEAlloc + GLWEAutomorphism<BE>,
-{
-    GLWE::bytes_of_from_infos(module, res_infos)
-        + (GLWE::rsh_tmp_bytes(module.n()) | module.glwe_automorphism_tmp_bytes(res_infos, res_infos, key_infos))
-}
-
 /// [combine] merges two ciphertexts together.
 fn combine<B, M, K, BE: Backend>(
     module: &M,
@@ -242,8 +311,9 @@ fn combine<B, M, K, BE: Backend>(
     scratch: &mut Scratch<BE>,
 ) where
     B: GLWEToRef + GLWEInfos,
-    K: AutomorphismKeyPreparedToRef<BE>,
+    M: GLWEAutomorphism<BE> + GaloisElement + GLWERotate<BE> + GLWESub + GLWEShift<BE> + GLWEAdd + GLWENormalize<BE>,
-    M: GLWEAutomorphism<BE> + GaloisElement + VecZnxRotateInplace<BE>,
+    B: GLWEToRef + GLWEInfos,
+    K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
     Scratch<BE>: ScratchTakeCore<BE>,
 {
     let log_n: usize = acc.data.n().log2();
@@ -272,50 +342,50 @@ fn combine<B, M, K, BE: Backend>(
             let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(module, a);
 
             // a = a * X^-t
-            a.rotate_inplace(module, -t, scratch_1);
+            module.glwe_rotate_inplace(-t, a, scratch_1);
 
             // tmp_b = a * X^-t - b
-            tmp_b.sub(module, a, b);
+            module.glwe_sub(&mut tmp_b, a, b);
-            tmp_b.rsh(module, 1, scratch_1);
+            module.glwe_rsh(1, &mut tmp_b, scratch_1);
 
             // a = a * X^-t + b
-            a.add_inplace(module, b);
+            module.glwe_add_inplace(a, b);
-            a.rsh(module, 1, scratch_1);
+            module.glwe_rsh(1, a, scratch_1);
 
-            tmp_b.normalize_inplace(module, scratch_1);
+            module.glwe_normalize_inplace(&mut tmp_b, scratch_1);
 
             // tmp_b = phi(a * X^-t - b)
-            if let Some(key) = auto_keys.get(&gal_el) {
+            if let Some(auto_key) = auto_keys.get(&gal_el) {
-                tmp_b.automorphism_inplace(module, key, scratch_1);
+                module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1);
             } else {
                 panic!("auto_key[{gal_el}] not found");
             }
 
             // a = a * X^-t + b - phi(a * X^-t - b)
-            a.sub_inplace_ab(module, &tmp_b);
+            module.glwe_sub_inplace(a, &tmp_b);
-            a.normalize_inplace(module, scratch_1);
+            module.glwe_normalize_inplace(a, scratch_1);
 
             // a = a + b * X^t - phi(a * X^-t - b) * X^t
             //   = a + b * X^t - phi(a * X^-t - b) * - phi(X^t)
             //   = a + b * X^t + phi(a - b * X^t)
-            a.rotate_inplace(module, t, scratch_1);
+            module.glwe_rotate_inplace(t, a, scratch_1);
         } else {
-            a.rsh(module, 1, scratch);
+            module.glwe_rsh(1, a, scratch);
             // a = a + phi(a)
-            if let Some(key) = auto_keys.get(&gal_el) {
+            if let Some(auto_key) = auto_keys.get(&gal_el) {
-                a.automorphism_add_inplace(module, key, scratch);
+                module.glwe_automorphism_add_inplace(a, auto_key, scratch);
             } else {
                 panic!("auto_key[{gal_el}] not found");
             }
         }
     } else if let Some(b) = b {
-        let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(a);
+        let (mut tmp_b, scratch_1) = scratch.take_glwe_ct(module, a);
-        tmp_b.rotate(module, 1 << (log_n - i - 1), b);
+        module.glwe_rotate(t, &mut tmp_b, b);
-        tmp_b.rsh(module, 1, scratch_1);
+        module.glwe_rsh(1, &mut tmp_b, scratch_1);
 
         // a = (b* X^t - phi(b* X^t))
-        if let Some(key) = auto_keys.get(&gal_el) {
+        if let Some(auto_key) = auto_keys.get(&gal_el) {
-            a.automorphism_sub_negate(module, &tmp_b, key, scratch_1);
+            module.glwe_automorphism_sub_negate(a, &tmp_b, auto_key, scratch_1);
         } else {
             panic!("auto_key[{gal_el}] not found");
         }
@@ -324,62 +394,6 @@ fn combine<B, M, K, BE: Backend>(
     }
 }
 
-pub trait GLWEPacking<BE: Backend>
-where
-    Self: GLWEAutomorphism<BE>
-        + GaloisElement
-        + ModuleLogN
-        + GLWERotate<BE>
-        + GLWESub
-        + GLWEShift<BE>
-        + GLWEAdd
-        + GLWENormalize<BE>,
-{
-    /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
-    /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
-    fn glwe_pack<R, K>(
-        &self,
-        cts: &mut HashMap<usize, &mut R>,
-        log_gap_out: usize,
-        keys: &HashMap<i64, K>,
-        scratch: &mut Scratch<BE>,
-    ) where
-        R: GLWEToMut + GLWEToRef + GLWEInfos,
-        K: AutomorphismKeyPreparedToRef<BE> + GetAutomorphismGaloisElement,
-        Scratch<BE>: ScratchTakeCore<BE>,
-    {
-        #[cfg(debug_assertions)]
-        {
-            assert!(*cts.keys().max().unwrap() < self.n())
-        }
-
-        let log_n: usize = self.log_n();
-
-        for i in 0..(log_n - log_gap_out){
-            let t: usize = (1 << log_n).min(1 << (log_n - 1 - i));
-
-            let key: &K = if i == 0 {
-                keys.get(&-1).unwrap()
-            } else {
-                keys.get(&self.galois_element(1 << (i - 1))).unwrap()
-            };
-
-            for j in 0..t{
-                let mut a: Option<&mut R> = cts.remove(&j);
-                let mut b: Option<&mut R> = cts.remove(&(j + t));
-
-                pack_internal(self, &mut a, &mut b, i, key, scratch);
-
-                if let Some(a) = a {
-                    cts.insert(j, a);
-                } else if let Some(b) = b {
-                    cts.insert(j, b);
-                }
-            };
-        };
-    }
-}
-
 #[allow(clippy::too_many_arguments)]
 fn pack_internal<M, A, B, K, BE: Backend>(
     module: &M,

poulpy-core/src/keyswitching/lwe_ct.rs

@@ -4,7 +4,7 @@ use poulpy_hal::{
 };
 
 use crate::{
-    ScratchTakeCore,
+    LWESampleExtract, ScratchTakeCore,
     keyswitching::glwe_ct::GLWEKeyswitch,
     layouts::{
         GGLWEInfos, GLWE, GLWEAlloc, GLWELayout, LWE, LWEInfos, LWEToMut, LWEToRef, Rank, TorusPrecision,
@@ -40,7 +40,7 @@ impl<BE: Backend> LWEKeySwitch<BE> for Module<BE> where Self: LWEKeySwitch<BE> {
 
 pub trait LWEKeySwitch<BE: Backend>
 where
-    Self: GLWEKeyswitch<BE> + GLWEAlloc,
+    Self: GLWEKeyswitch<BE> + GLWEAlloc + LWESampleExtract,
 {
     fn lwe_keyswitch_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
     where
@@ -121,6 +121,6 @@ where
         }
 
         self.glwe_keyswitch(&mut glwe_out, &glwe_in, &ksk.0, scratch_1);
-        res.sample_extract(&glwe_out);
+        self.lwe_sample_extract(res, &glwe_out);
     }
 }

poulpy-core/src/operations/glwe.rs

@@ -277,6 +277,10 @@ pub trait GLWEShift<BE: Backend>
 where
     Self: ModuleN + VecZnxRshInplace<BE>,
 {
+    fn glwe_rsh_tmp_byte(&self) -> usize {
+        VecZnx::rsh_tmp_bytes(self.n())
+    }
+
     fn glwe_rsh<R>(&self, k: usize, res: &mut R, scratch: &mut Scratch<BE>)
     where
         R: GLWEToMut,
@@ -291,8 +295,11 @@ where
 }
 
 impl GLWE<Vec<u8>> {
-    pub fn rsh_tmp_bytes(n: usize) -> usize {
+    pub fn rsh_tmp_bytes<M, BE: Backend>(module: &M) -> usize
-        VecZnx::rsh_tmp_bytes(n)
+    where
+        M: GLWEShift<BE>,
+    {
+        module.glwe_rsh_tmp_byte()
     }
 }