wip

poulpy-core/src/layouts/gglwe_ksk.rs

@@ -67,6 +67,36 @@ pub struct GLWESwitchingKey<D: Data> {
     pub(crate) sk_out_n: usize, // Degree of sk_out
 }
 
+pub(crate) trait GLWESwitchingKeySetMetaData {
+    fn set_sk_in_n(&mut self, sk_in_n: usize);
+    fn set_sk_out_n(&mut self, sk_out_n: usize);
+}
+
+impl<D: DataMut> GLWESwitchingKeySetMetaData for GLWESwitchingKey<D> {
+    fn set_sk_in_n(&mut self, sk_in_n: usize) {
+        self.sk_in_n = sk_in_n
+    }
+
+    fn set_sk_out_n(&mut self, sk_out_n: usize) {
+        self.sk_out_n = sk_out_n
+    }
+}
+
+pub(crate) trait GLWESwtichingKeyGetMetaData {
+    fn sk_in_n(&self) -> usize;
+    fn sk_out_n(&self) -> usize;
+}
+
+impl<D: DataRef> GLWESwtichingKeyGetMetaData for GLWESwitchingKey<D> {
+    fn sk_in_n(&self) -> usize {
+        self.sk_in_n
+    }
+
+    fn sk_out_n(&self) -> usize {
+        self.sk_out_n
+    }
+}
+
 impl<D: Data> LWEInfos for GLWESwitchingKey<D> {
     fn n(&self) -> Degree {
         self.key.n()

poulpy-core/src/layouts/lwe_ksk.rs

@@ -193,11 +193,11 @@ impl<D: DataRef> WriterTo for LWESwitchingKey<D> {
     }
 }
 
-pub trait LWEToLWEKeyToRef {
+pub trait LWESwitchingKeyToRef {
     fn to_ref(&self) -> LWESwitchingKey<&[u8]>;
 }
 
-impl<D: DataRef> LWEToLWEKeyToRef for LWESwitchingKey<D>
+impl<D: DataRef> LWESwitchingKeyToRef for LWESwitchingKey<D>
 where
     GLWESwitchingKey<D>: GLWESwitchingKeyToRef,
 {
@@ -206,11 +206,11 @@ where
     }
 }
 
-pub trait LWEToLWEKeyToMut {
+pub trait LWESwitchingKeyToMut {
     fn to_mut(&mut self) -> LWESwitchingKey<&mut [u8]>;
 }
 
-impl<D: DataMut> LWEToLWEKeyToMut for LWESwitchingKey<D>
+impl<D: DataMut> LWESwitchingKeyToMut for LWESwitchingKey<D>
 where
     GLWESwitchingKey<D>: GLWESwitchingKeyToMut,
 {

poulpy-core/src/layouts/prepared/gglwe_atk.rs

@@ -1,12 +1,11 @@
-use poulpy_hal::{
+use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
-    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
-    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
-};
 
 use crate::layouts::{
     AutomorphismKeyToRef, Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision,
     prepared::{
-        GLWESwitchingKeyPrepareTmpBytes, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToMut, GLWESwitchingKeyPreparedToRef,
+        GLWESwitchingKeyPrepare, GLWESwitchingKeyPrepareTmpBytes, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedAlloc,
+        GLWESwitchingKeyPreparedAllocBytes, GLWESwitchingKeyPreparedAllocBytesFromInfos, GLWESwitchingKeyPreparedAllocFromInfos,
+        GLWESwitchingKeyPreparedToMut, GLWESwitchingKeyPreparedToRef,
     },
 };
 
@@ -16,12 +15,6 @@ pub struct AutomorphismKeyPrepared<D: Data, B: Backend> {
     pub(crate) p: i64,
 }
 
-impl<D: Data, B: Backend> AutomorphismKeyPrepared<D, B> {
-    pub fn p(&self) -> i64 {
-        self.p
-    }
-}
-
 impl<D: Data, B: Backend> LWEInfos for AutomorphismKeyPrepared<D, B> {
     fn n(&self) -> Degree {
         self.key.n()
@@ -40,11 +33,21 @@ impl<D: Data, B: Backend> LWEInfos for AutomorphismKeyPrepared<D, B> {
     }
 }
 
-pub trait SetP {
+pub trait GetAutomorphismGaloisElement {
+    fn p(&self) -> i64;
+}
+
+impl<D: Data, B: Backend> GetAutomorphismGaloisElement for AutomorphismKeyPrepared<D, B> {
+    fn p(&self) -> i64 {
+        self.p
+    }
+}
+
+pub trait SetAutomorphismGaloisElement {
     fn set_p(&mut self, p: i64);
 }
 
-impl<D: Data, B: Backend> SetP for AutomorphismKeyPrepared<D, B> {
+impl<D: Data, B: Backend> SetAutomorphismGaloisElement for AutomorphismKeyPrepared<D, B> {
     fn set_p(&mut self, p: i64) {
         self.p = p
     }
@@ -74,56 +77,146 @@ impl<D: Data, B: Backend> GGLWEInfos for AutomorphismKeyPrepared<D, B> {
     }
 }
 
-impl<B: Backend> AutomorphismKeyPrepared<Vec<u8>, B> {
+pub trait AutomorphismKeyPreparedAlloc<B: Backend> {
-    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
+    fn automorphism_key_prepared_alloc(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> AutomorphismKeyPrepared<Vec<u8>, B>;
+}
+
+impl<B: Backend> AutomorphismKeyPreparedAlloc<B> for Module<B>
+where
+    Module<B>: GLWESwitchingKeyPreparedAlloc<B>,
+{
+    fn automorphism_key_prepared_alloc(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> AutomorphismKeyPrepared<Vec<u8>, B> {
+        AutomorphismKeyPrepared::<Vec<u8>, B> {
+            key: self.glwe_switching_key_prepared_alloc(base2k, k, rank, rank, dnum, dsize),
+            p: 0,
+        }
+    }
+}
+
+pub trait AutomorphismKeyPreparedAllocFromInfos<B: Backend> {
+    fn automorphism_key_prepared_alloc_from_infos<A>(&self, infos: &A) -> AutomorphismKeyPrepared<Vec<u8>, B>
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> AutomorphismKeyPreparedAllocFromInfos<B> for Module<B>
+where
+    Module<B>: GLWESwitchingKeyPreparedAllocFromInfos<B>,
+{
+    fn automorphism_key_prepared_alloc_from_infos<A>(&self, infos: &A) -> AutomorphismKeyPrepared<Vec<u8>, B>
+    where
+        A: GGLWEInfos,
+    {
+        assert_eq!(
+            infos.rank_in(),
+            infos.rank_out(),
+            "rank_in != rank_out is not supported for AutomorphismKeyPrepared"
+        );
+        AutomorphismKeyPrepared {
+            key: self.glwe_switching_key_prepared_alloc_from_infos(infos),
+            p: 0,
+        }
+    }
+}
+
+pub trait AutomorphismKeyPreparedAllocBytes<B: Backend> {
+    fn automorphism_key_prepared_alloc_bytes(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> usize;
+}
+
+impl<B: Backend> AutomorphismKeyPreparedAllocBytes<B> for Module<B>
+where
+    Module<B>: GLWESwitchingKeyPreparedAllocBytes<B>,
+{
+    fn automorphism_key_prepared_alloc_bytes(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> usize {
+        self.glwe_switching_key_prepared_alloc_bytes(base2k, k, rank, rank, dnum, dsize)
+    }
+}
+
+pub trait AutomorphismKeyPreparedAllocBytesFromInfos<B: Backend> {
+    fn automorphism_key_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> AutomorphismKeyPreparedAllocBytesFromInfos<B> for Module<B>
+where
+    Module<B>: GLWESwitchingKeyPreparedAllocBytesFromInfos<B>,
+{
+    fn automorphism_key_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>: VmpPMatAlloc<B>,
     {
         assert_eq!(
             infos.rank_in(),
             infos.rank_out(),
             "rank_in != rank_out is not supported for GGLWEAutomorphismKeyPrepared"
         );
-        AutomorphismKeyPrepared::<Vec<u8>, B> {
+        self.glwe_switching_key_prepared_alloc_bytes_from_infos(infos)
-            key: GLWESwitchingKeyPrepared::alloc(module, infos),
-            p: 0,
     }
+}
+
+impl<B: Backend> AutomorphismKeyPrepared<Vec<u8>, B> {
+    pub fn alloc_from_infos<A>(module: &Module<B>, infos: &A) -> Self
+    where
+        A: GGLWEInfos,
+        Module<B>: AutomorphismKeyPreparedAllocFromInfos<B>,
+    {
+        module.automorphism_key_prepared_alloc_from_infos(infos)
     }
 
     pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
     where
-        Module<B>: VmpPMatAlloc<B>,
+        Module<B>: AutomorphismKeyPreparedAlloc<B>,
     {
-        AutomorphismKeyPrepared {
+        module.automorphism_key_prepared_alloc(base2k, k, rank, dnum, dsize)
-            key: GLWESwitchingKeyPrepared::alloc_with(module, base2k, k, rank, rank, dnum, dsize),
-            p: 0,
-        }
     }
 
     pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>: VmpPMatAllocBytes,
+        Module<B>: AutomorphismKeyPreparedAllocBytesFromInfos<B>,
     {
-        assert_eq!(
+        module.automorphism_key_prepared_alloc_bytes_from_infos(infos)
-            infos.rank_in(),
-            infos.rank_out(),
-            "rank_in != rank_out is not supported for GGLWEAutomorphismKeyPrepared"
-        );
-        GLWESwitchingKeyPrepared::alloc_bytes(module, infos)
     }
 
     pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
     where
-        Module<B>: VmpPMatAllocBytes,
+        Module<B>: AutomorphismKeyPreparedAllocBytes<B>,
     {
-        GLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, rank, rank, dnum, dsize)
+        module.automorphism_key_prepared_alloc_bytes(base2k, k, rank, dnum, dsize)
     }
 }
 
 pub trait AutomorphismKeyPrepareTmpBytes {
-    fn automorphism_key_prepare_tmp_bytes<A>(&self, infos: &A)
+    fn automorphism_key_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
     where
         A: GGLWEInfos;
 }
@@ -132,7 +225,7 @@ impl<B: Backend> AutomorphismKeyPrepareTmpBytes for Module<B>
 where
     Module<B>: GLWESwitchingKeyPrepareTmpBytes,
 {
-    fn automorphism_key_prepare_tmp_bytes<A>(&self, infos: &A)
+    fn automorphism_key_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
     where
         A: GGLWEInfos,
     {
@@ -145,45 +238,28 @@ impl<D: DataRef, B: Backend> AutomorphismKeyPrepared<D, B> {
     where
         Module<B>: AutomorphismKeyPrepareTmpBytes,
     {
-        module.automorphism_key_prepare_tmp_bytes(self);
+        module.automorphism_key_prepare_tmp_bytes(self)
     }
 }
 
 pub trait AutomorphismKeyPrepare<B: Backend> {
-    fn automorphism_key_prepare<R, O>(&self, res: &R, other: &O, scratch: &Scratch<B>)
+    fn automorphism_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
     where
-        R: AutomorphismKeyPreparedToMut<B>,
+        R: AutomorphismKeyPreparedToMut<B> + SetAutomorphismGaloisElement,
-        O: AutomorphismKeyToRef;
+        O: AutomorphismKeyToRef + GetAutomorphismGaloisElement;
 }
 
-impl<B: Backend> AutomorphismKeyPrepare<B> for Module<B> {
+impl<B: Backend> AutomorphismKeyPrepare<B> for Module<B>
-    fn automorphism_key_prepare<R, O>(&self, res: &R, other: &O, scratch: &Scratch<B>)
-    where
-        R: AutomorphismKeyPreparedToMut<B>,
-        O: AutomorphismKeyToRef,
-    {
-        self.key.prepare(self, &other.to_ref().key, scratch);
-        self.p = other.p;
-    }
-}
-
-pub trait AutomorphismKeyPrepareAlloc<B: Backend> {
-    fn automorphism_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>) -> AutomorphismKeyPrepared<Vec<u8>, B>
-    where
-        O: AutomorphismKeyToRef;
-}
-
-impl<B: Backend> AutomorphismKeyPrepareAlloc<B> for Module<B>
 where
-    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
+    Module<B>: GLWESwitchingKeyPrepare<B>,
 {
-    fn automorphism_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>) -> AutomorphismKeyPrepared<Vec<u8>, B>
+    fn automorphism_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
     where
-        O: AutomorphismKeyToRef,
+        R: AutomorphismKeyPreparedToMut<B> + SetAutomorphismGaloisElement,
+        O: AutomorphismKeyToRef + GetAutomorphismGaloisElement,
     {
-        let mut atk_prepared: AutomorphismKeyPrepared<Vec<u8>, B> = AutomorphismKeyPrepared::alloc(self, &other.to_ref());
+        self.glwe_switching_prepare(&mut res.to_mut().key, &other.to_ref().key, scratch);
-        atk_prepared.prepare(self, &other.to_ref(), scratch);
+        res.set_p(other.p());
-        atk_prepared
     }
 }
 

poulpy-core/src/layouts/prepared/gglwe_ct.rs

@@ -160,15 +160,70 @@ impl<D: Data, B: Backend> GGLWEPreparedBuilder<D, B> {
     }
 }
 
-impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
+pub trait GGLWEPreparedAlloc<B: Backend> {
-    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
+    fn gglwe_prepared_alloc(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank_in: Rank,
+        rank_out: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> GGLWEPrepared<Vec<u8>, B>;
+}
+
+impl<B: Backend> GGLWEPreparedAlloc<B> for Module<B>
+where
+    Module<B>: VmpPMatAlloc<B>,
+{
+    fn gglwe_prepared_alloc(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank_in: Rank,
+        rank_out: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> GGLWEPrepared<Vec<u8>, B> {
+        let size: usize = k.0.div_ceil(base2k.0) as usize;
+        debug_assert!(
+            size as u32 > dsize.0,
+            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
+            dsize.0
+        );
+
+        assert!(
+            dnum.0 * dsize.0 <= size as u32,
+            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
+            dnum.0,
+            dsize.0,
+        );
+
+        GGLWEPrepared {
+            data: self.vmp_pmat_alloc(dnum.into(), rank_in.into(), (rank_out + 1).into(), size),
+            k,
+            base2k,
+            dsize,
+        }
+    }
+}
+
+pub trait GGLWEPreparedAllocFromInfos<B: Backend> {
+    fn gglwe_prepared_alloc_from_infos<A>(&self, infos: &A) -> GGLWEPrepared<Vec<u8>, B>
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> GGLWEPreparedAllocFromInfos<B> for Module<B>
+where
+    Module<B>: GGLWEPreparedAlloc<B>,
+{
+    fn gglwe_prepared_alloc_from_infos<A>(&self, infos: &A) -> GGLWEPrepared<Vec<u8>, B>
     where
         A: GGLWEInfos,
-        Module<B>: VmpPMatAlloc<B>,
     {
-        debug_assert_eq!(module.n(), infos.n().0 as usize, "module.n() != infos.n()");
+        assert_eq!(self.n() as u32, infos.n(), "module.n() != infos.n()");
-        Self::alloc_with(
+        self.gglwe_prepared_alloc(
-            module,
             infos.base2k(),
             infos.k(),
             infos.rank_in(),
@@ -177,8 +232,87 @@ impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
             infos.dsize(),
         )
     }
+}
 
-    pub fn alloc_with(
+pub trait GGLWEPreparedAllocBytes<B: Backend> {
+    fn gglwe_prepared_alloc_bytes(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank_in: Rank,
+        rank_out: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> usize;
+}
+
+impl<B: Backend> GGLWEPreparedAllocBytes<B> for Module<B>
+where
+    Module<B>: VmpPMatAllocBytes,
+{
+    fn gglwe_prepared_alloc_bytes(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank_in: Rank,
+        rank_out: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> usize {
+        let size: usize = k.0.div_ceil(base2k.0) as usize;
+        debug_assert!(
+            size as u32 > dsize.0,
+            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
+            dsize.0
+        );
+
+        assert!(
+            dnum.0 * dsize.0 <= size as u32,
+            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
+            dnum.0,
+            dsize.0,
+        );
+
+        self.vmp_pmat_alloc_bytes(dnum.into(), rank_in.into(), (rank_out + 1).into(), size)
+    }
+}
+
+pub trait GGLWEPreparedAllocBytesFromInfos<B: Backend> {
+    fn gglwe_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> GGLWEPreparedAllocBytesFromInfos<B> for Module<B>
+where
+    Module<B>: GGLWEPreparedAllocBytes<B>,
+{
+    fn gglwe_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos,
+    {
+        assert_eq!(self.n() as u32, infos.n(), "module.n() != infos.n()");
+        self.gglwe_prepared_alloc_bytes(
+            infos.base2k(),
+            infos.k(),
+            infos.rank_in(),
+            infos.rank_out(),
+            infos.dnum(),
+            infos.dsize(),
+        )
+    }
+}
+
+impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
+    pub fn alloc_from_infos<A>(module: &Module<B>, infos: &A) -> Self
+    where
+        A: GGLWEInfos,
+        Module<B>: GGLWEPreparedAllocFromInfos<B>,
+    {
+        module.gglwe_prepared_alloc_from_infos(infos)
+    }
+
+    pub fn alloc(
         module: &Module<B>,
         base2k: Base2K,
         k: TorusPrecision,
@@ -188,48 +322,20 @@ impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
         dsize: Dsize,
     ) -> Self
     where
-        Module<B>: VmpPMatAlloc<B>,
+        Module<B>: GGLWEPreparedAlloc<B>,
     {
-        let size: usize = k.0.div_ceil(base2k.0) as usize;
+        module.gglwe_prepared_alloc(base2k, k, rank_in, rank_out, dnum, dsize)
-        debug_assert!(
-            size as u32 > dsize.0,
-            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
-            dsize.0
-        );
-
-        assert!(
-            dnum.0 * dsize.0 <= size as u32,
-            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
-            dnum.0,
-            dsize.0,
-        );
-
-        Self {
-            data: module.vmp_pmat_alloc(dnum.into(), rank_in.into(), (rank_out + 1).into(), size),
-            k,
-            base2k,
-            dsize,
-        }
     }
 
-    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
+    pub fn alloc_bytes_from_infos<A>(module: &Module<B>, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>: VmpPMatAllocBytes,
+        Module<B>: GGLWEPreparedAllocBytesFromInfos<B>,
     {
-        debug_assert_eq!(module.n(), infos.n().0 as usize, "module.n() != infos.n()");
+        module.gglwe_prepared_alloc_bytes_from_infos(infos)
-        Self::alloc_bytes_with(
-            module,
-            infos.base2k(),
-            infos.k(),
-            infos.rank_in(),
-            infos.rank_out(),
-            infos.dnum(),
-            infos.dsize(),
-        )
     }
 
-    pub fn alloc_bytes_with(
+    pub fn alloc_bytes(
         module: &Module<B>,
         base2k: Base2K,
         k: TorusPrecision,
@@ -239,23 +345,9 @@ impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
         dsize: Dsize,
     ) -> usize
     where
-        Module<B>: VmpPMatAllocBytes,
+        Module<B>: GGLWEPreparedAllocBytes<B>,
     {
-        let size: usize = k.0.div_ceil(base2k.0) as usize;
+        module.gglwe_prepared_alloc_bytes(base2k, k, rank_in, rank_out, dnum, dsize)
-        debug_assert!(
-            size as u32 > dsize.0,
-            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
-            dsize.0
-        );
-
-        assert!(
-            dnum.0 * dsize.0 <= size as u32,
-            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
-            dnum.0,
-            dsize.0,
-        );
-
-        module.vmp_pmat_alloc_bytes(dnum.into(), rank_in.into(), (rank_out + 1).into(), size)
     }
 }
 
@@ -283,7 +375,7 @@ where
 }
 
 impl<B: Backend> GGLWEPrepared<Vec<u8>, B> {
-    pub fn prepare_tmp_bytes(&self, module: &Module<B>)
+    pub fn prepare_tmp_bytes(&self, module: &Module<B>) -> usize
     where
         Module<B>: GGLWEPrepareTmpBytes,
     {
@@ -307,7 +399,7 @@ where
         R: GGLWEPreparedToMut<B>,
         O: GGLWEToRef,
     {
-        let mut res: GGLWEPrepared<&mut [u8], B> = self.to_mut();
+        let mut res: GGLWEPrepared<&mut [u8], B> = res.to_mut();
         let other: GGLWE<&[u8]> = other.to_ref();
 
         assert_eq!(res.n(), self.n() as u32);
@@ -332,30 +424,6 @@ where
     }
 }
 
-pub trait GGLWEPrepareAlloc<B: Backend> {
-    fn gglwe_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>) -> GGLWEPrepared<Vec<u8>, B>;
-}
-
-impl<B: Backend> GGLWEPrepareAlloc<B> for Module<B>
-where
-    Module<B>: VmpPMatAlloc<B> + VmpPrepare<B>,
-{
-    fn gglwe_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>) -> GGLWEPrepared<Vec<u8>, B> {
-        let mut ct_prepared: GGLWEPrepared<Vec<u8>, B> = GGLWEPrepared::alloc(self, &other.to_ref());
-        ct_prepared.prepare(self, &other.to_ref(), scratch);
-        ct_prepared
-    }
-}
-
-impl<D: DataRef> GGLWE<D> {
-    fn prepare_alloc<B: Backend>(&self, module: &Module<B>, scratch: &Scratch<B>) -> GGLWEPrepared<Vec<u8>, B>
-    where
-        Module<B>: GGLWEPrepareAlloc<B>,
-    {
-        module.gglwe_prepare_alloc(self, scratch)
-    }
-}
-
 pub trait GGLWEPreparedToMut<B: Backend> {
     fn to_mut(&mut self) -> GGLWEPrepared<&mut [u8], B>;
 }

poulpy-core/src/layouts/prepared/gglwe_ksk.rs

@@ -1,11 +1,12 @@
-use poulpy_hal::{
+use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
-    api::{VmpPMatAlloc, VmpPMatAllocBytes},
-    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
-};
 
 use crate::layouts::{
-    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, GLWESwitchingKey, GLWESwitchingKeyToRef, LWEInfos, Rank, TorusPrecision,
+    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, GLWESwitchingKeySetMetaData, GLWESwitchingKeyToRef,
-    prepared::{GGLWEPrepare, GGLWEPrepareTmpBytes, GGLWEPrepared, GGLWEPreparedToMut, GGLWEPreparedToRef},
+    GLWESwtichingKeyGetMetaData, LWEInfos, Rank, TorusPrecision,
+    prepared::{
+        GGLWEPrepare, GGLWEPrepareTmpBytes, GGLWEPrepared, GGLWEPreparedAlloc, GGLWEPreparedAllocBytes,
+        GGLWEPreparedAllocBytesFromInfos, GGLWEPreparedAllocFromInfos, GGLWEPreparedToMut, GGLWEPreparedToRef,
+    },
 };
 
 #[derive(PartialEq, Eq)]
@@ -15,18 +16,23 @@ pub struct GLWESwitchingKeyPrepared<D: Data, B: Backend> {
     pub(crate) sk_out_n: usize, // Degree of sk_out
 }
 
-pub(crate) trait GLWESwitchingKeyPreparedSetMetaData {
+impl<D: DataMut, B: Backend> GLWESwitchingKeySetMetaData for GLWESwitchingKeyPrepared<D, B> {
-    fn set_sk_in_n(&mut self, sk_in_n: usize);
-    fn set_sk_out_n(&mut self, sk_out_n: usize);
-}
-
-impl<D: DataMut, B: Backend> GLWESwitchingKeyPreparedSetMetaData for GLWESwitchingKeyPrepared<D, B> {
     fn set_sk_in_n(&mut self, sk_in_n: usize) {
         self.sk_in_n = sk_in_n
     }
 
     fn set_sk_out_n(&mut self, sk_out_n: usize) {
-        self.sk_out_n = self.sk_out_n
+        self.sk_out_n = sk_out_n
+    }
+}
+
+impl<D: DataRef, B: Backend> GLWESwtichingKeyGetMetaData for GLWESwitchingKeyPrepared<D, B> {
+    fn sk_in_n(&self) -> usize {
+        self.sk_in_n
+    }
+
+    fn sk_out_n(&self) -> usize {
+        self.sk_out_n
     }
 }
 
@@ -72,21 +78,118 @@ impl<D: Data, B: Backend> GGLWEInfos for GLWESwitchingKeyPrepared<D, B> {
     }
 }
 
-impl<B: Backend> GLWESwitchingKeyPrepared<Vec<u8>, B> {
+pub trait GLWESwitchingKeyPreparedAlloc<B: Backend> {
-    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
+    fn glwe_switching_key_prepared_alloc(
-    where
+        &self,
-        A: GGLWEInfos,
+        base2k: Base2K,
-        Module<B>: VmpPMatAlloc<B>,
+        k: TorusPrecision,
-    {
+        rank_in: Rank,
-        debug_assert_eq!(module.n() as u32, infos.n(), "module.n() != infos.n()");
+        rank_out: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> GLWESwitchingKeyPrepared<Vec<u8>, B>;
+}
+
+impl<B: Backend> GLWESwitchingKeyPreparedAlloc<B> for Module<B>
+where
+    Module<B>: GGLWEPreparedAlloc<B>,
+{
+    fn glwe_switching_key_prepared_alloc(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank_in: Rank,
+        rank_out: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> GLWESwitchingKeyPrepared<Vec<u8>, B> {
         GLWESwitchingKeyPrepared::<Vec<u8>, B> {
-            key: GGLWEPrepared::alloc(module, infos),
+            key: self.gglwe_prepared_alloc(base2k, k, rank_in, rank_out, dnum, dsize),
             sk_in_n: 0,
             sk_out_n: 0,
         }
     }
+}
 
-    pub fn alloc_with(
+pub trait GLWESwitchingKeyPreparedAllocFromInfos<B: Backend> {
+    fn glwe_switching_key_prepared_alloc_from_infos<A>(&self, infos: &A) -> GLWESwitchingKeyPrepared<Vec<u8>, B>
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> GLWESwitchingKeyPreparedAllocFromInfos<B> for Module<B>
+where
+    Module<B>: GGLWEPreparedAllocFromInfos<B>,
+{
+    fn glwe_switching_key_prepared_alloc_from_infos<A>(&self, infos: &A) -> GLWESwitchingKeyPrepared<Vec<u8>, B>
+    where
+        A: GGLWEInfos,
+    {
+        GLWESwitchingKeyPrepared::<Vec<u8>, B> {
+            key: self.gglwe_prepared_alloc_from_infos(infos),
+            sk_in_n: 0,
+            sk_out_n: 0,
+        }
+    }
+}
+
+pub trait GLWESwitchingKeyPreparedAllocBytes<B: Backend> {
+    fn glwe_switching_key_prepared_alloc_bytes(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank_in: Rank,
+        rank_out: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> usize;
+}
+
+impl<B: Backend> GLWESwitchingKeyPreparedAllocBytes<B> for Module<B>
+where
+    Module<B>: GGLWEPreparedAllocBytes<B>,
+{
+    fn glwe_switching_key_prepared_alloc_bytes(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank_in: Rank,
+        rank_out: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> usize {
+        self.gglwe_prepared_alloc_bytes(base2k, k, rank_in, rank_out, dnum, dsize)
+    }
+}
+
+pub trait GLWESwitchingKeyPreparedAllocBytesFromInfos<B: Backend> {
+    fn glwe_switching_key_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> GLWESwitchingKeyPreparedAllocBytesFromInfos<B> for Module<B>
+where
+    Module<B>: GGLWEPreparedAllocBytesFromInfos<B>,
+{
+    fn glwe_switching_key_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos,
+    {
+        self.gglwe_prepared_alloc_bytes_from_infos(infos)
+    }
+}
+
+impl<B: Backend> GLWESwitchingKeyPrepared<Vec<u8>, B> {
+    pub fn alloc_from_infos<A>(module: &Module<B>, infos: &A) -> Self
+    where
+        A: GGLWEInfos,
+        Module<B>: GLWESwitchingKeyPreparedAllocFromInfos<B>,
+    {
+        module.glwe_switching_key_prepared_alloc_from_infos(infos)
+    }
+
+    pub fn alloc(
         module: &Module<B>,
         base2k: Base2K,
         k: TorusPrecision,
@@ -96,25 +199,20 @@ impl<B: Backend> GLWESwitchingKeyPrepared<Vec<u8>, B> {
         dsize: Dsize,
     ) -> Self
     where
-        Module<B>: VmpPMatAlloc<B>,
+        Module<B>: GLWESwitchingKeyPreparedAlloc<B>,
     {
-        GLWESwitchingKeyPrepared::<Vec<u8>, B> {
+        module.glwe_switching_key_prepared_alloc(base2k, k, rank_in, rank_out, dnum, dsize)
-            key: GGLWEPrepared::alloc_with(module, base2k, k, rank_in, rank_out, dnum, dsize),
-            sk_in_n: 0,
-            sk_out_n: 0,
-        }
     }
 
-    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
+    pub fn alloc_bytes_from_infos<A>(module: &Module<B>, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>: VmpPMatAllocBytes,
+        Module<B>: GLWESwitchingKeyPreparedAllocBytesFromInfos<B>,
     {
-        debug_assert_eq!(module.n() as u32, infos.n(), "module.n() != infos.n()");
+        module.glwe_switching_key_prepared_alloc_bytes_from_infos(infos)
-        GGLWEPrepared::alloc_bytes(module, infos)
     }
 
-    pub fn alloc_bytes_with(
+    pub fn alloc_bytes(
         module: &Module<B>,
         base2k: Base2K,
         k: TorusPrecision,
@@ -124,9 +222,9 @@ impl<B: Backend> GLWESwitchingKeyPrepared<Vec<u8>, B> {
         dsize: Dsize,
     ) -> usize
     where
-        Module<B>: VmpPMatAllocBytes,
+        Module<B>: GLWESwitchingKeyPreparedAllocBytes<B>,
     {
-        GGLWEPrepared::alloc_bytes_with(module, base2k, k, rank_in, rank_out, dnum, dsize)
+        module.glwe_switching_key_prepared_alloc_bytes(base2k, k, rank_in, rank_out, dnum, dsize)
     }
 }
 
@@ -158,66 +256,37 @@ where
 }
 
 pub trait GLWESwitchingKeyPrepare<B: Backend> {
-    fn glwe_switching_prepare<R, O>(&self, res: &R, other: &O, scratch: &mut Scratch<B>)
+    fn glwe_switching_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
     where
-        R: GLWESwitchingKeyPreparedToMut<B> + GLWESwitchingKeyPreparedSetMetaData,
+        R: GLWESwitchingKeyPreparedToMut<B> + GLWESwitchingKeySetMetaData,
-        O: GLWESwitchingKeyToRef;
+        O: GLWESwitchingKeyToRef + GLWESwtichingKeyGetMetaData;
 }
 
 impl<B: Backend> GLWESwitchingKeyPrepare<B> for Module<B>
 where
     Module<B>: GGLWEPrepare<B>,
 {
-    fn glwe_switching_prepare<R, O>(&self, res: &R, other: &O, scratch: &mut Scratch<B>)
+    fn glwe_switching_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
     where
-        R: GLWESwitchingKeyPreparedToMut<B> + GLWESwitchingKeyPreparedSetMetaData,
+        R: GLWESwitchingKeyPreparedToMut<B> + GLWESwitchingKeySetMetaData,
-        O: GLWESwitchingKeyToRef,
+        O: GLWESwitchingKeyToRef + GLWESwtichingKeyGetMetaData,
     {
-        self.gglwe_prepare(&res.to_mut(), other, scratch);
+        self.gglwe_prepare(&mut res.to_mut().key, &other.to_ref().key, scratch);
-        res.set_sk_in_n(other.sk_in_n);
+        res.set_sk_in_n(other.sk_in_n());
-        res.set_sk_out_n(other.sk_out_n);
+        res.set_sk_out_n(other.sk_out_n());
     }
 }
 
 impl<D: DataMut, B: Backend> GLWESwitchingKeyPrepared<D, B> {
     pub fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &mut Scratch<B>)
     where
-        O: GLWESwitchingKeyToRef,
+        O: GLWESwitchingKeyToRef + GLWESwtichingKeyGetMetaData,
         Module<B>: GLWESwitchingKeyPrepare<B>,
     {
         module.glwe_switching_prepare(self, other, scratch);
     }
 }
 
-pub trait GLWESwitchingKeyPrepareAlloc<B: Backend> {
-    fn glwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
-    where
-        O: GLWESwitchingKeyToRef;
-}
-
-impl<B: Backend> GLWESwitchingKeyPrepareAlloc<B> for Module<B>
-where
-    Module<B>: GLWESwitchingKeyPrepare<B>,
-{
-    fn glwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
-    where
-        O: GLWESwitchingKeyToRef,
-    {
-        let mut ct_prepared: GLWESwitchingKeyPrepared<Vec<u8>, B> = GLWESwitchingKeyPrepared::alloc(self, self);
-        self.glwe_switching_prepare(&mut ct_prepared, other, scratch);
-        ct_prepared
-    }
-}
-
-impl<D: DataRef> GLWESwitchingKey<D> {
-    pub fn prepare_alloc<B: Backend>(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GLWESwitchingKeyPrepared<Vec<u8>, B>
-    where
-        Module<B>: GLWESwitchingKeyPrepareAlloc<B>,
-    {
-        module.glwe_switching_key_prepare_alloc(self, scratch);
-    }
-}
-
 pub trait GLWESwitchingKeyPreparedToMut<B: Backend> {
     fn to_mut(&mut self) -> GLWESwitchingKeyPrepared<&mut [u8], B>;
 }

poulpy-core/src/layouts/prepared/gglwe_tsk.rs

@@ -87,7 +87,7 @@ impl<B: Backend> TensorKeyPrepared<Vec<u8>, B> {
         let mut keys: Vec<GLWESwitchingKeyPrepared<Vec<u8>, B>> = Vec::new();
         let pairs: u32 = (((rank.0 + 1) * rank.0) >> 1).max(1);
         (0..pairs).for_each(|_| {
-            keys.push(GLWESwitchingKeyPrepared::alloc_with(
+            keys.push(GLWESwitchingKeyPrepared::alloc(
                 module,
                 base2k,
                 k,
@@ -113,7 +113,7 @@ impl<B: Backend> TensorKeyPrepared<Vec<u8>, B> {
         let rank_out: usize = infos.rank_out().into();
         let pairs: usize = (((rank_out + 1) * rank_out) >> 1).max(1);
         pairs
-            * GLWESwitchingKeyPrepared::alloc_bytes_with(
+            * GLWESwitchingKeyPrepared::alloc_bytes(
                 module,
                 infos.base2k(),
                 infos.k(),
@@ -129,7 +129,7 @@ impl<B: Backend> TensorKeyPrepared<Vec<u8>, B> {
         Module<B>: VmpPMatAllocBytes,
     {
         let pairs: usize = (((rank.0 + 1) * rank.0) >> 1).max(1) as usize;
-        pairs * GLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, Rank(1), rank, dnum, dsize)
+        pairs * GLWESwitchingKeyPrepared::alloc_bytes(module, base2k, k, Rank(1), rank, dnum, dsize)
     }
 }
 

poulpy-core/src/layouts/prepared/glwe_to_lwe_ksk.rs

@@ -1,11 +1,15 @@
 use poulpy_hal::{
-    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
+    api::{VmpPMatAlloc, VmpPMatAllocBytes},
     layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 };
 
 use crate::layouts::{
-    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, GLWEToLWESwitchingKey, LWEInfos, Rank, TorusPrecision,
+    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, GLWEToLWESwitchingKeyToMut, GLWEToLWESwitchingKeyToRef, LWEInfos, Rank,
-    prepared::GLWESwitchingKeyPrepared,
+    TorusPrecision,
+    prepared::{
+        GLWESecretPrepareTmpBytes, GLWESwitchingKeyPrepare, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToMut,
+        GLWESwitchingKeyPreparedToRef,
+    },
 };
 
 #[derive(PartialEq, Eq)]
@@ -69,14 +73,14 @@ impl<B: Backend> GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> {
             1,
             "dsize > 1 is not supported for GLWEToLWESwitchingKeyPrepared"
         );
-        Self(GLWESwitchingKeyPrepared::alloc(module, infos))
+        Self(GLWESwitchingKeyPrepared::alloc_from_infos(module, infos))
     }
 
     pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank_in: Rank, dnum: Dnum) -> Self
     where
         Module<B>: VmpPMatAlloc<B>,
     {
-        Self(GLWESwitchingKeyPrepared::alloc_with(
+        Self(GLWESwitchingKeyPrepared::alloc(
             module,
             base2k,
             k,
@@ -102,42 +106,117 @@ impl<B: Backend> GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> {
             1,
             "dsize > 1 is not supported for GLWEToLWESwitchingKeyPrepared"
         );
-        GLWESwitchingKeyPrepared::alloc_bytes(module, infos)
+        GLWESwitchingKeyPrepared::alloc_bytes_from_infos(module, infos)
     }
 
     pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank_in: Rank, dnum: Dnum) -> usize
     where
         Module<B>: VmpPMatAllocBytes,
     {
-        GLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, rank_in, Rank(1), dnum, Dsize(1))
+        GLWESwitchingKeyPrepared::alloc_bytes(module, base2k, k, rank_in, Rank(1), dnum, Dsize(1))
     }
 }
 
-impl<B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for GLWEToLWESwitchingKeyPrepared<Vec<u8>, B>
+pub trait GLWEToLWESwitchingKeyPrepareTmpBytes {
+    fn glwe_to_lwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> GLWEToLWESwitchingKeyPrepareTmpBytes for Module<B>
 where
-    GLWESwitchingKeyPrepared<Vec<u8>, B>: PrepareScratchSpace<B, A>,
+    Module<B>: GLWEToLWESwitchingKeyPrepareTmpBytes,
 {
-    fn prepare_scratch_space(module: &Module<B>, infos: &A) -> usize {
+    fn glwe_to_lwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
-        GLWESwitchingKeyPrepared::prepare_scratch_space(module, infos)
+    where
+        A: GGLWEInfos,
+    {
+        self.glwe_to_lwe_switching_key_prepare_tmp_bytes(infos);
     }
 }
 
-impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWEToLWESwitchingKeyPrepared<Vec<u8>, B>> for GLWEToLWESwitchingKey<D>
+impl<B: Backend> GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> {
-where
+    pub fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A)
-    Module<B>: VmpPrepare<B> + VmpPMatAlloc<B>,
+    where
-{
+        A: GLWEInfos,
-    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> {
+        Module<B>: GLWEToLWESwitchingKeyPrepareTmpBytes,
-        let mut ksk_prepared: GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> = GLWEToLWESwitchingKeyPrepared::alloc(module, self);
+    {
-        ksk_prepared.prepare(module, self, scratch);
+        module.glwe_secret_prepare_tmp_bytes(infos);
-        ksk_prepared
     }
 }
 
-impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWEToLWESwitchingKey<DR>> for GLWEToLWESwitchingKeyPrepared<DM, B>
+pub trait GLWEToLWESwitchingKeyPrepare<B: Backend> {
+    fn glwe_to_lwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
+    where
+        R: GLWEToLWESwitchingKeyPreparedToMut<B>,
+        O: GLWEToLWESwitchingKeyToRef;
+}
+
+impl<B: Backend> GLWEToLWESwitchingKeyPrepare<B> for Module<B>
 where
-    Module<B>: VmpPrepare<B>,
+    Module<B>: GLWESwitchingKeyPrepare<B>,
 {
-    fn prepare(&mut self, module: &Module<B>, other: &GLWEToLWESwitchingKey<DR>, scratch: &mut Scratch<B>) {
+    fn glwe_to_lwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
-        self.0.prepare(module, &other.0, scratch);
+    where
+        R: GLWEToLWESwitchingKeyPreparedToMut<B>,
+        O: GLWEToLWESwitchingKeyToRef,
+    {
+        self.glwe_switching_prepare(&mut res.to_mut().0, other, scratch);
+    }
+}
+
+impl<D: DataMut, B: Backend> GLWEToLWESwitchingKeyPrepared<D, B> {
+    fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &Scratch<B>)
+    where
+        O: GLWEToLWESwitchingKeyToRef,
+        Module<B>: GLWEToLWESwitchingKeyPrepare<B>,
+    {
+        module.glwe_to_lwe_switching_key_prepare(self, other, scratch);
+    }
+}
+
+pub trait GLWEToLWESwitchingKeyPrepareAlloc<B: Backend> {
+    fn glwe_to_lwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: GLWEToLWESwitchingKeyToRef;
+}
+
+impl<B: Backend> GLWEToLWESwitchingKeyPrepareAlloc<B> for Module<B>
+where
+    Module<B>: GLWEToLWESwitchingKeyPrepare<B>,
+{
+    fn glwe_to_lwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: GLWEToLWESwitchingKeyToRef,
+    {
+        let mut ct_prep: GLWEToLWESwitchingKeyPrepared<Vec<u8>, B> = GLWEToLWESwitchingKeyPrepared::alloc(self, &other.to_ref());
+        self.glwe_to_lwe_switching_key_prepare(&mut ct_prep, other, scratch);
+        ct_prep
+    }
+}
+
+pub trait GLWEToLWESwitchingKeyPreparedToRef<B: Backend> {
+    fn to_ref(&self) -> GLWEToLWESwitchingKeyPrepared<&[u8], B>;
+}
+
+impl<D: DataRef, B: Backend> GLWEToLWESwitchingKeyPreparedToRef<B> for GLWEToLWESwitchingKeyPrepared<D, B>
+where
+    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToRef<B>,
+{
+    fn to_ref(&self) -> GLWEToLWESwitchingKeyPrepared<&[u8], B> {
+        GLWEToLWESwitchingKeyPrepared(self.0.to_ref())
+    }
+}
+
+pub trait GLWEToLWESwitchingKeyPreparedToMut<B: Backend> {
+    fn to_mut(&mut self) -> GLWEToLWESwitchingKeyPrepared<&mut [u8], B>;
+}
+
+impl<D: DataMut, B: Backend> GLWEToLWESwitchingKeyPreparedToMut<B> for GLWEToLWESwitchingKeyPrepared<D, B>
+where
+    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToMut<B>,
+{
+    fn to_mut(&mut self) -> GLWEToLWESwitchingKeyPrepared<&mut [u8], B> {
+        GLWEToLWESwitchingKeyPrepared(self.0.to_mut())
     }
 }

poulpy-core/src/layouts/prepared/lwe_ksk.rs

@@ -1,11 +1,11 @@
 use poulpy_hal::{
-    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
+    api::{VmpPMatAlloc, VmpPMatAllocBytes},
     layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 };
 
 use crate::layouts::{
-    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWESwitchingKey, Rank, TorusPrecision,
+    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWESwitchingKeyToRef, Rank, TorusPrecision,
-    prepared::GLWESwitchingKeyPrepared,
+    prepared::{GLWESwitchingKeyPrepare, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToMut, GLWESwitchingKeyPreparedToRef},
 };
 
 #[derive(PartialEq, Eq)]
@@ -73,14 +73,14 @@ impl<B: Backend> LWESwitchingKeyPrepared<Vec<u8>, B> {
             1,
             "rank_out > 1 is not supported for LWESwitchingKey"
         );
-        Self(GLWESwitchingKeyPrepared::alloc(module, infos))
+        Self(GLWESwitchingKeyPrepared::alloc_from_infos(module, infos))
     }
 
     pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum) -> Self
     where
         Module<B>: VmpPMatAlloc<B>,
     {
-        Self(GLWESwitchingKeyPrepared::alloc_with(
+        Self(GLWESwitchingKeyPrepared::alloc(
             module,
             base2k,
             k,
@@ -111,42 +111,117 @@ impl<B: Backend> LWESwitchingKeyPrepared<Vec<u8>, B> {
             1,
             "rank_out > 1 is not supported for LWESwitchingKey"
         );
-        GLWESwitchingKeyPrepared::alloc_bytes(module, infos)
+        GLWESwitchingKeyPrepared::alloc_bytes_from_infos(module, infos)
     }
 
     pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum) -> usize
     where
         Module<B>: VmpPMatAllocBytes,
     {
-        GLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, Rank(1), Rank(1), dnum, Dsize(1))
+        GLWESwitchingKeyPrepared::alloc_bytes(module, base2k, k, Rank(1), Rank(1), dnum, Dsize(1))
     }
 }
 
-impl<B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for LWESwitchingKeyPrepared<Vec<u8>, B>
+pub trait LWESwitchingKeyPrepareTmpBytes {
+    fn lwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> LWESwitchingKeyPrepareTmpBytes for Module<B>
 where
-    GLWESwitchingKeyPrepared<Vec<u8>, B>: PrepareScratchSpace<B, A>,
+    Module<B>: LWESwitchingKeyPrepareTmpBytes,
 {
-    fn prepare_scratch_space(module: &Module<B>, infos: &A) -> usize {
+    fn lwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
-        GLWESwitchingKeyPrepared::prepare_scratch_space(module, infos)
+    where
+        A: GGLWEInfos,
+    {
+        self.lwe_switching_key_prepare_tmp_bytes(infos);
     }
 }
 
-impl<D: DataRef, B: Backend> PrepareAlloc<B, LWESwitchingKeyPrepared<Vec<u8>, B>> for LWESwitchingKey<D>
+impl<B: Backend> LWESwitchingKeyPrepared<Vec<u8>, B> {
-where
+    pub fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A)
-    Module<B>: VmpPrepare<B> + VmpPMatAlloc<B>,
+    where
-{
+        A: GLWEInfos,
-    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> LWESwitchingKeyPrepared<Vec<u8>, B> {
+        Module<B>: LWESwitchingKeyPrepareTmpBytes,
-        let mut ksk_prepared: LWESwitchingKeyPrepared<Vec<u8>, B> = LWESwitchingKeyPrepared::alloc(module, self);
+    {
-        ksk_prepared.prepare(module, self, scratch);
+        module.glwe_secret_prepare_tmp_bytes(infos);
-        ksk_prepared
     }
 }
 
-impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, LWESwitchingKey<DR>> for LWESwitchingKeyPrepared<DM, B>
+pub trait LWESwitchingKeyPrepare<B: Backend> {
+    fn lwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
+    where
+        R: LWESwitchingKeyPreparedToMut<B>,
+        O: LWESwitchingKeyToRef;
+}
+
+impl<B: Backend> LWESwitchingKeyPrepare<B> for Module<B>
 where
-    Module<B>: VmpPrepare<B>,
+    Module<B>: GLWESwitchingKeyPrepare<B>,
 {
-    fn prepare(&mut self, module: &Module<B>, other: &LWESwitchingKey<DR>, scratch: &mut Scratch<B>) {
+    fn lwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
-        self.0.prepare(module, &other.0, scratch);
+    where
+        R: LWESwitchingKeyPreparedToMut<B>,
+        O: LWESwitchingKeyToRef,
+    {
+        self.glwe_switching_prepare(&mut res.to_mut().0, other, scratch);
+    }
+}
+
+impl<D: DataMut, B: Backend> LWESwitchingKeyPrepared<D, B> {
+    fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &Scratch<B>)
+    where
+        O: LWESwitchingKeyToRef,
+        Module<B>: LWESwitchingKeyPrepare<B>,
+    {
+        module.lwe_switching_key_prepare(self, other, scratch);
+    }
+}
+
+pub trait LWESwitchingKeyPrepareAlloc<B: Backend> {
+    fn lwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: LWESwitchingKeyToRef;
+}
+
+impl<B: Backend> LWESwitchingKeyPrepareAlloc<B> for Module<B>
+where
+    Module<B>: LWESwitchingKeyPrepare<B>,
+{
+    fn lwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: LWESwitchingKeyToRef,
+    {
+        let mut ct_prep: LWESwitchingKeyPrepared<Vec<u8>, B> = LWESwitchingKeyPrepared::alloc(self, &other.to_ref());
+        self.lwe_switching_key_prepare(&mut ct_prep, other, scratch);
+        ct_prep
+    }
+}
+
+pub trait LWESwitchingKeyPreparedToRef<B: Backend> {
+    fn to_ref(&self) -> LWESwitchingKeyPrepared<&[u8], B>;
+}
+
+impl<D: DataRef, B: Backend> LWESwitchingKeyPreparedToRef<B> for LWESwitchingKeyPrepared<D, B>
+where
+    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToRef<B>,
+{
+    fn to_ref(&self) -> LWESwitchingKeyPrepared<&[u8], B> {
+        LWESwitchingKeyPrepared(self.0.to_ref())
+    }
+}
+
+pub trait LWESwitchingKeyPreparedToMut<B: Backend> {
+    fn to_mut(&mut self) -> LWESwitchingKeyPrepared<&mut [u8], B>;
+}
+
+impl<D: DataMut, B: Backend> LWESwitchingKeyPreparedToMut<B> for LWESwitchingKeyPrepared<D, B>
+where
+    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToMut<B>,
+{
+    fn to_mut(&mut self) -> LWESwitchingKeyPrepared<&mut [u8], B> {
+        LWESwitchingKeyPrepared(self.0.to_mut())
     }
 }

poulpy-core/src/layouts/prepared/lwe_to_glwe_ksk.rs

@@ -1,11 +1,11 @@
 use poulpy_hal::{
-    api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
+    api::{VmpPMatAlloc, VmpPMatAllocBytes},
     layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 };
 
 use crate::layouts::{
-    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWEToGLWESwitchingKey, Rank, TorusPrecision,
+    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, LWEInfos, LWEToGLWESwitchingKeyToRef, Rank, TorusPrecision,
-    prepared::GLWESwitchingKeyPrepared,
+    prepared::{GLWESwitchingKeyPrepare, GLWESwitchingKeyPrepared, GLWESwitchingKeyPreparedToMut, GLWESwitchingKeyPreparedToRef},
 };
 
 /// A special [GLWESwitchingKey] required to for the conversion from [LWECiphertext] to [GLWECiphertext].
@@ -70,14 +70,14 @@ impl<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
             1,
             "dsize > 1 is not supported for LWEToGLWESwitchingKey"
         );
-        Self(GLWESwitchingKeyPrepared::alloc(module, infos))
+        Self(GLWESwitchingKeyPrepared::alloc_from_infos(module, infos))
     }
 
     pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> Self
     where
         Module<B>: VmpPMatAlloc<B>,
     {
-        Self(GLWESwitchingKeyPrepared::alloc_with(
+        Self(GLWESwitchingKeyPrepared::alloc(
             module,
             base2k,
             k,
@@ -103,42 +103,117 @@ impl<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
             1,
             "dsize > 1 is not supported for LWEToGLWESwitchingKey"
         );
-        GLWESwitchingKeyPrepared::alloc_bytes(module, infos)
+        GLWESwitchingKeyPrepared::alloc_bytes_from_infos(module, infos)
     }
 
     pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, dnum: Dnum, rank_out: Rank) -> usize
     where
         Module<B>: VmpPMatAllocBytes,
     {
-        GLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, Rank(1), rank_out, dnum, Dsize(1))
+        GLWESwitchingKeyPrepared::alloc_bytes(module, base2k, k, Rank(1), rank_out, dnum, Dsize(1))
     }
 }
 
-impl<B: Backend, A: GGLWEInfos> PrepareScratchSpace<B, A> for LWEToGLWESwitchingKeyPrepared<Vec<u8>, B>
+pub trait LWEToGLWESwitchingKeyPrepareTmpBytes {
+    fn lwe_to_glwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
+    where
+        A: GGLWEInfos;
+}
+
+impl<B: Backend> LWEToGLWESwitchingKeyPrepareTmpBytes for Module<B>
 where
-    GLWESwitchingKeyPrepared<Vec<u8>, B>: PrepareScratchSpace<B, A>,
+    Module<B>: LWEToGLWESwitchingKeyPrepareTmpBytes,
 {
-    fn prepare_scratch_space(module: &Module<B>, infos: &A) -> usize {
+    fn lwe_to_glwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A)
-        GLWESwitchingKeyPrepared::prepare_scratch_space(module, infos)
+    where
+        A: GGLWEInfos,
+    {
+        self.lwe_to_glwe_switching_key_prepare_tmp_bytes(infos);
     }
 }
 
-impl<D: DataRef, B: Backend> PrepareAlloc<B, LWEToGLWESwitchingKeyPrepared<Vec<u8>, B>> for LWEToGLWESwitchingKey<D>
+impl<B: Backend> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
-where
+    pub fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A)
-    Module<B>: VmpPrepare<B> + VmpPMatAlloc<B>,
+    where
-{
+        A: GLWEInfos,
-    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> {
+        Module<B>: LWEToGLWESwitchingKeyPrepareTmpBytes,
-        let mut ksk_prepared: LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> = LWEToGLWESwitchingKeyPrepared::alloc(module, self);
+    {
-        ksk_prepared.prepare(module, self, scratch);
+        module.glwe_secret_prepare_tmp_bytes(infos);
-        ksk_prepared
     }
 }
 
-impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, LWEToGLWESwitchingKey<DR>> for LWEToGLWESwitchingKeyPrepared<DM, B>
+pub trait LWEToGLWESwitchingKeyPrepare<B: Backend> {
+    fn lwe_to_glwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
+    where
+        R: LWEToGLWESwitchingKeyPreparedToMut<B>,
+        O: LWEToGLWESwitchingKeyToRef;
+}
+
+impl<B: Backend> LWEToGLWESwitchingKeyPrepare<B> for Module<B>
 where
-    Module<B>: VmpPrepare<B>,
+    Module<B>: GLWESwitchingKeyPrepare<B>,
 {
-    fn prepare(&mut self, module: &Module<B>, other: &LWEToGLWESwitchingKey<DR>, scratch: &mut Scratch<B>) {
+    fn lwe_to_glwe_switching_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
-        self.0.prepare(module, &other.0, scratch);
+    where
+        R: LWEToGLWESwitchingKeyPreparedToMut<B>,
+        O: LWEToGLWESwitchingKeyToRef,
+    {
+        self.glwe_switching_prepare(&mut res.to_mut().0, other, scratch);
+    }
+}
+
+impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyPrepared<D, B> {
+    fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &Scratch<B>)
+    where
+        O: LWEToGLWESwitchingKeyToRef,
+        Module<B>: LWEToGLWESwitchingKeyPrepare<B>,
+    {
+        module.lwe_to_glwe_switching_key_prepare(self, other, scratch);
+    }
+}
+
+pub trait LWEToGLWESwitchingKeyPrepareAlloc<B: Backend> {
+    fn lwe_to_glwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: LWEToGLWESwitchingKeyToRef;
+}
+
+impl<B: Backend> LWEToGLWESwitchingKeyPrepareAlloc<B> for Module<B>
+where
+    Module<B>: LWEToGLWESwitchingKeyPrepare<B>,
+{
+    fn lwe_to_glwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: LWEToGLWESwitchingKeyToRef,
+    {
+        let mut ct_prep: LWEToGLWESwitchingKeyPrepared<Vec<u8>, B> = LWEToGLWESwitchingKeyPrepared::alloc(self, &other.to_ref());
+        self.lwe_to_glwe_switching_key_prepare(&mut ct_prep, other, scratch);
+        ct_prep
+    }
+}
+
+pub trait LWEToGLWESwitchingKeyPreparedToRef<B: Backend> {
+    fn to_ref(&self) -> LWEToGLWESwitchingKeyPrepared<&[u8], B>;
+}
+
+impl<D: DataRef, B: Backend> LWEToGLWESwitchingKeyPreparedToRef<B> for LWEToGLWESwitchingKeyPrepared<D, B>
+where
+    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToRef<B>,
+{
+    fn to_ref(&self) -> LWEToGLWESwitchingKeyPrepared<&[u8], B> {
+        LWEToGLWESwitchingKeyPrepared(self.0.to_ref())
+    }
+}
+
+pub trait LWEToGLWESwitchingKeyPreparedToMut<B: Backend> {
+    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyPrepared<&mut [u8], B>;
+}
+
+impl<D: DataMut, B: Backend> LWEToGLWESwitchingKeyPreparedToMut<B> for LWEToGLWESwitchingKeyPrepared<D, B>
+where
+    GLWESwitchingKeyPrepared<D, B>: GLWESwitchingKeyPreparedToMut<B>,
+{
+    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyPrepared<&mut [u8], B> {
+        LWEToGLWESwitchingKeyPrepared(self.0.to_mut())
     }
 }

poulpy-core/src/tests/test_suite/automorphism/gglwe_atk.rs

@@ -154,7 +154,7 @@ where
             );
 
             let mut auto_key_apply_prepared: AutomorphismKeyPrepared<Vec<u8>, B> =
-                AutomorphismKeyPrepared::alloc(module, &auto_key_apply_infos);
+                AutomorphismKeyPrepared::alloc_from_infos(module, &auto_key_apply_infos);
 
             auto_key_apply_prepared.prepare(module, &auto_key_apply, scratch.borrow());
 
@@ -348,7 +348,7 @@ where
             );
 
             let mut auto_key_apply_prepared: AutomorphismKeyPrepared<Vec<u8>, B> =
-                AutomorphismKeyPrepared::alloc(module, &auto_key_apply_layout);
+                AutomorphismKeyPrepared::alloc_from_infos(module, &auto_key_apply_layout);
 
             auto_key_apply_prepared.prepare(module, &auto_key_apply, scratch.borrow());
 

poulpy-core/src/tests/test_suite/automorphism/ggsw_ct.rs

@@ -179,7 +179,7 @@ where
             );
 
             let mut auto_key_prepared: AutomorphismKeyPrepared<Vec<u8>, B> =
-                AutomorphismKeyPrepared::alloc(module, &auto_key_layout);
+                AutomorphismKeyPrepared::alloc_from_infos(module, &auto_key_layout);
             auto_key_prepared.prepare(module, &auto_key, scratch.borrow());
 
             let mut tsk_prepared: TensorKeyPrepared<Vec<u8>, B> = TensorKeyPrepared::alloc(module, &tensor_key_layout);
@@ -359,7 +359,7 @@ where
             );
 
             let mut auto_key_prepared: AutomorphismKeyPrepared<Vec<u8>, B> =
-                AutomorphismKeyPrepared::alloc(module, &auto_key_layout);
+                AutomorphismKeyPrepared::alloc_from_infos(module, &auto_key_layout);
             auto_key_prepared.prepare(module, &auto_key, scratch.borrow());
 
             let mut tsk_prepared: TensorKeyPrepared<Vec<u8>, B> = TensorKeyPrepared::alloc(module, &tensor_key_layout);

poulpy-core/src/tests/test_suite/automorphism/glwe_ct.rs

@@ -141,7 +141,7 @@ where
             );
 
             let mut autokey_prepared: AutomorphismKeyPrepared<Vec<u8>, B> =
-                AutomorphismKeyPrepared::alloc(module, &autokey_infos);
+                AutomorphismKeyPrepared::alloc_from_infos(module, &autokey_infos);
             autokey_prepared.prepare(module, &autokey, scratch.borrow());
 
             ct_out.automorphism(module, &ct_in, &autokey_prepared, scratch.borrow());
@@ -274,7 +274,8 @@ where
                 scratch.borrow(),
             );
 
-            let mut autokey_prepared: AutomorphismKeyPrepared<Vec<u8>, B> = AutomorphismKeyPrepared::alloc(module, &autokey);
+            let mut autokey_prepared: AutomorphismKeyPrepared<Vec<u8>, B> =
+                AutomorphismKeyPrepared::alloc_from_infos(module, &autokey);
             autokey_prepared.prepare(module, &autokey, scratch.borrow());
 
             ct.automorphism_inplace(module, &autokey_prepared, scratch.borrow());