refactor(integer): add compression key types

tfhe/src/high_level_api/backward_compatibility/keys.rs

@@ -211,14 +211,35 @@ impl Upgrade<IntegerClientKeyV1> for IntegerClientKeyV0 {
     }
 }
 
-impl Upgrade<IntegerClientKey> for IntegerClientKeyV1 {
+impl Upgrade<IntegerClientKeyV2> for IntegerClientKeyV1 {
+    type Error = Infallible;
+
+    fn upgrade(self) -> Result<IntegerClientKeyV2, Self::Error> {
+        Ok(IntegerClientKeyV2 {
+            key: self.key,
+            dedicated_compact_private_key: self.dedicated_compact_private_key,
+            compression_key: self.compression_key,
+        })
+    }
+}
+
+#[derive(Version)]
+pub(crate) struct IntegerClientKeyV2 {
+    pub(crate) key: crate::integer::ClientKey,
+    pub(crate) dedicated_compact_private_key: Option<CompactPrivateKey>,
+    pub(crate) compression_key: Option<crate::shortint::list_compression::CompressionPrivateKeys>,
+}
+
+impl Upgrade<IntegerClientKey> for IntegerClientKeyV2 {
     type Error = Infallible;
 
     fn upgrade(self) -> Result<IntegerClientKey, Self::Error> {
         Ok(IntegerClientKey {
             key: self.key,
             dedicated_compact_private_key: self.dedicated_compact_private_key,
-            compression_key: self.compression_key,
+            compression_key: self
+                .compression_key
+                .map(|key| crate::integer::compression_keys::CompressionPrivateKeys { key }),
         })
     }
 }
@@ -228,7 +249,8 @@ impl Upgrade<IntegerClientKey> for IntegerClientKeyV1 {
 pub(crate) enum IntegerClientKeyVersions {
     V0(IntegerClientKeyV0),
     V1(IntegerClientKeyV1),
-    V2(IntegerClientKey),
+    V2(IntegerClientKeyV2),
+    V3(IntegerClientKey),
 }
 
 #[derive(Version)]
@@ -261,11 +283,11 @@ impl Upgrade<IntegerServerKeyV1> for IntegerServerKeyV0 {
     }
 }
 
-impl Upgrade<IntegerServerKey> for IntegerServerKeyV1 {
+impl Upgrade<IntegerServerKeyV2> for IntegerServerKeyV1 {
     type Error = Infallible;
 
-    fn upgrade(self) -> Result<IntegerServerKey, Self::Error> {
-        Ok(IntegerServerKey {
+    fn upgrade(self) -> Result<IntegerServerKeyV2, Self::Error> {
+        Ok(IntegerServerKeyV2 {
             key: self.key,
             cpk_key_switching_key_material: self.cpk_key_switching_key_material,
             compression_key: self.compression_key,
@@ -274,11 +296,38 @@ impl Upgrade<IntegerServerKey> for IntegerServerKeyV1 {
     }
 }
 
+#[derive(Version)]
+pub struct IntegerServerKeyV2 {
+    pub(crate) key: crate::integer::ServerKey,
+    pub(crate) cpk_key_switching_key_material:
+        Option<crate::integer::key_switching_key::KeySwitchingKeyMaterial>,
+    pub(crate) compression_key: Option<crate::shortint::list_compression::CompressionKey>,
+    pub(crate) decompression_key: Option<crate::shortint::list_compression::DecompressionKey>,
+}
+
+impl Upgrade<IntegerServerKey> for IntegerServerKeyV2 {
+    type Error = Infallible;
+
+    fn upgrade(self) -> Result<IntegerServerKey, Self::Error> {
+        Ok(IntegerServerKey {
+            key: self.key,
+            cpk_key_switching_key_material: self.cpk_key_switching_key_material,
+            compression_key: self
+                .compression_key
+                .map(|key| crate::integer::compression_keys::CompressionKey { key }),
+            decompression_key: self
+                .decompression_key
+                .map(|key| crate::integer::compression_keys::DecompressionKey { key }),
+        })
+    }
+}
+
 #[derive(VersionsDispatch)]
 pub enum IntegerServerKeyVersions {
     V0(IntegerServerKeyV0),
     V1(IntegerServerKeyV1),
-    V2(IntegerServerKey),
+    V2(IntegerServerKeyV2),
+    V3(IntegerServerKey),
 }
 
 #[derive(Version)]

tfhe/src/high_level_api/keys/client.rs

@@ -6,8 +6,8 @@ use super::{CompressedServerKey, ServerKey};
 use crate::high_level_api::backward_compatibility::keys::ClientKeyVersions;
 use crate::high_level_api::config::Config;
 use crate::high_level_api::keys::{CompactPrivateKey, IntegerClientKey};
+use crate::integer::compression_keys::CompressionPrivateKeys;
 use crate::prelude::Tagged;
-use crate::shortint::list_compression::CompressionPrivateKeys;
 use crate::shortint::MessageModulus;
 use crate::Tag;
 use concrete_csprng::seeders::Seed;

tfhe/src/high_level_api/keys/inner.rs

@@ -1,12 +1,12 @@
 use crate::core_crypto::commons::generators::DeterministicSeeder;
 use crate::core_crypto::prelude::ActivatedRandomGenerator;
 use crate::high_level_api::backward_compatibility::keys::*;
-use crate::integer::public_key::CompactPublicKey;
-use crate::integer::CompressedCompactPublicKey;
-use crate::shortint::list_compression::{
+use crate::integer::compression_keys::{
     CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, CompressionPrivateKeys,
     DecompressionKey,
 };
+use crate::integer::public_key::CompactPublicKey;
+use crate::integer::CompressedCompactPublicKey;
 use crate::shortint::parameters::list_compression::CompressionParameters;
 use crate::shortint::MessageModulus;
 use crate::Error;
@@ -103,11 +103,11 @@ impl IntegerClientKey {
         let cks = crate::shortint::engine::ShortintEngine::new_from_seeder(&mut seeder)
             .new_client_key(config.block_parameters.into());
 
+        let key = crate::integer::ClientKey::from(cks);
+
         let compression_key = config
             .compression_parameters
-            .map(|params| cks.new_compression_private_key(params));
-
-        let key = crate::integer::ClientKey::from(cks);
+            .map(|params| key.new_compression_private_key(params));
 
         let dedicated_compact_private_key = config
             .dedicated_compact_public_key_parameters
@@ -193,7 +193,7 @@ impl From<IntegerConfig> for IntegerClientKey {
 
         let compression_key = config
             .compression_parameters
-            .map(|params| key.key.new_compression_private_key(params));
+            .map(|params| key.new_compression_private_key(params));
 
         Self {
             key,
@@ -225,7 +225,7 @@ impl IntegerServerKey {
             || (None, None),
             |a| {
                 let (compression_key, decompression_key) =
-                    cks.key.new_compression_decompression_keys(a);
+                    cks.new_compression_decompression_keys(a);
                 (Some(compression_key), Some(decompression_key))
             },
         );
@@ -312,7 +312,6 @@ impl IntegerCompressedServerKey {
                 .as_ref()
                 .map_or((None, None), |compression_private_key| {
                     let (compression_keys, decompression_keys) = client_key
-                        .key
                         .key
                         .new_compressed_compression_decompression_keys(compression_private_key);
 

tfhe/src/high_level_api/keys/server.rs

@@ -5,10 +5,10 @@ use crate::backward_compatibility::keys::{CompressedServerKeyVersions, ServerKey
 #[cfg(feature = "gpu")]
 use crate::core_crypto::gpu::{synchronize_devices, CudaStreams};
 use crate::high_level_api::keys::{IntegerCompressedServerKey, IntegerServerKey};
-use crate::prelude::Tagged;
-use crate::shortint::list_compression::{
+use crate::integer::compression_keys::{
     CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, DecompressionKey,
 };
+use crate::prelude::Tagged;
 use crate::shortint::MessageModulus;
 use crate::Tag;
 use std::sync::Arc;

tfhe/src/integer/backward_compatibility/list_compression.rs

@@ -0,0 +1,30 @@
+use crate::integer::compression_keys::{
+    CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, CompressionPrivateKeys,
+    DecompressionKey,
+};
+use tfhe_versionable::VersionsDispatch;
+
+#[derive(VersionsDispatch)]
+pub enum CompressionKeyVersions {
+    V0(CompressionKey),
+}
+
+#[derive(VersionsDispatch)]
+pub enum DecompressionKeyVersions {
+    V0(DecompressionKey),
+}
+
+#[derive(VersionsDispatch)]
+pub enum CompressedCompressionKeyVersions {
+    V0(CompressedCompressionKey),
+}
+
+#[derive(VersionsDispatch)]
+pub enum CompressedDecompressionKeyVersions {
+    V0(CompressedDecompressionKey),
+}
+
+#[derive(VersionsDispatch)]
+pub enum CompressionPrivateKeysVersions {
+    V0(CompressionPrivateKeys),
+}

tfhe/src/integer/backward_compatibility/mod.rs

@@ -3,6 +3,7 @@
 pub mod ciphertext;
 pub mod client_key;
 pub mod key_switching_key;
+pub mod list_compression;
 pub mod public_key;
 pub mod server_key;
 pub mod wopbs;

tfhe/src/integer/ciphertext/compressed_ciphertext_list.rs

@@ -1,8 +1,8 @@
 use super::{DataKind, Expandable, RadixCiphertext, SignedRadixCiphertext};
 use crate::integer::backward_compatibility::ciphertext::CompressedCiphertextListVersions;
+use crate::integer::compression_keys::{CompressionKey, DecompressionKey};
 use crate::integer::BooleanBlock;
 use crate::shortint::ciphertext::CompressedCiphertextList as ShortintCompressedCiphertextList;
-use crate::shortint::list_compression::{CompressionKey, DecompressionKey};
 use crate::shortint::Ciphertext;
 use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
@@ -84,7 +84,9 @@ impl CompressedCiphertextListBuilder {
     }
 
     pub fn build(&self, comp_key: &CompressionKey) -> CompressedCiphertextList {
-        let packed_list = comp_key.compress_ciphertexts_into_list(&self.ciphertexts);
+        let packed_list = comp_key
+            .key
+            .compress_ciphertexts_into_list(&self.ciphertexts);
 
         CompressedCiphertextList {
             packed_list,
@@ -128,7 +130,7 @@ impl CompressedCiphertextList {
         Some((
             (start_block_index..end_block_index)
                 .into_par_iter()
-                .map(|i| decomp_key.unpack(&self.packed_list, i).unwrap())
+                .map(|i| decomp_key.key.unpack(&self.packed_list, i).unwrap())
                 .collect(),
             current_info,
         ))

tfhe/src/integer/client_key/mod.rs

@@ -20,9 +20,9 @@ use crate::integer::block_decomposition::BlockRecomposer;
 use crate::integer::ciphertext::boolean_value::BooleanBlock;
 use crate::integer::ciphertext::{CompressedCrtCiphertext, CrtCiphertext};
 use crate::integer::client_key::utils::i_crt;
+use crate::integer::compression_keys::{CompressionKey, CompressionPrivateKeys, DecompressionKey};
 use crate::integer::encryption::{encrypt_crt, encrypt_words_radix_impl};
 use crate::shortint::ciphertext::Degree;
-use crate::shortint::list_compression::{CompressionKey, CompressionPrivateKeys, DecompressionKey};
 use crate::shortint::parameters::{CompressionParameters, MessageModulus};
 use crate::shortint::{
     Ciphertext, ClientKey as ShortintClientKey, ShortintParameterSet as ShortintParameters,
@@ -720,14 +720,22 @@ impl ClientKey {
         &self,
         params: CompressionParameters,
     ) -> CompressionPrivateKeys {
-        self.key.new_compression_private_key(params)
+        CompressionPrivateKeys {
+            key: self.key.new_compression_private_key(params),
+        }
     }
 
     pub fn new_compression_decompression_keys(
         &self,
         private_compression_key: &CompressionPrivateKeys,
     ) -> (CompressionKey, DecompressionKey) {
-        self.key
-            .new_compression_decompression_keys(private_compression_key)
+        let (comp_key, decomp_key) = self
+            .key
+            .new_compression_decompression_keys(&private_compression_key.key);
+
+        (
+            CompressionKey { key: comp_key },
+            DecompressionKey { key: decomp_key },
+        )
     }
 }

tfhe/src/integer/client_key/radix.rs

@@ -5,10 +5,10 @@ use crate::core_crypto::prelude::{SignedNumeric, UnsignedNumeric};
 use crate::integer::backward_compatibility::client_key::RadixClientKeyVersions;
 use crate::integer::block_decomposition::{DecomposableInto, RecomposableFrom};
 use crate::integer::ciphertext::{RadixCiphertext, SignedRadixCiphertext};
-use crate::integer::BooleanBlock;
-use crate::shortint::list_compression::{
+use crate::integer::compression_keys::{
     CompressedCompressionKey, CompressedDecompressionKey, CompressionPrivateKeys,
 };
+use crate::integer::BooleanBlock;
 use crate::shortint::{Ciphertext as ShortintCiphertext, PBSParameters as ShortintParameters};
 use serde::{Deserialize, Serialize};
 use tfhe_versionable::Versionize;
@@ -139,9 +139,15 @@ impl RadixClientKey {
         &self,
         private_compression_key: &CompressionPrivateKeys,
     ) -> (CompressedCompressionKey, CompressedDecompressionKey) {
-        self.key
+        let (comp_key, decomp_key) = self
             .key
-            .new_compressed_compression_decompression_keys(private_compression_key)
+            .key
+            .new_compressed_compression_decompression_keys(&private_compression_key.key);
+
+        (
+            CompressedCompressionKey { key: comp_key },
+            CompressedDecompressionKey { key: decomp_key },
+        )
     }
 }
 

tfhe/src/integer/compression_keys.rs

@@ -0,0 +1,66 @@
+use super::ClientKey;
+use crate::integer::backward_compatibility::list_compression::*;
+use serde::{Deserialize, Serialize};
+use tfhe_versionable::Versionize;
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(CompressionPrivateKeysVersions)]
+pub struct CompressionPrivateKeys {
+    pub(crate) key: crate::shortint::list_compression::CompressionPrivateKeys,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(CompressionKeyVersions)]
+pub struct CompressionKey {
+    pub(crate) key: crate::shortint::list_compression::CompressionKey,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(DecompressionKeyVersions)]
+pub struct DecompressionKey {
+    pub(crate) key: crate::shortint::list_compression::DecompressionKey,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(CompressedCompressionKeyVersions)]
+pub struct CompressedCompressionKey {
+    pub(crate) key: crate::shortint::list_compression::CompressedCompressionKey,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(CompressedDecompressionKeyVersions)]
+pub struct CompressedDecompressionKey {
+    pub(crate) key: crate::shortint::list_compression::CompressedDecompressionKey,
+}
+
+impl CompressedCompressionKey {
+    pub fn decompress(&self) -> CompressionKey {
+        CompressionKey {
+            key: self.key.decompress(),
+        }
+    }
+}
+
+impl CompressedDecompressionKey {
+    pub fn decompress(&self) -> DecompressionKey {
+        DecompressionKey {
+            key: self.key.decompress(),
+        }
+    }
+}
+
+impl ClientKey {
+    pub fn new_compressed_compression_decompression_keys(
+        &self,
+        private_compression_key: &CompressionPrivateKeys,
+    ) -> (CompressedCompressionKey, CompressedDecompressionKey) {
+        let (comp_key, decomp_key) = self
+            .key
+            .new_compressed_compression_decompression_keys(&private_compression_key.key);
+
+        (
+            CompressedCompressionKey { key: comp_key },
+            CompressedDecompressionKey { key: decomp_key },
+        )
+    }
+}