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/layouts/prepared/lwe_ksk.rs

@@ -1,65 +1,124 @@
 use poulpy_hal::{
     api::{VmpPMatAlloc, VmpPMatAllocBytes, VmpPrepare},
-    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch, VmpPMat},
+    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
 };
 
 use crate::layouts::{
-    Infos, LWESwitchingKey,
+    Base2K, Degree, Digits, GGLWELayoutInfos, GLWEInfos, LWEInfos, LWESwitchingKey, Rank, Rows, TorusPrecision,
     prepared::{GGLWESwitchingKeyPrepared, Prepare, PrepareAlloc},
 };
 
 #[derive(PartialEq, Eq)]
 pub struct LWESwitchingKeyPrepared<D: Data, B: Backend>(pub(crate) GGLWESwitchingKeyPrepared<D, B>);
 
-impl<D: Data, B: Backend> Infos for LWESwitchingKeyPrepared<D, B> {
-    type Inner = VmpPMat<D, B>;
-
-    fn inner(&self) -> &Self::Inner {
-        self.0.inner()
+impl<D: Data, B: Backend> LWEInfos for LWESwitchingKeyPrepared<D, B> {
+    fn base2k(&self) -> Base2K {
+        self.0.base2k()
     }
 
-    fn basek(&self) -> usize {
-        self.0.basek()
-    }
-
-    fn k(&self) -> usize {
+    fn k(&self) -> TorusPrecision {
         self.0.k()
     }
+
+    fn n(&self) -> Degree {
+        self.0.n()
+    }
+
+    fn size(&self) -> usize {
+        self.0.size()
+    }
+}
+impl<D: Data, B: Backend> GLWEInfos for LWESwitchingKeyPrepared<D, B> {
+    fn rank(&self) -> Rank {
+        self.rank_out()
+    }
 }
 
-impl<D: Data, B: Backend> LWESwitchingKeyPrepared<D, B> {
-    pub fn digits(&self) -> usize {
+impl<D: Data, B: Backend> GGLWELayoutInfos for LWESwitchingKeyPrepared<D, B> {
+    fn digits(&self) -> Digits {
         self.0.digits()
     }
 
-    pub fn rank(&self) -> usize {
-        self.0.rank()
-    }
-
-    pub fn rank_in(&self) -> usize {
+    fn rank_in(&self) -> Rank {
         self.0.rank_in()
     }
 
-    pub fn rank_out(&self) -> usize {
+    fn rank_out(&self) -> Rank {
         self.0.rank_out()
     }
+
+    fn rows(&self) -> Rows {
+        self.0.rows()
+    }
 }
 
 impl<B: Backend> LWESwitchingKeyPrepared<Vec<u8>, B> {
-    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize) -> Self
+    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
+    where
+        A: GGLWELayoutInfos,
+        Module<B>: VmpPMatAlloc<B>,
+    {
+        debug_assert_eq!(
+            infos.digits().0,
+            1,
+            "digits > 1 is not supported for LWESwitchingKey"
+        );
+        debug_assert_eq!(
+            infos.rank_in().0,
+            1,
+            "rank_in > 1 is not supported for LWESwitchingKey"
+        );
+        debug_assert_eq!(
+            infos.rank_out().0,
+            1,
+            "rank_out > 1 is not supported for LWESwitchingKey"
+        );
+        Self(GGLWESwitchingKeyPrepared::alloc(module, infos))
+    }
+
+    pub fn alloc_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rows: Rows) -> Self
     where
         Module<B>: VmpPMatAlloc<B>,
     {
-        Self(GGLWESwitchingKeyPrepared::alloc(
-            module, basek, k, rows, 1, 1, 1,
+        Self(GGLWESwitchingKeyPrepared::alloc_with(
+            module,
+            base2k,
+            k,
+            rows,
+            Digits(1),
+            Rank(1),
+            Rank(1),
         ))
     }
 
-    pub fn bytes_of(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize) -> usize
+    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
+    where
+        A: GGLWELayoutInfos,
+        Module<B>: VmpPMatAllocBytes,
+    {
+        debug_assert_eq!(
+            infos.digits().0,
+            1,
+            "digits > 1 is not supported for LWESwitchingKey"
+        );
+        debug_assert_eq!(
+            infos.rank_in().0,
+            1,
+            "rank_in > 1 is not supported for LWESwitchingKey"
+        );
+        debug_assert_eq!(
+            infos.rank_out().0,
+            1,
+            "rank_out > 1 is not supported for LWESwitchingKey"
+        );
+        GGLWESwitchingKeyPrepared::alloc_bytes(module, infos)
+    }
+
+    pub fn alloc_bytes_with(module: &Module<B>, base2k: Base2K, k: TorusPrecision, rows: Rows) -> usize
     where
         Module<B>: VmpPMatAllocBytes,
     {
-        GGLWESwitchingKeyPrepared::<Vec<u8>, B>::bytes_of(module, basek, k, rows, digits, 1, 1)
+        GGLWESwitchingKeyPrepared::alloc_bytes_with(module, base2k, k, rows, Digits(1), Rank(1), Rank(1))
     }
 }
 
@@ -68,8 +127,7 @@ where
     Module<B>: VmpPrepare<B> + VmpPMatAlloc<B>,
 {
     fn prepare_alloc(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> LWESwitchingKeyPrepared<Vec<u8>, B> {
-        let mut ksk_prepared: LWESwitchingKeyPrepared<Vec<u8>, B> =
-            LWESwitchingKeyPrepared::alloc(module, self.0.basek(), self.0.k(), self.0.rows());
+        let mut ksk_prepared: LWESwitchingKeyPrepared<Vec<u8>, B> = LWESwitchingKeyPrepared::alloc(module, self);
         ksk_prepared.prepare(module, self, scratch);
         ksk_prepared
     }