Add cross-basek normalization (#90)

* added cross_basek_normalization

* updated method signatures to take layouts

* fixed cross-base normalization

fix #91
fix #93

poulpy-core/src/conversion/glwe_to_lwe.rs

@@ -1,31 +1,46 @@
 use poulpy_hal::{
     api::{
-        ScratchAvailable, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
-        VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VmpApplyDftToDft, VmpApplyDftToDftAdd,
-        VmpApplyDftToDftTmpBytes,
+        ScratchAvailable, TakeVecZnx, TakeVecZnxDft, VecZnxBigAddSmallInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
+        VecZnxDftAllocBytes, VecZnxDftApply, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes, VmpApplyDftToDft,
+        VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
     },
     layouts::{Backend, DataMut, DataRef, Module, Scratch, ZnxView, ZnxViewMut, ZnxZero},
 };
 
 use crate::{
     TakeGLWECt,
-    layouts::{GLWECiphertext, Infos, LWECiphertext, prepared::GLWEToLWESwitchingKeyPrepared},
+    layouts::{
+        GGLWELayoutInfos, GLWECiphertext, GLWECiphertextLayout, GLWEInfos, LWECiphertext, LWEInfos, Rank,
+        prepared::GLWEToLWESwitchingKeyPrepared,
+    },
 };
 
 impl LWECiphertext<Vec<u8>> {
-    pub fn from_glwe_scratch_space<B: Backend>(
+    pub fn from_glwe_scratch_space<B: Backend, OUT, IN, KEY>(
         module: &Module<B>,
-        basek: usize,
-        k_lwe: usize,
-        k_glwe: usize,
-        k_ksk: usize,
-        rank: usize,
+        lwe_infos: &OUT,
+        glwe_infos: &IN,
+        key_infos: &KEY,
     ) -> usize
     where
-        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes,
+        OUT: LWEInfos,
+        IN: GLWEInfos,
+        KEY: GGLWELayoutInfos,
+        Module<B>: VecZnxDftAllocBytes + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
     {
-        GLWECiphertext::bytes_of(module.n(), basek, k_lwe, 1)
-            + GLWECiphertext::keyswitch_scratch_space(module, basek, k_lwe, k_glwe, k_ksk, 1, rank, 1)
+        let glwe_layout: GLWECiphertextLayout = GLWECiphertextLayout {
+            n: module.n().into(),
+            base2k: lwe_infos.base2k(),
+            k: lwe_infos.k(),
+            rank: Rank(1),
+        };
+
+        GLWECiphertext::alloc_bytes_with(
+            module.n().into(),
+            lwe_infos.base2k(),
+            lwe_infos.k(),
+            1u32.into(),
+        ) + GLWECiphertext::keyswitch_scratch_space(module, &glwe_layout, glwe_infos, key_infos)
     }
 }
 
@@ -34,10 +49,11 @@ impl<DLwe: DataMut> LWECiphertext<DLwe> {
         #[cfg(debug_assertions)]
         {
             assert!(self.n() <= a.n());
+            assert!(self.base2k() == a.base2k());
         }
 
         let min_size: usize = self.size().min(a.size());
-        let n: usize = self.n();
+        let n: usize = self.n().into();
 
         self.data.zero();
         (0..min_size).for_each(|i| {
@@ -64,15 +80,26 @@ impl<DLwe: DataMut> LWECiphertext<DLwe> {
             + VecZnxDftApply<B>
             + VecZnxIdftApplyConsume<B>
             + VecZnxBigAddSmallInplace<B>
-            + VecZnxBigNormalize<B>,
-        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeGLWECt,
+            + VecZnxBigNormalize<B>
+            + VecZnxNormalize<B>
+            + VecZnxNormalizeTmpBytes,
+        Scratch<B>: ScratchAvailable + TakeVecZnxDft<B> + TakeGLWECt + TakeVecZnx,
     {
         #[cfg(debug_assertions)]
         {
-            assert_eq!(self.basek(), a.basek());
-            assert_eq!(a.n(), ks.n());
+            assert_eq!(a.n(), module.n() as u32);
+            assert_eq!(ks.n(), module.n() as u32);
+            assert!(self.n() <= module.n() as u32);
         }
-        let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(a.n(), a.basek(), self.k(), 1);
+
+        let glwe_layout: GLWECiphertextLayout = GLWECiphertextLayout {
+            n: module.n().into(),
+            base2k: self.base2k(),
+            k: self.k(),
+            rank: Rank(1),
+        };
+
+        let (mut tmp_glwe, scratch_1) = scratch.take_glwe_ct(&glwe_layout);
         tmp_glwe.keyswitch(module, a, &ks.0, scratch_1);
         self.sample_extract(&tmp_glwe);
     }