glwe swk + autokey compressed encrypt sk

poulpy-core/src/encryption/compressed/gglwe_atk.rs

@@ -1,66 +1,84 @@
 use poulpy_hal::{
-    api::{
+    api::{ModuleN, ScratchAvailable, VecZnxAutomorphism},
-        ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPPolAlloc, SvpPPolBytesOf, VecZnxAutomorphism, VecZnxDftBytesOf,
+    layouts::{Backend, DataMut, GaloisElement, Module, Scratch},
-        VecZnxNormalizeTmpBytes,
-    },
-    layouts::{Backend, DataMut, DataRef, Module, Scratch},
     source::Source,
 };
 
 use crate::{
     ScratchTakeCore,
-    encryption::compressed::gglwe_ksk::GGLWEKeyCompressedEncryptSk,
+    encryption::compressed::gglwe_ksk::GLWESwitchingKeyCompressedEncryptSk,
     layouts::{
-        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos,
+        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretAlloc, GLWESecretToRef, LWEInfos,
-        compressed::{AutomorphismKeyCompressed, AutomorphismKeyCompressedToMut, GLWESwitchingKeyCompressed},
+        compressed::{AutomorphismKeyCompressed, AutomorphismKeyCompressedToMut},
     },
 };
 
 impl AutomorphismKeyCompressed<Vec<u8>> {
-    pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
+    pub fn encrypt_sk_tmp_bytes<M, BE: Backend, A>(module: &M, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>:
+        M: AutomorphismKeyCompressedEncryptSk<BE>,
-            ModuleN + SvpPPolAlloc<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolBytesOf,
     {
-        assert_eq!(module.n() as u32, infos.n());
+        module.automorphism_key_compressed_encrypt_sk_tmp_bytes(infos)
-        GLWESwitchingKeyCompressed::encrypt_sk_tmp_bytes(module, infos) + GLWESecret::bytes_of(module, infos.rank_out())
     }
 }
 
-pub trait GGLWEAutomorphismKeyCompressedEncryptSk<B: Backend> {
+impl<DataSelf: DataMut> AutomorphismKeyCompressed<DataSelf> {
-    fn gglwe_automorphism_key_compressed_encrypt_sk<R, S>(
+    #[allow(clippy::too_many_arguments)]
+    pub fn encrypt_sk<M, S, BE: Backend>(
+        &mut self,
+        module: &M,
+        p: i64,
+        sk: &S,
+        seed_xa: [u8; 32],
+        source_xe: &mut Source,
+        scratch: &mut Scratch<BE>,
+    ) where
+        S: GLWESecretToRef,
+        M: AutomorphismKeyCompressedEncryptSk<BE>,
+    {
+        module.automorphism_key_compressed_encrypt_sk(self, p, sk, seed_xa, source_xe, scratch);
+    }
+}
+
+pub trait AutomorphismKeyCompressedEncryptSk<BE: Backend> {
+    fn automorphism_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos;
+
+    fn automorphism_key_compressed_encrypt_sk<R, S>(
         &self,
         res: &mut R,
         p: i64,
         sk: &S,
         seed_xa: [u8; 32],
         source_xe: &mut Source,
-        scratch: &mut Scratch<B>,
+        scratch: &mut Scratch<BE>,
     ) where
         R: AutomorphismKeyCompressedToMut,
         S: GLWESecretToRef;
 }
 
-impl<B: Backend> GGLWEAutomorphismKeyCompressedEncryptSk<B> for Module<B>
+impl<BE: Backend> AutomorphismKeyCompressedEncryptSk<BE> for Module<BE>
 where
-    Module<B>: ModuleN
+    Module<BE>: ModuleN + GaloisElement + VecZnxAutomorphism + GLWESwitchingKeyCompressedEncryptSk<BE>,
-        + GGLWEKeyCompressedEncryptSk<B>
+    Scratch<BE>: ScratchTakeCore<BE>,
-        + VecZnxNormalizeTmpBytes
-        + VecZnxDftBytesOf
-        + SvpPPolBytesOf
-        + VecZnxAutomorphism
-        + SvpPPolAlloc<B>,
-    Scratch<B>: ScratchAvailable + ScratchTakeBasic + ScratchTakeCore<B>,
 {
-    fn gglwe_automorphism_key_compressed_encrypt_sk<R, S>(
+    fn automorphism_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos,
+    {
+        self.glwe_switching_key_compressed_encrypt_sk_tmp_bytes(infos) + self.bytes_of_glwe_secret(infos.rank())
+    }
+
+    fn automorphism_key_compressed_encrypt_sk<R, S>(
         &self,
         res: &mut R,
         p: i64,
         sk: &S,
         seed_xa: [u8; 32],
         source_xe: &mut Source,
-        scratch: &mut Scratch<B>,
+        scratch: &mut Scratch<BE>,
     ) where
         R: AutomorphismKeyCompressedToMut,
         S: GLWESecretToRef,
@@ -68,8 +86,6 @@ where
         let res: &mut AutomorphismKeyCompressed<&mut [u8]> = &mut res.to_mut();
         let sk: &GLWESecret<&[u8]> = &sk.to_ref();
 
-        #[cfg(debug_assertions)]
-        {
         assert_eq!(res.n(), sk.n());
         assert_eq!(res.rank_out(), res.rank_in());
         assert_eq!(sk.rank(), res.rank_out());
@@ -78,13 +94,12 @@ where
             "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {}",
             scratch.available(),
             AutomorphismKeyCompressed::encrypt_sk_tmp_bytes(self, res)
-            )
+        );
-        }
 
         let (mut sk_out, scratch_1) = scratch.take_glwe_secret(self, sk.rank());
 
         {
-            (0..res.rank_out().into()).for_each(|i| {
+            for i in 0..res.rank_out().into(){
                 self.vec_znx_automorphism(
                     self.galois_element_inv(p),
                     &mut sk_out.data.as_vec_znx_mut(),
@@ -92,28 +107,11 @@ where
                     &sk.data.as_vec_znx(),
                     i,
                 );
-            });
+            };
         }
 
-        self.gglwe_key_compressed_encrypt_sk(&mut res.key, sk, &sk_out, seed_xa, source_xe, scratch_1);
+        self.glwe_switching_key_compressed_encrypt_sk(&mut res.key, sk, &sk_out, seed_xa, source_xe, scratch_1);
 
         res.p = p;
     }
 }
-
-impl<DataSelf: DataMut> AutomorphismKeyCompressed<DataSelf> {
-    #[allow(clippy::too_many_arguments)]
-    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
-        &mut self,
-        module: &Module<B>,
-        p: i64,
-        sk: &GLWESecret<DataSk>,
-        seed_xa: [u8; 32],
-        source_xe: &mut Source,
-        scratch: &mut Scratch<B>,
-    ) where
-        Module<B>: GGLWEAutomorphismKeyCompressedEncryptSk<B>,
-    {
-        module.gglwe_automorphism_key_compressed_encrypt_sk(self, p, sk, seed_xa, source_xe, scratch);
-    }
-}

