wip all of Encryption

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

@@ -1,10 +1,11 @@
 use poulpy_hal::{
-    api::{ScratchAvailable, SvpPPolBytesOf, VecZnxAutomorphism, VecZnxDftBytesOf, VecZnxNormalizeTmpBytes},
+    api::{ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPPolAlloc, SvpPPolBytesOf, VecZnxAutomorphism, VecZnxDftBytesOf, VecZnxNormalizeTmpBytes},
     layouts::{Backend, DataMut, DataRef, Module, Scratch},
     source::Source,
 };
 
 use crate::{
+    ScratchTakeCore,
     encryption::compressed::gglwe_ksk::GGLWEKeyCompressedEncryptSk,
     layouts::{
         GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos,
@@ -16,10 +17,10 @@ impl AutomorphismKeyCompressed<Vec<u8>> {
     pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolBytesOf,
+        Module<B>: ModuleN + SvpPPolAlloc<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolBytesOf,
     {
         assert_eq!(module.n() as u32, infos.n());
-        GLWESwitchingKeyCompressed::encrypt_sk_tmp_bytes(module, infos) + GLWESecret::bytes_of(infos.n(), infos.rank_out())
+        GLWESwitchingKeyCompressed::encrypt_sk_tmp_bytes(module, infos) + GLWESecret::bytes_of(module, infos.rank_out())
     }
 }
 
@@ -39,8 +40,14 @@ pub trait GGLWEAutomorphismKeyCompressedEncryptSk<B: Backend> {
 
 impl<B: Backend> GGLWEAutomorphismKeyCompressedEncryptSk<B> for Module<B>
 where
-    Module<B>: GGLWEKeyCompressedEncryptSk<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + SvpPPolBytesOf + VecZnxAutomorphism,
-    Scratch<B>: ScratchAvailable,
+    Module<B>: ModuleN
+        + GGLWEKeyCompressedEncryptSk<B>
+        + VecZnxNormalizeTmpBytes
+        + VecZnxDftBytesOf
+        + SvpPPolBytesOf
+        + VecZnxAutomorphism
+        + SvpPPolAlloc<B>,
+    Scratch<B>: ScratchAvailable + ScratchTakeBasic + ScratchTakeCore<B>,
 {
     fn gglwe_automorphism_key_compressed_encrypt_sk<R, S>(
         &self,
@@ -70,7 +77,7 @@ where
             )
         }
 
-        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(sk.n(), sk.rank());
+        let (mut sk_out, scratch_1) = scratch.take_glwe_secret(self, sk.rank());
 
         {
             (0..res.rank_out().into()).for_each(|i| {

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

@@ -1,6 +1,6 @@
 use poulpy_hal::{
     api::{
-        ScratchAvailable, VecZnxAddScalarInplace, VecZnxDftBytesOf, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes,
+        ModuleN, ScratchAvailable, VecZnxAddScalarInplace, VecZnxDftBytesOf, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes,
         ZnNormalizeInplace,
     },
     layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxZero},
@@ -8,6 +8,7 @@ use poulpy_hal::{
 };
 
 use crate::{
+    ScratchTakeCore,
     encryption::{SIGMA, glwe_ct::GLWEEncryptSkInternal},
     layouts::{
         GGLWE, GGLWEInfos, LWEInfos,
@@ -60,13 +61,14 @@ pub trait GGLWECompressedEncryptSk<B: Backend> {
 
 impl<B: Backend> GGLWECompressedEncryptSk<B> for Module<B>
 where
-    Module<B>: GLWEEncryptSkInternal<B>
+    Module<B>: ModuleN
+        + GLWEEncryptSkInternal<B>
         + VecZnxNormalizeInplace<B>
         + VecZnxNormalizeTmpBytes
         + VecZnxDftBytesOf
         + VecZnxAddScalarInplace
         + ZnNormalizeInplace<B>,
-    Scratch<B>: ScratchAvailable,
+    Scratch<B>: ScratchAvailable + ScratchTakeCore<B>,
 {
     fn gglwe_compressed_encrypt_sk<R, P, S>(
         &self,
@@ -130,7 +132,7 @@ where
 
         let mut source_xa = Source::new(seed);
 
-        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(res);
+        let (mut tmp_pt, scrach_1) = scratch.take_glwe_pt(self, res);
         (0..rank_in).for_each(|col_i| {
             (0..dnum).for_each(|d_i| {
                 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt

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

@@ -1,10 +1,11 @@
 use poulpy_hal::{
-    api::{ScratchAvailable, SvpPPolBytesOf, SvpPrepare, VecZnxDftBytesOf, VecZnxNormalizeTmpBytes, VecZnxSwitchRing},
+    api::{ModuleN, ScratchAvailable, ScratchTakeBasic, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxDftBytesOf, VecZnxNormalizeTmpBytes, VecZnxSwitchRing},
     layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch},
     source::Source,
 };
 
 use crate::{
+    ScratchTakeCore,
     encryption::compressed::gglwe_ct::GGLWECompressedEncryptSk,
     layouts::{
         GGLWE, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos, RingDegree,
@@ -17,7 +18,7 @@ impl GLWESwitchingKeyCompressed<Vec<u8>> {
     pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>: VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolBytesOf,
+        Module<B>: ModuleN + SvpPPolAlloc<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + SvpPPolBytesOf,
     {
         (GGLWE::encrypt_sk_tmp_bytes(module, infos) | ScalarZnx::bytes_of(module.n(), 1))
             + ScalarZnx::bytes_of(module.n(), infos.rank_in().into())
@@ -59,13 +60,15 @@ pub trait GGLWEKeyCompressedEncryptSk<B: Backend> {
 
 impl<B: Backend> GGLWEKeyCompressedEncryptSk<B> for Module<B>
 where
-    Module<B>: GGLWECompressedEncryptSk<B>
+    Module<B>: ModuleN
+        + GGLWECompressedEncryptSk<B>
         + SvpPPolBytesOf
         + VecZnxNormalizeTmpBytes
         + VecZnxDftBytesOf
         + VecZnxSwitchRing
-        + SvpPrepare<B>,
-    Scratch<B>: ScratchAvailable,
+        + SvpPrepare<B>
+        + SvpPPolAlloc<B>,
+    Scratch<B>: ScratchAvailable + ScratchTakeBasic + ScratchTakeCore<B>,
 {
     fn gglwe_key_compressed_encrypt_sk<R, SI, SO>(
         &self,
@@ -100,7 +103,7 @@ where
 
         let n: usize = sk_in.n().max(sk_out.n()).into();
 
-        let (mut sk_in_tmp, scratch_1) = scratch.take_scalar_znx(n, sk_in.rank().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| {
             self.vec_znx_switch_ring(
                 &mut sk_in_tmp.as_vec_znx_mut(),
@@ -110,9 +113,9 @@ where
             );
         });
 
-        let (mut sk_out_tmp, scratch_2) = scratch_1.take_glwe_secret_prepared(RingDegree(n as u32), sk_out.rank());
+        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(n, 1);
+            let (mut tmp, _) = scratch_2.take_scalar_znx(self, 1);
             (0..sk_out.rank().into()).for_each(|i| {
                 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);

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

@@ -1,16 +1,18 @@
 use poulpy_hal::{
     api::{
-        SvpApplyDftToDft, SvpPPolBytesOf, SvpPrepare, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf,
+        ModuleN, ScratchTakeBasic, SvpApplyDftToDft, SvpPPolAlloc, SvpPPolBytesOf, SvpPrepare, VecZnxBigBytesOf, VecZnxBigNormalize, VecZnxDftApply, VecZnxDftBytesOf,
         VecZnxIdftApplyTmpA, VecZnxNormalizeTmpBytes,
     },
+    oep::{SvpPPolAllocBytesImpl, VecZnxBigAllocBytesImpl, VecZnxDftAllocBytesImpl},
     layouts::{Backend, DataMut, DataRef, Module, Scratch},
     source::Source,
 };
 
 use crate::{
+    ScratchTakeCore,
     encryption::compressed::gglwe_ksk::GGLWEKeyCompressedEncryptSk,
     layouts::{
-        GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos, Rank, TensorKey,
+        GetDist, GGLWEInfos, GLWEInfos, GLWESecret, GLWESecretToRef, LWEInfos, Rank, TensorKey,
         compressed::{TensorKeyCompressed, TensorKeyCompressedToMut},
     },
 };
@@ -19,7 +21,7 @@ impl TensorKeyCompressed<Vec<u8>> {
     pub fn encrypt_sk_tmp_bytes<B: Backend, A>(module: &Module<B>, infos: &A) -> usize
     where
         A: GGLWEInfos,
-        Module<B>: SvpPPolBytesOf + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigBytesOf,
+        Module<B>: ModuleN + SvpPPolBytesOf + SvpPPolAlloc<B> + VecZnxNormalizeTmpBytes + VecZnxDftBytesOf + VecZnxNormalizeTmpBytes + VecZnxBigBytesOf,
     {
         TensorKey::encrypt_sk_tmp_bytes(module, infos)
     }
@@ -35,18 +37,25 @@ pub trait GGLWETensorKeyCompressedEncryptSk<B: Backend> {
         scratch: &mut Scratch<B>,
     ) where
         R: TensorKeyCompressedToMut,
-        S: GLWESecretToRef;
+        S: GLWESecretToRef + GetDist;
 }
 
 impl<B: Backend> GGLWETensorKeyCompressedEncryptSk<B> for Module<B>
 where
-    Module<B>: GGLWEKeyCompressedEncryptSk<B>
+    Module<B>: ModuleN
+        + GGLWEKeyCompressedEncryptSk<B>
         + VecZnxDftApply<B>
         + SvpApplyDftToDft<B>
         + VecZnxIdftApplyTmpA<B>
         + VecZnxBigNormalize<B>
-        + SvpPrepare<B>,
-    Scratch<B>:,
+        + SvpPrepare<B>
+        + SvpPPolAllocBytesImpl<B>
+        + SvpPPolBytesOf
+        + VecZnxDftAllocBytesImpl<B>
+        + VecZnxBigAllocBytesImpl<B>
+        + VecZnxDftBytesOf
+        + VecZnxBigBytesOf,
+    Scratch<B>: ScratchTakeBasic + ScratchTakeCore<B>,
 {
     fn gglwe_tensor_key_encrypt_sk<R, S>(
         &self,
@@ -57,9 +66,13 @@ where
         scratch: &mut Scratch<B>,
     ) where
         R: TensorKeyCompressedToMut,
-        S: GLWESecretToRef,
+        S: GLWESecretToRef + GetDist,
     {
         let res: &mut TensorKeyCompressed<&mut [u8]> = &mut res.to_mut();
+
+        let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(self, res.rank_out());
+        sk_dft_prep.prepare(self, sk);
+
         let sk: &GLWESecret<&[u8]> = &sk.to_ref();
 
         #[cfg(debug_assertions)]
@@ -68,21 +81,18 @@ where
             assert_eq!(res.n(), sk.n());
         }
 
-        let n: usize = sk.n().into();
+        // let n: usize = sk.n().into();
         let rank: usize = res.rank_out().into();
 
-        let (mut sk_dft_prep, scratch_1) = scratch.take_glwe_secret_prepared(sk.n(), res.rank_out());
-        sk_dft_prep.prepare(self, sk, scratch_1);
-
-        let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(n, rank, 1);
+        let (mut sk_dft, scratch_2) = scratch_1.take_vec_znx_dft(self, rank, 1);
 
         for i in 0..rank {
             self.vec_znx_dft_apply(1, 0, &mut sk_dft, i, &sk.data.as_vec_znx(), i);
         }
 
-        let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(n, 1, 1);
-        let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(sk.n(), Rank(1));
-        let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(n, 1, 1);
+        let (mut sk_ij_big, scratch_3) = scratch_2.take_vec_znx_big(self, 1, 1);
+        let (mut sk_ij, scratch_4) = scratch_3.take_glwe_secret(self, Rank(1));
+        let (mut sk_ij_dft, scratch_5) = scratch_4.take_vec_znx_dft(self, 1, 1);
 
         let mut source_xa: Source = Source::new(seed_xa);
 
@@ -125,6 +135,7 @@ impl<DataSelf: DataMut> TensorKeyCompressed<DataSelf> {
         source_xe: &mut Source,
         scratch: &mut Scratch<B>,
     ) where
+        GLWESecret<DataSk>: GetDist,
         Module<B>: GGLWETensorKeyCompressedEncryptSk<B>,
     {
         module.gglwe_tensor_key_encrypt_sk(self, sk, seed_xa, source_xe, scratch);

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

@@ -1,10 +1,11 @@
 use poulpy_hal::{
-    api::{VecZnxAddScalarInplace, VecZnxDftBytesOf, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes},
+    api::{ModuleN, VecZnxAddScalarInplace, VecZnxDftBytesOf, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes},
     layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, ScalarZnxToRef, Scratch, ZnxZero},
     source::Source,
 };
 
 use crate::{
+    ScratchTakeCore,
     encryption::{SIGMA, glwe_ct::GLWEEncryptSkInternal},
     layouts::{
         GGSW, GGSWInfos, GLWEInfos, LWEInfos,
@@ -40,8 +41,8 @@ pub trait GGSWCompressedEncryptSk<B: Backend> {
 
 impl<B: Backend> GGSWCompressedEncryptSk<B> for Module<B>
 where
-    Module<B>: GLWEEncryptSkInternal<B> + VecZnxAddScalarInplace + VecZnxNormalizeInplace<B>,
-    Scratch<B>:,
+    Module<B>: ModuleN + GLWEEncryptSkInternal<B> + VecZnxAddScalarInplace + VecZnxNormalizeInplace<B>,
+    Scratch<B>: ScratchTakeCore<B>,
 {
     fn ggsw_compressed_encrypt_sk<R, P, S>(
         &self,
@@ -74,7 +75,7 @@ where
         let cols: usize = rank + 1;
         let dsize: usize = res.dsize().into();
 
-        let (mut tmp_pt, scratch_1) = scratch.take_glwe_pt(&res.glwe_layout());
+        let (mut tmp_pt, scratch_1) = scratch.take_glwe_pt(self, &res.glwe_layout());
 
         let mut source = Source::new(seed_xa);