M poulpy-core/src/layouts/prepared/glwe_pk.rs → poulpy-core/src/layouts/prepared/glwe_pk.rs

@@ -6,7 +6,7 @@ };
 
 use crate::{
     dist::Distribution,
-    layouts::{Base2K, BuildError, Degree, GLWEInfos, GLWEPublicKey, LWEInfos, Rank, TorusPrecision},
+    layouts::{Base2K, BuildError, Degree, GLWEInfos, GLWEPublicKey, GLWEPublicKeyToRef, LWEInfos, Rank, TorusPrecision},
 };
 
 #[derive(PartialEq, Eq)]
@@ -15,6 +15,16 @@ pub(crate) data: VecZnxDft<D, B>,
     pub(crate) base2k: Base2K,
     pub(crate) k: TorusPrecision,
     pub(crate) dist: Distribution,
+}
+
+pub(crate) trait SetDist {
+    fn set_dist(&mut self, dist: Distribution);
+}
+
+impl<D: Data, B: Backend> SetDist for GLWEPublicKeyPrepared<D, B> {
+    fn set_dist(&mut self, dist: Distribution) {
+        self.dist = dist
+    }
 }
 
 impl<D: Data, B: Backend> LWEInfos for GLWEPublicKeyPrepared<D, B> {
@@ -166,40 +176,104 @@ module.vec_znx_dft_alloc_bytes((rank + 1).into(), k.0.div_ceil(base2k.0) as usize)
     }
 }
 
-impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWEPublicKeyPrepared<Vec<u8>, B>> for GLWEPublicKey<D>
+pub trait GLWEPublicKeyPrepareTmpBytes {
+    fn glwe_public_key_prepare_tmp_bytes<A>(&self, infos: &A)
+    where
+        A: GLWEInfos;
+}
+
+impl<B: Backend> GLWEPublicKeyPrepareTmpBytes for Module<B> {
+    fn glwe_public_key_prepare_tmp_bytes<A>(&self, infos: &A)
+    where
+        A: GLWEInfos,
+    {
+        0
+    }
+}
+
+impl<B: Backend> GLWEPublicKeyPrepared<Vec<u8>, B> {
+    pub fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A)
+    where
+        A: GLWEInfos,
+        Module<B>: GLWEPublicKeyPrepareTmpBytes,
+    {
+        module.glwe_public_key_prepare_tmp_bytes(infos);
+    }
+}
+
+pub trait GLWEPublicKeyPrepare<B: Backend> {
+    fn glwe_public_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
+    where
+        R: GLWEPublicKeyPreparedToMut<B> + SetDist,
+        O: GLWEPublicKeyToRef;
+}
+
+impl<B: Backend> GLWEPublicKeyPrepare<B> for Module<B>
 where
     Module<B>: VecZnxDftAlloc<B> + VecZnxDftApply<B>,
 {
-    fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GLWEPublicKeyPrepared<Vec<u8>, B> {
-        let mut pk_prepared: GLWEPublicKeyPrepared<Vec<u8>, B> = GLWEPublicKeyPrepared::alloc(module, self);
-        pk_prepared.prepare(module, self, scratch);
-        pk_prepared
+    fn glwe_public_key_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &Scratch<B>)
+    where
+        R: GLWEPublicKeyPreparedToMut<B> + SetDist,
+        O: GLWEPublicKeyToRef,
+    {
+        {
+            let res: GLWEPublicKeyPrepared<&mut [u8], B> = res.to_mut();
+            let other: GLWEPublicKey<&[u8]> = other.to_ref();
+
+            assert_eq!(res.n(), self.n() as u32);
+            assert_eq!(other.n(), self.n() as u32);
+            assert_eq!(res.size(), other.size());
+            assert_eq!(res.k(), other.k());
+            assert_eq!(res.base2k(), other.base2k());
+
+            for i in 0..(self.rank() + 1).into() {
+                self.vec_znx_dft_apply(1, 0, &mut self.data, i, &other.data, i);
+            }
+        }
+
+        res.set_dist(other.dist);
     }
 }
 
