Replace hasmap of glweautomorphismkeys by helper trait, enabling not having to pass, for example, but full CBT key for ops that do not require it

poulpy-core/src/layouts/compressed/glwe_automorphism_key.rs

@@ -138,7 +138,7 @@ impl<D: DataRef> WriterTo for GLWEAutomorphismKeyCompressed<D> {
     }
 }
 
-pub trait AutomorphismKeyDecompress
+pub trait GLWEAutomorphismKeyDecompress
 where
     Self: GGLWEDecompress,
 {
@@ -152,7 +152,7 @@ where
     }
 }
 
-impl<B: Backend> AutomorphismKeyDecompress for Module<B> where Self: GLWEDecompress {}
+impl<B: Backend> GLWEAutomorphismKeyDecompress for Module<B> where Self: GLWEDecompress {}
 
 impl<D: DataMut> GLWEAutomorphismKey<D>
 where
@@ -161,7 +161,7 @@ where
     pub fn decompress<O, M>(&mut self, module: &M, other: &O)
     where
         O: GGLWECompressedToRef + GetGaloisElement,
-        M: AutomorphismKeyDecompress,
+        M: GLWEAutomorphismKeyDecompress,
     {
         module.decompress_automorphism_key(self, other);
     }

poulpy-core/src/layouts/glwe_automorphism_key.rs

@@ -1,15 +1,21 @@
 use poulpy_hal::{
-    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo},
+    layouts::{Backend, Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo},
     source::Source,
 };
 
 use crate::layouts::{
-    Base2K, Degree, Dnum, Dsize, GGLWE, GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWE, GLWEInfos, LWEInfos, Rank, TorusPrecision,
+    Base2K, Degree, Dnum, Dsize, GGLWE, GGLWEInfos, GGLWELayout, GGLWEToMut, GGLWEToRef, GLWE, GLWEInfos, LWEInfos, Rank,
+    TorusPrecision,
 };
 use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 
 use std::fmt;
 
+pub trait GLWEAutomorphismKeyHelper<K, BE: Backend> {
+    fn get_automorphism_key(&self, k: i64) -> Option<&K>;
+    fn automorphism_key_infos(&self) -> GGLWELayout;
+}
+
 #[derive(PartialEq, Eq, Copy, Clone, Debug)]
 pub struct GLWEAutomorphismKeyLayout {
     pub n: Degree,

poulpy-core/src/layouts/prepared/glwe_automorphism_key.rs

@@ -1,10 +1,28 @@
+use std::collections::HashMap;
+
 use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
 
 use crate::layouts::{
-    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GGLWEPrepared, GGLWEPreparedFactory, GGLWEPreparedToMut, GGLWEPreparedToRef,
-    GGLWEToRef, GLWEInfos, GetGaloisElement, LWEInfos, Rank, SetGaloisElement, TorusPrecision,
+    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GGLWELayout, GGLWEPrepared, GGLWEPreparedFactory, GGLWEPreparedToMut,
+    GGLWEPreparedToRef, GGLWEToRef, GLWEAutomorphismKeyHelper, GLWEInfos, GetGaloisElement, LWEInfos, Rank, SetGaloisElement,
+    TorusPrecision,
 };
 
+impl<K, BE: Backend> GLWEAutomorphismKeyHelper<K, BE> for HashMap<i64, K>
+where
+    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+{
+    fn get_automorphism_key(&self, k: i64) -> Option<&K> {
+        self.get(&k)
+    }
+
+    fn automorphism_key_infos(&self) -> GGLWELayout {
+        self.get(self.keys().next().unwrap())
+            .unwrap()
+            .gglwe_layout()
+    }
+}
+
 #[derive(PartialEq, Eq)]
 pub struct GLWEAutomorphismKeyPrepared<D: Data, B: Backend> {
     pub(crate) key: GGLWEPrepared<D, B>,