refactor(integer): compression uses ClientKey instead of RadixClientKey

zama-ai · Jul 24, 2024 · 752cb46 · 752cb46
1 parent 91f05b0
commit 752cb46
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 33 deletions.
diff --git a/tfhe/src/integer/ciphertext/compressed_ciphertext_list.rs b/tfhe/src/integer/ciphertext/compressed_ciphertext_list.rs
@@ -150,25 +150,23 @@ impl CompressedCiphertextList {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::integer::{RadixCiphertext, RadixClientKey};
+    use crate::integer::ClientKey;
     use crate::shortint::parameters::list_compression::COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
     use crate::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
 
     #[test]
     fn test_heterogeneous_ciphertext_compression_ci_run_filter() {
-        let num_blocks = 2;
-
-        let cks = RadixClientKey::new(PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64, num_blocks);
+        let cks = ClientKey::new(PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64);
 
         let private_compression_key =
             cks.new_compression_private_key(COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64);
 
         let (compression_key, decompression_key) =
             cks.new_compression_decompression_keys(&private_compression_key);
 
-        let ct1 = cks.encrypt(3_u32);
+        let ct1 = cks.encrypt_radix(3_u32, 16);
 
-        let ct2 = cks.encrypt_signed(-2);
+        let ct2 = cks.encrypt_signed_radix(-2, 16);
 
         let ct3 = cks.encrypt_bool(true);
 
@@ -178,20 +176,20 @@ mod tests {
             .push(ct3)
             .build(&compression_key);
 
-        let a = compressed.blocks_of(0, &decompression_key).unwrap();
+        let decompressed1 = compressed.get(0, &decompression_key).unwrap().unwrap();
+
+        let decrypted: u32 = cks.decrypt_radix(&decompressed1);
 
-        let decrypted: u32 = cks.decrypt(&RadixCiphertext::from_expanded_blocks(a.0, a.1).unwrap());
         assert_eq!(decrypted, 3_u32);
 
-        let b = compressed.blocks_of(1, &decompression_key).unwrap();
+        let decompressed2 = compressed.get(1, &decompression_key).unwrap().unwrap();
 
-        let decrypted: i32 =
-            cks.decrypt_signed(&SignedRadixCiphertext::from_expanded_blocks(b.0, b.1).unwrap());
+        let decrypted2: i32 = cks.decrypt_signed_radix(&decompressed2);
 
-        assert_eq!(decrypted, -2);
+        assert_eq!(decrypted2, -2);
 
-        let c = compressed.blocks_of(2, &decompression_key).unwrap();
+        let decompressed3 = compressed.get(2, &decompression_key).unwrap().unwrap();
 
-        assert!(cks.decrypt_bool(&BooleanBlock::from_expanded_blocks(c.0, c.1).unwrap()));
+        assert!(cks.decrypt_bool(&decompressed3));
     }
 }
diff --git a/tfhe/src/integer/client_key/mod.rs b/tfhe/src/integer/client_key/mod.rs
@@ -22,7 +22,8 @@ use crate::integer::ciphertext::{CompressedCrtCiphertext, CrtCiphertext};
 use crate::integer::client_key::utils::i_crt;
 use crate::integer::encryption::{encrypt_crt, encrypt_words_radix_impl};
 use crate::shortint::ciphertext::Degree;
-use crate::shortint::parameters::MessageModulus;
+use crate::shortint::list_compression::{CompressionKey, CompressionPrivateKeys, DecompressionKey};
+use crate::shortint::parameters::{CompressionParameters, MessageModulus};
 use crate::shortint::{
     Ciphertext, ClientKey as ShortintClientKey, ShortintParameterSet as ShortintParameters,
 };
@@ -714,4 +715,19 @@ impl ClientKey {
     {
         encrypt_crt(&self.key, message, base_vec, encrypt_block)
     }
+
+    pub fn new_compression_private_key(
+        &self,
+        params: CompressionParameters,
+    ) -> 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)
+    }
 }
diff --git a/tfhe/src/integer/client_key/radix.rs b/tfhe/src/integer/client_key/radix.rs
@@ -6,8 +6,6 @@ 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::{CompressionKey, CompressionPrivateKeys, DecompressionKey};
-use crate::shortint::parameters::list_compression::CompressionParameters;
 use crate::shortint::{Ciphertext as ShortintCiphertext, PBSParameters as ShortintParameters};
 use serde::{Deserialize, Serialize};
 use tfhe_versionable::Versionize;
@@ -133,22 +131,6 @@ impl RadixClientKey {
     pub fn num_blocks(&self) -> usize {
         self.num_blocks
     }
-
-    pub fn new_compression_private_key(
-        &self,
-        params: CompressionParameters,
-    ) -> CompressionPrivateKeys {
-        self.key.key.new_compression_private_key(params)
-    }
-
-    pub fn new_compression_decompression_keys(
-        &self,
-        private_compression_key: &CompressionPrivateKeys,
-    ) -> (CompressionKey, DecompressionKey) {
-        self.key
-            .key
-            .new_compression_decompression_keys(private_compression_key)
-    }
 }
 
 impl From<(ClientKey, usize)> for RadixClientKey {