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/glwe_packer.rs

@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use poulpy_hal::{
     api::ModuleLogN,
     layouts::{Backend, GaloisElement, Module, Scratch},
@@ -8,7 +6,10 @@ use poulpy_hal::{
 use crate::{
     GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore,
     glwe_trace::GLWETrace,
-    layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos},
+    layouts::{
+        GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement,
+        LWEInfos,
+    },
 };
 
 /// [GLWEPacker] enables only the fly GLWE packing
@@ -110,12 +111,13 @@ impl GLWEPacker {
     /// * `res`: space to append fully packed ciphertext. Only when the number
     ///   of packed ciphertexts reaches N/2^log_batch is a result written.
     /// * `a`: ciphertext to pack. Can optionally give None to pack a 0 ciphertext.
-    /// * `auto_keys`: a [HashMap] containing the [AutomorphismKeyExec]s.
+    /// * `auto_keys`: an implementation of [GLWEAutomorphismKeyHelper], containing [GLWEAutomorphismKeyPrepared] with index of [Self::galois_elements].
     /// * `scratch`: scratch space of size at least [Self::tmp_bytes].
-    pub fn add<A, K, M, BE: Backend>(&mut self, module: &M, a: Option<&A>, auto_keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
+    pub fn add<A, K, H, M, BE: Backend>(&mut self, module: &M, a: Option<&A>, auto_keys: &H, scratch: &mut Scratch<BE>)
     where
         A: GLWEToRef + GLWEInfos,
         K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
         M: GLWEPackerOps<BE>,
         Scratch<BE>: ScratchTakeCore<BE>,
     {
@@ -187,28 +189,23 @@ where
         + GLWEAdd
         + GLWENormalize<BE>,
 {
-    fn packer_add<A, K>(
+    fn packer_add<A, K, H>(&self, packer: &mut GLWEPacker, a: Option<&A>, i: usize, auto_keys: &H, scratch: &mut Scratch<BE>)
-        &self,
+    where
-        packer: &mut GLWEPacker,
-        a: Option<&A>,
-        i: usize,
-        auto_keys: &HashMap<i64, K>,
-        scratch: &mut Scratch<BE>,
-    ) where
         A: GLWEToRef + GLWEInfos,
         K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
         Scratch<BE>: ScratchTakeCore<BE>,
     {
         pack_core(self, a, &mut packer.accumulators, i, auto_keys, scratch)
     }
 }
 
-fn pack_core<A, K, M, BE: Backend>(
+fn pack_core<A, K, H, M, BE: Backend>(
     module: &M,
     a: Option<&A>,
     accumulators: &mut [Accumulator],
     i: usize,
-    auto_keys: &HashMap<i64, K>,
+    auto_keys: &H,
     scratch: &mut Scratch<BE>,
 ) where
     A: GLWEToRef + GLWEInfos,
@@ -229,6 +226,7 @@ fn pack_core<A, K, M, BE: Backend>(
         + GLWEAdd
         + GLWENormalize<BE>,
     K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+    H: GLWEAutomorphismKeyHelper<K, BE>,
     Scratch<BE>: ScratchTakeCore<BE>,
 {
     let log_n: usize = module.log_n();
@@ -280,12 +278,12 @@ fn pack_core<A, K, M, BE: Backend>(
     }
 }
 
-fn combine<B, K, M, BE: Backend>(
+fn combine<B, K, H, M, BE: Backend>(
     module: &M,
     acc: &mut Accumulator,
     b: Option<&B>,
     i: usize,
-    auto_keys: &HashMap<i64, K>,
+    auto_keys: &H,
     scratch: &mut Scratch<BE>,
 ) where
     B: GLWEToRef + GLWEInfos,
@@ -307,6 +305,7 @@ fn combine<B, K, M, BE: Backend>(
         + GLWEAdd
         + GLWENormalize<BE>,
     K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+    H: GLWEAutomorphismKeyHelper<K, BE>,
     Scratch<BE>: ScratchTakeCore<BE>,
 {
     let log_n: usize = acc.data.n().log2();
@@ -348,7 +347,7 @@ fn combine<B, K, M, BE: Backend>(
             module.glwe_normalize_inplace(&mut tmp_b, scratch_1);
 
             // tmp_b = phi(a * X^-t - b)
-            if let Some(auto_key) = auto_keys.get(&gal_el) {
+            if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
                 module.glwe_automorphism_inplace(&mut tmp_b, auto_key, scratch_1);
             } else {
                 panic!("auto_key[{gal_el}] not found");
@@ -365,7 +364,7 @@ fn combine<B, K, M, BE: Backend>(
         } else {
             module.glwe_rsh(1, a, scratch);
             // a = a + phi(a)
-            if let Some(auto_key) = auto_keys.get(&gal_el) {
+            if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
                 module.glwe_automorphism_add_inplace(a, auto_key, scratch);
             } else {
                 panic!("auto_key[{gal_el}] not found");
@@ -377,7 +376,7 @@ fn combine<B, K, M, BE: Backend>(
         module.glwe_rsh(1, &mut tmp_b, scratch_1);
 
         // a = (b* X^t - phi(b* X^t))
-        if let Some(auto_key) = auto_keys.get(&gal_el) {
+        if let Some(auto_key) = auto_keys.get_automorphism_key(gal_el) {
             module.glwe_automorphism_sub_negate(a, &tmp_b, auto_key, scratch_1);
         } else {
             panic!("auto_key[{gal_el}] not found");

poulpy-core/src/glwe_packing.rs

@@ -7,20 +7,16 @@ use poulpy_hal::{
 
 use crate::{
     GLWEAdd, GLWEAutomorphism, GLWECopy, GLWENormalize, GLWERotate, GLWEShift, GLWESub, ScratchTakeCore,
-    layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement},
+    layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWEAutomorphismKeyHelper, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement},
 };
 pub trait GLWEPacking<BE: Backend> {
     /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
     /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
-    fn glwe_pack<R, K>(
+    fn glwe_pack<R, K, H>(&self, cts: &mut HashMap<usize, &mut R>, log_gap_out: usize, keys: &H, scratch: &mut Scratch<BE>)
-        &self,
+    where
-        cts: &mut HashMap<usize, &mut R>,
-        log_gap_out: usize,
-        keys: &HashMap<i64, K>,
-        scratch: &mut Scratch<BE>,
-    ) where
         R: GLWEToMut + GLWEToRef + GLWEInfos,
-        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
+        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>;
 }
 
 impl<BE: Backend> GLWEPacking<BE> for Module<BE>
@@ -38,15 +34,11 @@ where
 {
     /// Packs [x_0: GLWE(m_0), x_1: GLWE(m_1), ..., x_i: GLWE(m_i)]
     /// to [0: GLWE(m_0 * X^x_0 + m_1 * X^x_1 + ... + m_i * X^x_i)]
-    fn glwe_pack<R, K>(
+    fn glwe_pack<R, K, H>(&self, cts: &mut HashMap<usize, &mut R>, log_gap_out: usize, keys: &H, scratch: &mut Scratch<BE>)
-        &self,
+    where
-        cts: &mut HashMap<usize, &mut R>,
-        log_gap_out: usize,
-        keys: &HashMap<i64, K>,
-        scratch: &mut Scratch<BE>,
-    ) where
         R: GLWEToMut + GLWEToRef + GLWEInfos,
         K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
     {
         #[cfg(debug_assertions)]
         {
@@ -59,9 +51,10 @@ where
             let t: usize = (1 << log_n).min(1 << (log_n - 1 - i));
 
             let key: &K = if i == 0 {
-                keys.get(&-1).unwrap()
+                keys.get_automorphism_key(-1).unwrap()
             } else {
-                keys.get(&self.galois_element(1 << (i - 1))).unwrap()
+                keys.get_automorphism_key(self.galois_element(1 << (i - 1)))
+                    .unwrap()
             };
 
             for j in 0..t {

poulpy-core/src/glwe_trace.rs

@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use poulpy_hal::{
     api::{ModuleLogN, VecZnxNormalize, VecZnxNormalizeTmpBytes},
     layouts::{Backend, CyclotomicOrder, DataMut, GaloisElement, Module, Scratch, VecZnx, galois_element},
@@ -8,7 +6,8 @@ use poulpy_hal::{
 use crate::{
     GLWEAutomorphism, GLWECopy, GLWEShift, ScratchTakeCore,
     layouts::{
-        Base2K, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWELayout, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos,
+        GGLWEInfos, GGLWELayout, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWELayout, GLWEToMut,
+        GLWEToRef, GetGaloisElement, LWEInfos,
     },
 };
 
@@ -32,32 +31,34 @@ impl GLWE<Vec<u8>> {
 }
 
 impl<D: DataMut> GLWE<D> {
-    pub fn trace<A, K, M, BE: Backend>(
+    pub fn trace<A, H, K, M, BE: Backend>(
         &mut self,
         module: &M,
         start: usize,
         end: usize,
         a: &A,
-        keys: &HashMap<i64, K>,
+        keys: &H,
         scratch: &mut Scratch<BE>,
     ) where
         A: GLWEToRef,
         K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
         Scratch<BE>: ScratchTakeCore<BE>,
         M: GLWETrace<BE>,
     {
         module.glwe_trace(self, start, end, a, keys, scratch);
     }
 
-    pub fn trace_inplace<K, M, BE: Backend>(
+    pub fn trace_inplace<H, K, M, BE: Backend>(
         &mut self,
         module: &M,
         start: usize,
         end: usize,
-        keys: &HashMap<i64, K>,
+        keys: &H,
         scratch: &mut Scratch<BE>,
     ) where
         K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
         Scratch<BE>: ScratchTakeCore<BE>,
         M: GLWETrace<BE>,
     {
@@ -113,52 +114,48 @@ where
         trace
     }
 
-    fn glwe_trace<R, A, K>(&self, res: &mut R, start: usize, end: usize, a: &A, keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
+    fn glwe_trace<R, A, K, H>(&self, res: &mut R, start: usize, end: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
     where
         R: GLWEToMut,
         A: GLWEToRef,
         K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
     {
         self.glwe_copy(res, a);
         self.glwe_trace_inplace(res, start, end, keys, scratch);
     }
 
-    fn glwe_trace_inplace<R, K>(&self, res: &mut R, start: usize, end: usize, keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
+    fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, start: usize, end: usize, keys: &H, scratch: &mut Scratch<BE>)
     where
         R: GLWEToMut,
         K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
     {
         let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
 
-        let basek_ksk: Base2K = keys.get(keys.keys().next().unwrap()).unwrap().base2k();
+        let ksk_infos: &GGLWELayout = &keys.automorphism_key_infos();
 
-        #[cfg(debug_assertions)]
+        assert_eq!(res.n(), self.n() as u32);
-        {
+        assert_eq!(ksk_infos.n(), self.n() as u32);
-            assert_eq!(res.n(), self.n() as u32);
+        assert!(start < end);
-            assert!(start < end);
+        assert!(end <= self.log_n());
-            assert!(end <= self.log_n());
+        assert_eq!(ksk_infos.rank_in(), res.rank());
-            for key in keys.values() {
+        assert_eq!(ksk_infos.rank_out(), res.rank());
-                assert_eq!(key.n(), self.n() as u32);
-                assert_eq!(key.base2k(), basek_ksk);
-                assert_eq!(key.rank_in(), res.rank());
-                assert_eq!(key.rank_out(), res.rank());
-            }
-        }
 
-        if res.base2k() != basek_ksk {
+        if res.base2k() != ksk_infos.base2k() {
             let (mut self_conv, scratch_1) = scratch.take_glwe(&GLWELayout {
                 n: self.n().into(),
-                base2k: basek_ksk,
+                base2k: ksk_infos.base2k(),
                 k: res.k(),
                 rank: res.rank(),
             });
 
             for j in 0..(res.rank() + 1).into() {
                 self.vec_znx_normalize(
-                    basek_ksk.into(),
+                    ksk_infos.base2k().into(),
                     &mut self_conv.data,
                     j,
-                    basek_ksk.into(),
+                    res.base2k().into(),
                     res.data(),
                     j,
                     scratch_1,
@@ -174,7 +171,7 @@ where
                     self.galois_element(1 << (i - 1))
                 };
 
-                if let Some(key) = keys.get(&p) {
+                if let Some(key) = keys.get_automorphism_key(p) {
                     self.glwe_automorphism_add_inplace(&mut self_conv, key, scratch_1);
                 } else {
                     panic!("keys[{p}] is empty")
@@ -186,7 +183,7 @@ where
                     res.base2k().into(),
                     res.data_mut(),
                     j,
-                    basek_ksk.into(),
+                    ksk_infos.base2k().into(),
                     &self_conv.data,
                     j,
                     scratch_1,
@@ -204,7 +201,7 @@ where
                     self.galois_element(1 << (i - 1))
                 };
 
-                if let Some(key) = keys.get(&p) {
+                if let Some(key) = keys.get_automorphism_key(p) {
                     self.glwe_automorphism_add_inplace(res, key, scratch);
                 } else {
                     panic!("keys[{p}] is empty")
@@ -223,21 +220,16 @@ pub trait GLWETrace<BE: Backend> {
         A: GLWEInfos,
         K: GGLWEInfos;
 
-    fn glwe_trace<R, A, K>(
+    fn glwe_trace<R, A, K, H>(&self, res: &mut R, start: usize, end: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
-        &self,
-        res: &mut R,
-        start: usize,
-        end: usize,
-        a: &A,
-        keys: &HashMap<i64, K>,
-        scratch: &mut Scratch<BE>,
-    ) where
-        R: GLWEToMut,
-        A: GLWEToRef,
-        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
-
-    fn glwe_trace_inplace<R, K>(&self, res: &mut R, start: usize, end: usize, keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
     where
         R: GLWEToMut,
-        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
+        A: GLWEToRef,
+        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>;
+
+    fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, start: usize, end: usize, keys: &H, scratch: &mut Scratch<BE>)
+    where
+        R: GLWEToMut,
+        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
+        H: GLWEAutomorphismKeyHelper<K, BE>;
 }

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,
+    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GGLWELayout, GGLWEPrepared, GGLWEPreparedFactory, GGLWEPreparedToMut,
-    GGLWEToRef, GLWEInfos, GetGaloisElement, LWEInfos, Rank, SetGaloisElement, TorusPrecision,
+    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>,

poulpy-core/src/tests/test_suite/encryption/gglwe_atk.rs

@@ -9,7 +9,7 @@ use crate::{
     GLWESwitchingKeyEncryptSk, ScratchTakeCore,
     encryption::SIGMA,
     layouts::{
-        AutomorphismKeyDecompress, GLWEAutomorphismKey, GLWEAutomorphismKeyLayout, GLWEInfos, GLWESecret,
+        GLWEAutomorphismKey, GLWEAutomorphismKeyDecompress, GLWEAutomorphismKeyLayout, GLWEInfos, GLWESecret,
         GLWESecretPreparedFactory, GLWESwitchingKeyDecompress, compressed::GLWEAutomorphismKeyCompressed,
         prepared::GLWESecretPrepared,
     },
@@ -97,7 +97,7 @@ where
         + GLWESecretPreparedFactory<BE>
         + GLWESwitchingKeyEncryptSk<BE>
         + GLWESwitchingKeyCompressedEncryptSk<BE>
-        + AutomorphismKeyDecompress
+        + GLWEAutomorphismKeyDecompress
         + VecZnxAutomorphism
         + VecZnxFillUniform
         + GGLWENoise<BE>,

poulpy-schemes/src/tfhe/bdd_arithmetic/ciphertexts/fhe_uint.rs

@@ -1,8 +1,8 @@
 use poulpy_core::{
     GLWEAdd, GLWECopy, GLWEDecrypt, GLWEEncryptSk, GLWEPacking, GLWERotate, GLWESub, GLWETrace, LWEFromGLWE, ScratchTakeCore,
     layouts::{
-        Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEPlaintextLayout, GLWESecretPreparedToRef, GLWEToMut,
+        Base2K, Degree, GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWEPlaintextLayout,
-        GLWEToRef, LWEInfos, LWEToMut, Rank, TorusPrecision,
+        GLWESecretPreparedToRef, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos, LWEToMut, Rank, TorusPrecision,
     },
 };
 use poulpy_hal::{
@@ -171,20 +171,20 @@ impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
     }
 
     #[allow(clippy::too_many_arguments)]
-    pub fn splice_u16<D0, A, B, BRA, M, BE: Backend>(
+    pub fn splice_u16<A, B, H, K, M, BE: Backend>(
         &mut self,
         module: &M,
         dst: usize,
         src: usize,
         a: &A,
         b: &B,
-        keys: &BDDKeyPrepared<D0, BRA, BE>,
+        keys: &H,
         scratch: &mut Scratch<BE>,
     ) where
-        D0: DataRef,
         A: GLWEToRef + GLWEInfos,
         B: GLWEToRef + GLWEInfos,
-        BRA: BlindRotationAlgo,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
+        K: GGLWEPreparedToRef<BE> + GGLWEInfos + GetGaloisElement,
         M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
         Scratch<BE>: ScratchTakeBDD<T, BE>,
     {
@@ -206,20 +206,20 @@ impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
 
     #[allow(clippy::too_many_arguments)]
     // Store on the receiver a where the byte_a-th byte of a has been replaced by byte_src2 of src2.
-    pub fn splice_u8<D0, A, B, BRA, M, BE: Backend>(
+    pub fn splice_u8<A, B, H, K, M, BE: Backend>(
         &mut self,
         module: &M,
         dst: usize,
         src: usize,
         a: &A,
         b: &B,
-        keys: &BDDKeyPrepared<D0, BRA, BE>,
+        keys: &H,
         scratch: &mut Scratch<BE>,
     ) where
-        D0: DataRef,
         A: GLWEToRef + GLWEInfos,
         B: GLWEToRef + GLWEInfos,
-        BRA: BlindRotationAlgo,
+        H: GLWEAutomorphismKeyHelper<K, BE>,
+        K: GGLWEPreparedToRef<BE> + GGLWEInfos + GetGaloisElement,
         M: ModuleLogN + GLWERotate<BE> + GLWETrace<BE> + GLWESub + GLWEAdd + GLWECopy,
         Scratch<BE>: ScratchTakeBDD<T, BE>,
     {
@@ -241,7 +241,7 @@ impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
             trace_start,
             module.log_n(),
             self,
-            &keys.cbt.atk,
+            keys,
             scratch_1,
         );
 
@@ -263,7 +263,7 @@ impl<D: DataMut, T: UnsignedInteger> FheUint<D, T> {
             &mut tmp_fhe_uint_byte,
             trace_start,
             module.log_n(),
-            &keys.cbt.atk,
+            keys,
             scratch_1,
         );
 

poulpy-schemes/src/tfhe/bdd_arithmetic/key.rs

@@ -8,6 +8,7 @@ use crate::tfhe::{
     },
 };
 
+use poulpy_core::layouts::{GLWEAutomorphismKeyHelper, GLWEAutomorphismKeyPrepared};
 use poulpy_core::{
     GLWEToLWESwitchingKeyEncryptSk, GetDistribution, LWEFromGLWE, ScratchTakeCore,
     layouts::{
@@ -134,6 +135,18 @@ where
     pub(crate) ks: GLWEToLWEKeyPrepared<D, BE>,
 }
 
+impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<GLWEAutomorphismKeyPrepared<D, BE>, BE>
+    for BDDKeyPrepared<D, BRA, BE>
+{
+    fn automorphism_key_infos(&self) -> poulpy_core::layouts::GGLWELayout {
+        self.cbt.automorphism_key_infos()
+    }
+
+    fn get_automorphism_key(&self, k: i64) -> Option<&GLWEAutomorphismKeyPrepared<D, BE>> {
+        self.cbt.get_automorphism_key(k)
+    }
+}
+
 pub trait BDDKeyPreparedFactory<BRA: BlindRotationAlgo, BE: Backend>
 where
     Self: Sized + CircuitBootstrappingKeyPreparedFactory<BRA, BE> + GLWEToLWEKeyPreparedFactory<BE>,

poulpy-schemes/src/tfhe/circuit_bootstrapping/circuit.rs

@@ -8,11 +8,12 @@ use poulpy_hal::{
 use poulpy_core::{
     GGSWFromGGLWE, GLWEDecrypt, GLWEPacking, GLWERotate, GLWETrace, ScratchTakeCore,
     layouts::{
-        Dsize, GGLWELayout, GGSWInfos, GGSWToMut, GLWEInfos, GLWESecretPreparedFactory, GLWEToMut, GLWEToRef, LWEInfos, LWEToRef,
+        Dsize, GGLWEInfos, GGLWELayout, GGLWEPreparedToRef, GGSWInfos, GGSWToMut, GLWEAutomorphismKeyHelper, GLWEInfos,
+        GLWESecretPreparedFactory, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos, LWEToRef,
     },
 };
 
-use poulpy_core::layouts::{GGSW, GLWE, LWE, prepared::GLWEAutomorphismKeyPrepared};
+use poulpy_core::layouts::{GGSW, GLWE, LWE};
 
 use crate::tfhe::{
     blind_rotation::{
@@ -323,18 +324,20 @@ pub fn circuit_bootstrap_core<R, L, D, M, BRA: BlindRotationAlgo, BE: Backend>(
 }
 
 #[allow(clippy::too_many_arguments)]
-fn post_process<R, A, M, BE: Backend>(
+fn post_process<R, A, M, H, K, BE: Backend>(
     module: &M,
     res: &mut R,
     a: &A,
     log_gap_in: usize,
     log_gap_out: usize,
     log_domain: usize,
-    auto_keys: &HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, BE>>,
+    auto_keys: &H,
     scratch: &mut Scratch<BE>,
 ) where
     R: GLWEToMut,
     A: GLWEToRef,
+    H: GLWEAutomorphismKeyHelper<K, BE>,
+    K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
     M: ModuleLogN + GLWETrace<BE> + GLWEPacking<BE> + GLWERotate<BE>,
     Scratch<BE>: ScratchTakeCore<BE>,
 {

poulpy-schemes/src/tfhe/circuit_bootstrapping/key_prepared.rs

@@ -1,8 +1,8 @@
 use poulpy_core::{
     layouts::{
-        GGLWEInfos, GGLWEToGGSWKeyPrepared, GGLWEToGGSWKeyPreparedFactory, GGSWInfos, GLWEAutomorphismKeyLayout,
+        GGLWEInfos, GGLWEToGGSWKeyPrepared, GGLWEToGGSWKeyPreparedFactory, GGSWInfos, GLWEAutomorphismKeyHelper,
-        GLWEAutomorphismKeyPreparedFactory, GLWEInfos, GLWETensorKeyLayout, GLWETensorKeyPreparedFactory, LWEInfos,
+        GLWEAutomorphismKeyLayout, GLWEAutomorphismKeyPreparedFactory, GLWEInfos, GLWETensorKeyLayout,
-        prepared::GLWEAutomorphismKeyPrepared,
+        GLWETensorKeyPreparedFactory, LWEInfos, prepared::GLWEAutomorphismKeyPrepared,
     },
     trace_galois_elements,
 };
@@ -105,8 +105,20 @@ where
 
 pub struct CircuitBootstrappingKeyPrepared<D: Data, BRA: BlindRotationAlgo, B: Backend> {
     pub(crate) brk: BlindRotationKeyPrepared<D, BRA, B>,
-    pub(crate) tsk: GGLWEToGGSWKeyPrepared<Vec<u8>, B>,
+    pub(crate) tsk: GGLWEToGGSWKeyPrepared<D, B>,
-    pub(crate) atk: HashMap<i64, GLWEAutomorphismKeyPrepared<Vec<u8>, B>>,
+    pub(crate) atk: HashMap<i64, GLWEAutomorphismKeyPrepared<D, B>>,
+}
+
+impl<D: DataRef, BRA: BlindRotationAlgo, BE: Backend> GLWEAutomorphismKeyHelper<GLWEAutomorphismKeyPrepared<D, BE>, BE>
+    for CircuitBootstrappingKeyPrepared<D, BRA, BE>
+{
+    fn get_automorphism_key(&self, k: i64) -> Option<&GLWEAutomorphismKeyPrepared<D, BE>> {
+        self.atk.get_automorphism_key(k)
+    }
+
+    fn automorphism_key_infos(&self) -> poulpy_core::layouts::GGLWELayout {
+        self.atk.automorphism_key_infos()
+    }
 }
 
 impl<D: DataRef, BRA: BlindRotationAlgo, B: Backend> CircuitBootstrappingKeyInfos for CircuitBootstrappingKeyPrepared<D, BRA, B> {