wip

poulpy-core/src/layouts/gglwe_atk.rs

@@ -1,10 +1,10 @@
 use poulpy_hal::{
-    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo},
+    layouts::{Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, WriterTo},
     source::Source,
 };
 
 use crate::layouts::{
-    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWECiphertext, GLWEInfos, GLWESwitchingKey, GLWESwitchingKeyToMut,
+    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWE, GLWEInfos, GLWESwitchingKey, GLWESwitchingKeyAlloc, GLWESwitchingKeyToMut,
     GLWESwitchingKeyToRef, LWEInfos, Rank, TorusPrecision,
 };
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
@@ -131,43 +131,105 @@ impl<D: DataRef> fmt::Display for AutomorphismKey<D> {
     }
 }
 
+impl<B: Backend> AutomorphismKeyAlloc for Module<B> where Self: GLWESwitchingKeyAlloc {}
+
+pub trait AutomorphismKeyAlloc
+where
+    Self: GLWESwitchingKeyAlloc,
+{
+    fn alloc_automorphism_key(
+        &self,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> AutomorphismKey<Vec<u8>> {
+        AutomorphismKey {
+            key: self.alloc_glwe_switching_key(base2k, k, rank, rank, dnum, dsize),
+            p: 0,
+        }
+    }
+
+    fn alloc_automorphism_key_from_infos<A>(&self, infos: &A) -> AutomorphismKey<Vec<u8>>
+    where
+        A: GGLWEInfos,
+    {
+        self.alloc_automorphism_key(
+            infos.base2k(),
+            infos.k(),
+            infos.rank(),
+            infos.dnum(),
+            infos.dsize(),
+        )
+    }
+
+    fn bytes_of_automorphism_key(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
+        self.bytes_of_glwe_switching_key(base2k, k, rank, rank, dnum, dsize)
+    }
+
+    fn bytes_of_automorphism_key_from_infos<A>(&self, infos: &A) -> usize
+    where
+        A: GGLWEInfos,
+    {
+        assert_eq!(
+            infos.rank_in(),
+            infos.rank_out(),
+            "rank_in != rank_out is not supported for AutomorphismKey"
+        );
+        self.bytes_of_automorphism_key(
+            infos.base2k(),
+            infos.k(),
+            infos.rank(),
+            infos.dnum(),
+            infos.dsize(),
+        )
+    }
+}
+
 impl AutomorphismKey<Vec<u8>> {
-    pub fn alloc<A>(infos: &A) -> Self
+    pub fn alloc_from_infos<A, B: Backend>(module: &Module<B>, infos: &A) -> Self
     where
         A: GGLWEInfos,
+        Module<B>: AutomorphismKeyAlloc,
     {
-        assert_eq!(
-            infos.rank_in(),
-            infos.rank_out(),
-            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
-        );
-        AutomorphismKey {
-            key: GLWESwitchingKey::alloc(infos),
-            p: 0,
-        }
+        module.alloc_automorphism_key_from_infos(infos)
     }
 
-    pub fn alloc_with(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
-        AutomorphismKey {
-            key: GLWESwitchingKey::alloc_with(n, base2k, k, rank, rank, dnum, dsize),
-            p: 0,
-        }
+    pub fn alloc_with<B: Backend>(
+        module: &Module<B>,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> Self
+    where
+        Module<B>: AutomorphismKeyAlloc,
+    {
+        module.alloc_automorphism_key(base2k, k, rank, dnum, dsize)
     }
 
-    pub fn alloc_bytes<A>(infos: &A) -> usize
+    pub fn bytes_of_from_infos<A, B: Backend>(module: &Module<B>, infos: &A) -> usize
     where
         A: GGLWEInfos,
+        Module<B>: AutomorphismKeyAlloc,
     {
-        assert_eq!(
-            infos.rank_in(),
-            infos.rank_out(),
-            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
-        );
-        GLWESwitchingKey::alloc_bytes(infos)
+        module.bytes_of_automorphism_key_from_infos(infos)
     }
 
-    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
-        GLWESwitchingKey::alloc_bytes_with(n, base2k, k, rank, rank, dnum, dsize)
+    pub fn bytes_of<B: Backend>(
+        module: &Module<B>,
+        base2k: Base2K,
+        k: TorusPrecision,
+        rank: Rank,
+        dnum: Dnum,
+        dsize: Dsize,
+    ) -> usize
+    where
+        Module<B>: AutomorphismKeyAlloc,
+    {
+        module.bytes_of_automorphism_key(base2k, k, rank, dnum, dsize)
     }
 }
 
@@ -204,13 +266,13 @@ where
 }
 
 impl<D: DataRef> AutomorphismKey<D> {
-    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
+    pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
         self.key.at(row, col)
     }
 }
 
 impl<D: DataMut> AutomorphismKey<D> {
-    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
+    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
         self.key.at_mut(row, col)
     }
 }