refactor(integer): add compression key types

tfhe/src/high_level_api/keys/client.rs

@@ -6,7 +6,7 @@ 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::shortint::list_compression::CompressionPrivateKeys;
+use crate::integer::compression_keys::CompressionPrivateKeys;
 use crate::shortint::MessageModulus;
 use concrete_csprng::seeders::Seed;
 use tfhe_versionable::Versionize;

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;
@@ -97,11 +97,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
@@ -187,7 +187,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,
@@ -219,7 +219,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))
             },
         );
@@ -308,7 +308,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

@@ -4,7 +4,7 @@ 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::shortint::list_compression::{
+use crate::integer::compression_keys::{
     CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, DecompressionKey,
 };
 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,7 @@ 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.0.compress_ciphertexts_into_list(&self.ciphertexts);
 
         CompressedCiphertextList {
             packed_list,
@@ -128,7 +128,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.0.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,17 @@ impl ClientKey {
         &self,
         params: CompressionParameters,
     ) -> CompressionPrivateKeys {
-        self.key.new_compression_private_key(params)
+        CompressionPrivateKeys(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.0);
+
+        (CompressionKey(comp_key), DecompressionKey(decomp_key))
     }
 }

tfhe/src/integer/compression_keys.rs

@@ -0,0 +1,56 @@
+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::shortint::list_compression::CompressionPrivateKeys);
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(CompressionKeyVersions)]
+pub struct CompressionKey(pub crate::shortint::list_compression::CompressionKey);
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(DecompressionKeyVersions)]
+pub struct DecompressionKey(pub crate::shortint::list_compression::DecompressionKey);
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(CompressedCompressionKeyVersions)]
+pub struct CompressedCompressionKey(
+    pub crate::shortint::list_compression::CompressedCompressionKey,
+);
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(CompressedDecompressionKeyVersions)]
+pub struct CompressedDecompressionKey(
+    pub crate::shortint::list_compression::CompressedDecompressionKey,
+);
+
+impl CompressedCompressionKey {
+    pub fn decompress(&self) -> CompressionKey {
+        CompressionKey(self.0.decompress())
+    }
+}
+
+impl CompressedDecompressionKey {
+    pub fn decompress(&self) -> DecompressionKey {
+        DecompressionKey(self.0.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.0);
+
+        (
+            CompressedCompressionKey(comp_key),
+            CompressedDecompressionKey(decomp_key),
+        )
+    }
+}

tfhe/src/integer/mod.rs

@@ -55,6 +55,7 @@ pub mod backward_compatibility;
 pub mod bigint;
 pub mod ciphertext;
 pub mod client_key;
+pub mod compression_keys;
 pub mod key_switching_key;
 #[cfg(any(test, feature = "internal-keycache"))]
 pub mod keycache;