-impl<B: Backend, A: GLWEInfos> PrepareScratchSpace<B, A> for GLWEPublicKeyPrepared<Vec<u8>, B> {
-    fn prepare_scratch_space(_module: &Module<B>, _infos: &A) -> usize {
-        0
+impl<D: DataMut, B: Backend> GLWEPublicKeyPrepared<D, B>
+where
+    Module<B>: GLWEPublicKeyPrepare<B>,
+{
+    pub fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: GLWEPublicKeyToRef,
+    {
+        module.glwe_public_key_prepare(self, other, scratch);
     }
 }
 
-impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWEPublicKey<DR>> for GLWEPublicKeyPrepared<DM, B>
+pub trait GLWEPublicKeyPrepareAlloc<B: Backend> {
+    fn glwe_public_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: GLWEPublicKeyToRef;
+}
+
+impl<B: Backend> GLWEPublicKeyPrepareAlloc<B> for Module<B>
 where
-    Module<B>: VecZnxDftApply<B>,
+    Module<B>: GLWEPublicKeyPrepare<B>,
 {
-    fn prepare(&mut self, module: &Module<B>, other: &GLWEPublicKey<DR>, _scratch: &mut Scratch<B>) {
-        #[cfg(debug_assertions)]
-        {
-            assert_eq!(self.n(), other.n());
-            assert_eq!(self.size(), other.size());
-        }
+    fn glwe_public_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: GLWEPublicKeyToRef,
+    {
+        let mut ct_prepared: GLWEPublicKeyPrepared<Vec<u8>, B> = GLWEPublicKeyPrepared::alloc(self, other);
+        self.glwe_public_key_prepare(&mut ct_prepared, ct_prepared, scratch);
+        ct_prepared
+    }
+}
 
-        (0..(self.rank() + 1).into()).for_each(|i| {
-            module.vec_znx_dft_apply(1, 0, &mut self.data, i, &other.data, i);
-        });
-        self.k = other.k();
-        self.base2k = other.base2k();
-        self.dist = other.dist;
+impl<D: DataRef> GLWEPublicKey<D> {
+    pub fn prepare_alloc<B: Backend>(&self, module: &Module<B>, scratch: &mut Scratch<B>)
+    where
+        Module<B>: GLWEPublicKeyPrepareAlloc<B>,
+    {
+        module.glwe_public_key_prepare_alloc(self, scratch);
     }
 }
 

M poulpy-core/src/layouts/prepared/glwe_sk.rs → poulpy-core/src/layouts/prepared/glwe_sk.rs

@@ -5,12 +5,21 @@ };
 
 use crate::{
     dist::Distribution,
-    layouts::{Base2K, Degree, GLWEInfos, GLWESecret, LWEInfos, Rank, TorusPrecision},
+    layouts::{
+        Base2K, Degree, GLWEInfos, GLWESecret, GLWESecretToMut, GLWESecretToRef, LWEInfos, Rank, TorusPrecision,
+        prepared::SetDist,
+    },
 };
 
 pub struct GLWESecretPrepared<D: Data, B: Backend> {
     pub(crate) data: SvpPPol<D, B>,
     pub(crate) dist: Distribution,
+}
+
+impl<D: DataRef, B: Backend> SetDist for GLWESecretPrepared<D, B> {
+    fn set_dist(&mut self, dist: Distribution) {
+        self.dist = dist
+    }
 }
 
 impl<D: Data, B: Backend> LWEInfos for GLWESecretPrepared<D, B> {
@@ -82,32 +91,79 @@ Rank(self.data.cols() as u32)
     }
 }
 
-impl<B: Backend, A: GLWEInfos> PrepareScratchSpace<B, A> for GLWESecretPrepared<Vec<u8>, B> {
-    fn prepare_scratch_space(_module: &Module<B>, _infos: &A) -> usize {
+pub trait GLWESecretPrepareTmpBytes {
+    fn glwe_secret_prepare_tmp_bytes<A>(&self, infos: &A)
+    where
+        A: GLWEInfos;
+}
+
+impl<B: Backend> GLWESecretPrepareTmpBytes for Module<B> {
+    fn glwe_secret_prepare_tmp_bytes<A>(&self, infos: &A)
+    where
+        A: GLWEInfos,
+    {
         0
     }
 }
 
-impl<D: DataRef, B: Backend> PrepareAlloc<B, GLWESecretPrepared<Vec<u8>, B>> for GLWESecret<D>
+impl<B: Backend> GLWESecretPrepared<Vec<u8>, B>
 where
-    Module<B>: SvpPrepare<B> + SvpPPolAlloc<B>,
+    Module<B>: GLWESecretPrepareTmpBytes,
 {
-    fn prepare_alloc(&self, module: &Module<B>, _scratch: &mut Scratch<B>) -> GLWESecretPrepared<Vec<u8>, B> {
-        let mut sk_dft: GLWESecretPrepared<Vec<u8>, B> = GLWESecretPrepared::alloc(module, self);
-        sk_dft.prepare(module, self, _scratch);
-        sk_dft
+    fn prepare_tmp_bytes<A>(&self, module: &Module<B>, infos: &A) -> usize
+    where
+        A: GLWEInfos,
+    {
+        0
     }
 }
 
-impl<DM: DataMut, DR: DataRef, B: Backend> Prepare<B, GLWESecret<DR>> for GLWESecretPrepared<DM, B>
+pub trait GLWESecretPrepare<B: Backend> {
+    fn glwe_secret_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
+    where
+        R: GLWESecretPreparedToMut<B> + SetDist,
+        O: GLWESecretToRef;
+}
+
+impl<B: Backend> GLWESecretPrepare<B> for Module<B>
 where
     Module<B>: SvpPrepare<B>,
 {
-    fn prepare(&mut self, module: &Module<B>, other: &GLWESecret<DR>, _scratch: &mut Scratch<B>) {
-        (0..self.rank().into()).for_each(|i| {
-            module.svp_prepare(&mut self.data, i, &other.data, i);
-        });
-        self.dist = other.dist
+    fn glwe_secret_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
+    where
+        R: GLWESecretPreparedToMut<B> + SetDist,
+        O: GLWESecretToRef,
+    {
+        {
+            let res: GLWESecretPrepared<&mut [u8], _> = res.to_mut();
+            let other: GLWESecret<&[u8]> = other.to_ref();
+
+            for i in 0..self.rank().into() {
+                self.svp_prepare(&mut res.data, i, &other.data, i);
+            }
+        }
+
+        res.set_dist(other.dist);
+    }
+}
+
+pub trait GLWESecretPrepareAlloc<B: Backend> {
+    fn glwe_secret_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: GLWESecretToMut;
+}
+
+impl<B: Backend> GLWESecretPrepareAlloc<B> for Module<B>
+where
+    Module<B>: GLWESecretPrepare<B>,
+{
+    fn glwe_secret_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
+    where
+        O: GLWESecretToMut,
+    {
+        let mut ct_prep: GLWESecretPrepared<Vec<u8>, B> = GLWESecretPrepared::alloc(self, self);
+        self.glwe_secret_prepare(&mut ct_prep, other, scratch);
+        ct_prep
     }
 }