wip

poulpy-core/src/encryption/gglwe_atk.rs

@@ -1,34 +1,33 @@
 use poulpy_hal::{
     api::{
-        ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeVecZnx, TakeVecZnxDft,
-        VecZnxAddInplace, VecZnxAddNormal, VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftAllocBytes,
-        VecZnxDftApply, VecZnxFillUniform, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace,
-        VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace, VecZnxSwitchRing,
+        ScratchAvailable, SvpApplyDftToDftInplace, SvpPPolBytesOf, SvpPrepare, VecZnxAddInplace, VecZnxAddNormal,
+        VecZnxAddScalarInplace, VecZnxAutomorphism, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf, VecZnxFillUniform,
+        VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubInplace,
+        VecZnxSwitchRing,
     },
-    layouts::{Backend, DataMut, DataRef, Module, Scratch},
+    layouts::{Backend, DataMut, Module, Scratch},
     source::Source,
 };
 
-use crate::{
-    TakeGLWESecret, TakeGLWESecretPrepared,
-    layouts::{GGLWEAutomorphismKey, GGLWEInfos, GGLWESwitchingKey, GLWEInfos, GLWESecret, LWEInfos},
+use crate::layouts::{
+    AutomorphismKey, AutomorphismKeyToMut, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, GLWESwitchingKey, LWEInfos,
 };
 
-impl GGLWEAutomorphismKey<Vec<u8>> {
-    pub fn encrypt_sk_scratch_space<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
+impl AutomorphismKey<Vec<u8>> {
+    pub fn encrypt_sk_tmp_bytes<BE: Backend, A>(module: &Module<BE>, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>: SvpPPolAllocBytes + VecZnxNormalizeTmpBytes + VecZnxDftAllocBytes + VecZnxNormalizeTmpBytes,
+        Module<BE>: SvpPPolBytesOf + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes,
     {
         assert_eq!(
             infos.rank_in(),
             infos.rank_out(),
             "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
         );
-        GGLWESwitchingKey::encrypt_sk_scratch_space(module, infos) + GLWESecret::alloc_bytes(&infos.glwe_layout())
+        GLWESwitchingKey::encrypt_sk_tmp_bytes(module, infos) + GLWESecret::bytes_of_from_infos(module, &infos.glwe_layout())
     }
 
-    pub fn encrypt_pk_scratch_space<B: Backend, A>(module: &Module<B>, _infos: &A) -> usize
+    pub fn encrypt_pk_tmp_bytes<BE: Backend, A>(module: &Module<BE>, _infos: &A) -> usize
     where
         A: GGLWEInfos,
     {
@@ -37,62 +36,102 @@ impl GGLWEAutomorphismKey<Vec<u8>> {
             _infos.rank_out(),
             "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
         );
-        GGLWESwitchingKey::encrypt_pk_scratch_space(module, _infos)
+        GLWESwitchingKey::encrypt_pk_tmp_bytes(module, _infos)
     }
 }
 
-impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
-    #[allow(clippy::too_many_arguments)]
-    pub fn encrypt_sk<DataSk: DataRef, B: Backend>(
-        &mut self,
-        module: &Module<B>,
+pub trait GGLWEAutomorphismKeyEncryptSk<BE: Backend> {
+    fn gglwe_automorphism_key_encrypt_sk<A, B>(
+        &self,
+        res: &mut A,
         p: i64,
-        sk: &GLWESecret<DataSk>,
+        sk: &B,
         source_xa: &mut Source,
         source_xe: &mut Source,
-        scratch: &mut Scratch<B>,
+        scratch: &mut Scratch<BE>,
     ) where
-        Module<B>: VecZnxAddScalarInplace
-            + VecZnxDftAllocBytes
-            + VecZnxBigNormalize<B>
-            + VecZnxDftApply<B>
-            + SvpApplyDftToDftInplace<B>
-            + VecZnxIdftApplyConsume<B>
-            + VecZnxNormalizeTmpBytes
-            + VecZnxFillUniform
-            + VecZnxSubInplace
-            + VecZnxAddInplace
-            + VecZnxNormalizeInplace<B>
-            + VecZnxAddNormal
-            + VecZnxNormalize<B>
-            + VecZnxSub
-            + SvpPrepare<B>
-            + VecZnxSwitchRing
-            + SvpPPolAllocBytes
-            + VecZnxAutomorphism,
-        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx + TakeScalarZnx + TakeGLWESecretPrepared<B>,
+        A: AutomorphismKeyToMut,
+        B: GLWESecretToRef;
+}
+
+impl<DM: DataMut> AutomorphismKey<DM>
+where
+    Self: AutomorphismKeyToMut,
+{
+    pub fn encrypt_sk<S, BE: Backend>(
+        &mut self,
+        module: &Module<BE>,
+        p: i64,
+        sk: &S,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        scratch: &mut Scratch<BE>,
+    ) where
+        S: GLWESecretToRef,
+        Module<BE>: GGLWEAutomorphismKeyEncryptSk<BE>,
     {
+        module.gglwe_automorphism_key_encrypt_sk(self, p, sk, source_xa, source_xe, scratch);
+    }
+}
+
+impl<BE: Backend> GGLWEAutomorphismKeyEncryptSk<BE> for Module<BE>
+where
+    Module<BE>: VecZnxAddScalarInplace
+        + VecZnxDftBytesOf
+        + VecZnxBigNormalize<BE>
+        + VecZnxDftApply<BE>
+        + SvpApplyDftToDftInplace<BE>
+        + VecZnxIdftApplyConsume<BE>
+        + VecZnxNormalizeTmpBytes
+        + VecZnxFillUniform
+        + VecZnxSubInplace
+        + VecZnxAddInplace
+        + VecZnxNormalizeInplace<BE>
+        + VecZnxAddNormal
+        + VecZnxNormalize<BE>
+        + VecZnxSub
+        + SvpPrepare<BE>
+        + VecZnxSwitchRing
+        + SvpPPolBytesOf
+        + VecZnxAutomorphism,
+    Scratch<BE>: ScratchAvailable,
+{
+    fn gglwe_automorphism_key_encrypt_sk<A, B>(
+        &self,
+        res: &mut A,
+        p: i64,
+        sk: &B,
+        source_xa: &mut Source,
+        source_xe: &mut Source,
+        scratch: &mut Scratch<BE>,
+    ) where
+        A: AutomorphismKeyToMut,
+        B: GLWESecretToRef,
+    {
+        let res: &mut AutomorphismKey<&mut [u8]> = &mut res.to_mut();
+        let sk: &GLWESecret<&[u8]> = &sk.to_ref();
+
         #[cfg(debug_assertions)]
         {
             use crate::layouts::{GLWEInfos, LWEInfos};
 
-            assert_eq!(self.n(), sk.n());
-            assert_eq!(self.rank_out(), self.rank_in());
-            assert_eq!(sk.rank(), self.rank_out());
+            assert_eq!(res.n(), sk.n());
+            assert_eq!(res.rank_out(), res.rank_in());
+            assert_eq!(sk.rank(), res.rank_out());
             assert!(
-                scratch.available() >= GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, self),
-                "scratch.available(): {} < AutomorphismKey::encrypt_sk_scratch_space: {:?}",
+                scratch.available() >= AutomorphismKey::encrypt_sk_tmp_bytes(self, res),
+                "scratch.available(): {} < AutomorphismKey::encrypt_sk_tmp_bytes: {:?}",
                 scratch.available(),
-                GGLWEAutomorphismKey::encrypt_sk_scratch_space(module, self)
+                AutomorphismKey::encrypt_sk_tmp_bytes(self, res)
             )
         }
 
         let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
 
         {
-            (0..self.rank_out().into()).for_each(|i| {
-                module.vec_znx_automorphism(
-                    module.galois_element_inv(p),
+            (0..res.rank_out().into()).for_each(|i| {
+                self.vec_znx_automorphism(
+                    self.galois_element_inv(p),
                     &mut sk_out.data.as_vec_znx_mut(),
                     i,
                     &sk.data.as_vec_znx(),
@@ -101,9 +140,9 @@ impl<DataSelf: DataMut> GGLWEAutomorphismKey<DataSelf> {
             });
         }
 
-        self.key
-            .encrypt_sk(module, sk, &sk_out, source_xa, source_xe, scratch_1);
+        res.key
+            .encrypt_sk(self, sk, &sk_out, source_xa, source_xe, scratch_1);
 
-        self.p = p;
+        res.p = p;
     }
 }