poulpy-core/src/encryption/compressed/gglwe_ksk.rs

@@ -1,9 +1,6 @@
 use poulpy_hal::{
-    api::{
+    api::{ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPrepare, VecZnxSwitchRing},
-        ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxDftBytesOf,
+    layouts::{Backend, DataMut, Module, ScalarZnx, Scratch},
-        VecZnxNormalizeTmpBytes, VecZnxSwitchRing,
-    },
-    layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
     source::Source,
 };
 
@@ -11,119 +8,118 @@ use crate::{
     ScratchTakeCore,
     encryption::compressed::gglwe_ct::GGLWECompressedEncryptSk,
     layouts::{
-        GGLWE, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos, RingDegree,
+        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos,
         compressed::{GLWESwitchingKeyCompressed, GLWESwitchingKeyCompressedToMut},
-        prepared::GLWESecretPrepared,
+        prepared::{GLWESecretPrepare, GLWESecretPrepared, GLWESecretPreparedAlloc},
     },
 };
 
 impl GLWESwitchingKeyCompressed<Vec<u8>> {
-    pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
+    pub fn encrypt_sk_tmp_bytes<M, A, BE: Backend>(module: &M, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>:
+        M: GLWESwitchingKeyCompressedEncryptSk<BE>,
-            ModuleN + SvpPPolAlloc<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolBytesOf,
     {
-        (GGLWE::encrypt_sk_tmp_bytes(module, infos) | ScalarZnx::bytes_of(module.n(), 1))
+        module.glwe_switching_key_compressed_encrypt_sk_tmp_bytes(infos)
-            + ScalarZnx::bytes_of(module.n(), infos.rank_in().into())
-            + GLWESecretPrepared::bytes_of(module, infos.rank_out())
     }
 }
 
-impl<DataSelf: DataMut> GLWESwitchingKeyCompressed<DataSelf> {
+impl<D: DataMut> GLWESwitchingKeyCompressed<D> {
     #[allow(clippy::too_many_arguments)]
-    pub fn encrypt_sk<DataSkIn: DataRef, DataSkOut: DataRef, B: Backend>(
+    pub fn encrypt_sk<M, S1, S2, BE: Backend>(
         &mut self,
-        module: &Module<B>,
+        module: &M,
-        sk_in: &GLWESecret<DataSkIn>,
+        sk_in: &S1,
-        sk_out: &GLWESecret<DataSkOut>,
+        sk_out: &S2,
         seed_xa: [u8; 32],
         source_xe: &mut Source,
-        scratch: &mut Scratch<B>,
+        scratch: &mut Scratch<BE>,
     ) where
-        Module<B>: GGLWEKeyCompressedEncryptSk<B>,
+        S1: GLWESecretToRef,
+        S2: GLWESecretToRef,
+        M: GLWESwitchingKeyCompressedEncryptSk<BE>,
     {
-        module.gglwe_key_compressed_encrypt_sk(self, sk_in, sk_out, seed_xa, source_xe, scratch);
+        module.glwe_switching_key_compressed_encrypt_sk(self, sk_in, sk_out, seed_xa, source_xe, scratch);
     }
 }
 
-pub trait GGLWEKeyCompressedEncryptSk<B: Backend> {
+pub trait GLWESwitchingKeyCompressedEncryptSk<BE: Backend> {
-    fn gglwe_key_compressed_encrypt_sk<R, SI, SO>(
+    fn glwe_switching_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos;
+
+    fn glwe_switching_key_compressed_encrypt_sk<R, S1, S2>(
         &self,
         res: &mut R,
-        sk_in: &SI,
+        sk_in: &S1,
-        sk_out: &SO,
+        sk_out: &S2,
         seed_xa: [u8; 32],
         source_xe: &mut Source,
-        scratch: &mut Scratch<B>,
+        scratch: &mut Scratch<BE>,
     ) where
         R: GLWESwitchingKeyCompressedToMut,
-        SI: GLWESecretToRef,
+        S1: GLWESecretToRef,
-        SO: GLWESecretToRef;
+        S2: GLWESecretToRef;
 }
 
-impl<B: Backend> GGLWEKeyCompressedEncryptSk<B> for Module<B>
+impl<BE: Backend> GLWESwitchingKeyCompressedEncryptSk<BE> for Module<BE>
 where
-    Module<B>: ModuleN
+    Module<BE>: ModuleN + GGLWECompressedEncryptSk<BE> + GLWESecretPreparedAlloc<BE> + GLWESecretPrepare<BE> + VecZnxSwitchRing,
-        + GGLWECompressedEncryptSk<B>
+    Scratch<BE>: ScratchTakeCore<BE>,
-        + SvpPPolBytesOf
-        + VecZnxNormalizeTmpBytes
-        + VecZnxDftBytesOf
-        + VecZnxSwitchRing
-        + SvpPrepare<B>
-        + SvpPPolAlloc<B>,
-    Scratch<B>: ScratchAvailable + ScratchTakeBasic + ScratchTakeCore<B>,
 {
-    fn gglwe_key_compressed_encrypt_sk<R, SI, SO>(
+    fn glwe_switching_key_compressed_encrypt_sk_tmp_bytes<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos,
+    {
+        self.gglwe_compressed_encrypt_sk_tmp_bytes(infos)
+            .max(ScalarZnx::bytes_of(self.n(), 1))
+            + ScalarZnx::bytes_of(self.n(), infos.rank_in().into())
+            + GLWESecretPrepared::bytes_of(self, infos.rank_out())
+    }
+
+    fn glwe_switching_key_compressed_encrypt_sk<R, S1, S2>(
         &self,
         res: &mut R,
-        sk_in: &SI,
+        sk_in: &S1,
-        sk_out: &SO,
+        sk_out: &S2,
         seed_xa: [u8; 32],
         source_xe: &mut Source,
-        scratch: &mut Scratch<B>,
+        scratch: &mut Scratch<BE>,
     ) where
         R: GLWESwitchingKeyCompressedToMut,
-        SI: GLWESecretToRef,
+        S1: GLWESecretToRef,
-        SO: GLWESecretToRef,
+        S2: GLWESecretToRef,
     {
         let res: &mut GLWESwitchingKeyCompressed<&mut [u8]> = &mut res.to_mut();
         let sk_in: &GLWESecret<&[u8]> = &sk_in.to_ref();
         let sk_out: &GLWESecret<&[u8]> = &sk_out.to_ref();
 
-        #[cfg(debug_assertions)]
-        {
-            use crate::layouts::GLWESwitchingKey;
-
         assert!(sk_in.n().0 <= self.n() as u32);
         assert!(sk_out.n().0 <= self.n() as u32);
         assert!(
-                scratch.available() >= GLWESwitchingKey::encrypt_sk_tmp_bytes(self, res),
+            scratch.available() >= self.gglwe_compressed_encrypt_sk_tmp_bytes(res),
             "scratch.available()={} < GLWESwitchingKey::encrypt_sk_tmp_bytes={}",
             scratch.available(),
-                GLWESwitchingKey::encrypt_sk_tmp_bytes(self, res)
+            self.gglwe_compressed_encrypt_sk_tmp_bytes(res)
-            )
+        );
-        }
-
-        let n: usize = sk_in.n().max(sk_out.n()).into();
 
         let (mut sk_in_tmp, scratch_1) = scratch.take_scalar_znx(self, sk_in.rank().into());
-        (0..sk_in.rank().into()).for_each(|i| {
+        for i in 0..sk_in.rank().into() {
             self.vec_znx_switch_ring(
                 &mut sk_in_tmp.as_vec_znx_mut(),
                 i,
                 &sk_in.data.as_vec_znx(),
                 i,
             );
-        });
+        }
 
         let (mut sk_out_tmp, scratch_2) = scratch_1.take_glwe_secret_prepared(self, sk_out.rank());
         {
             let (mut tmp, _) = scratch_2.take_scalar_znx(self, 1);
-            (0..sk_out.rank().into()).for_each(|i| {
+            for i in 0..sk_out.rank().into() {
                 self.vec_znx_switch_ring(&mut tmp.as_vec_znx_mut(), 0, &sk_out.data.as_vec_znx(), i);
                 self.svp_prepare(&mut sk_out_tmp.data, i, &tmp, 0);
-            });
+            }
         }
 
         self.gglwe_compressed_encrypt_sk(

poulpy-core/src/encryption/compressed/gglwe_tsk.rs

@@ -10,7 +10,7 @@ use poulpy_hal::{
 
 use crate::{
     ScratchTakeCore,
-    encryption::compressed::gglwe_ksk::GGLWEKeyCompressedEncryptSk,
+    encryption::compressed::gglwe_ksk::GLWESwitchingKeyCompressedEncryptSk,
     layouts::{
         GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GetDist, LWEInfos, Rank, TensorKey,
         compressed::{TensorKeyCompressed, TensorKeyCompressedToMut},
@@ -49,7 +49,7 @@ pub trait GGLWETensorKeyCompressedEncryptSk<B: Backend> {
 impl<B: Backend> GGLWETensorKeyCompressedEncryptSk<B> for Module<B>
 where
     Module<B>: ModuleN
-        + GGLWEKeyCompressedEncryptSk<B>
+        + GLWESwitchingKeyCompressedEncryptSk<B>
         + VecZnxDftApply<B>
         + SvpApplyDftToDft<B>
         + VecZnxIdftApplyTmpA<B>
@@ -119,7 +119,7 @@ where
 
                 let (seed_xa_tmp, _) = source_xa.branch();
 
-                self.gglwe_key_compressed_encrypt_sk(
+                self.glwe_switching_key_compressed_encrypt_sk(
                     res.at_mut(i, j),
                     &sk_ij,
                     sk,