feat(shortint): add compression for squashed noise ciphertexts

Author: nsarlin-zama

tfhe/src/core_crypto/entities/lwe_packing_keyswitch_key.rs

@@ -402,17 +402,19 @@ impl<Scalar: UnsignedInteger, C: ContainerMut<Element = Scalar>> ContiguousEntit
         Self: 'this;
 }
 
-pub struct LwePackingKeyswitchKeyConformanceParams {
+pub struct LwePackingKeyswitchKeyConformanceParams<Scalar: UnsignedInteger> {
     pub decomp_base_log: DecompositionBaseLog,
     pub decomp_level_count: DecompositionLevelCount,
     pub input_lwe_dimension: LweDimension,
     pub output_glwe_size: GlweSize,
     pub output_polynomial_size: PolynomialSize,
-    pub ciphertext_modulus: CiphertextModulus<u64>,
+    pub ciphertext_modulus: CiphertextModulus<Scalar>,
 }
 
-impl<C: Container<Element = u64>> ParameterSetConformant for LwePackingKeyswitchKey<C> {
-    type ParameterSet = LwePackingKeyswitchKeyConformanceParams;
+impl<Scalar: UnsignedInteger, C: Container<Element = Scalar>> ParameterSetConformant
+    for LwePackingKeyswitchKey<C>
+{
+    type ParameterSet = LwePackingKeyswitchKeyConformanceParams<Scalar>;
 
     fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
         let Self {

tfhe/src/core_crypto/entities/seeded_lwe_packing_keyswitch_key.rs

@@ -461,8 +461,10 @@ impl<Scalar: UnsignedInteger, C: ContainerMut<Element = Scalar>> ContiguousEntit
         Self: 'this;
 }
 
-impl<C: Container<Element = u64>> ParameterSetConformant for SeededLwePackingKeyswitchKey<C> {
-    type ParameterSet = LwePackingKeyswitchKeyConformanceParams;
+impl<Scalar: UnsignedInteger, C: Container<Element = Scalar>> ParameterSetConformant
+    for SeededLwePackingKeyswitchKey<C>
+{
+    type ParameterSet = LwePackingKeyswitchKeyConformanceParams<Scalar>;
 
     fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
         let Self {

tfhe/src/integer/ciphertext/compressed_ciphertext_list.rs

@@ -172,29 +172,26 @@ mod tests {
 
     #[test]
     fn test_empty_list_compression() {
-        for params in [
-            V1_0_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128.into(),
-            V1_0_PARAM_MULTI_BIT_GROUP_2_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64.into(),
-        ] {
-            let (cks, _) = gen_keys::<ShortintParameterSet>(params, IntegerKeyKind::Radix);
+        let (cks, _) = gen_keys::<ShortintParameterSet>(
+            TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128.into(),
+            IntegerKeyKind::Radix,
+        );
 
-            let private_compression_key = cks.new_compression_private_key(
-                V1_0_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128,
-            );
+        let private_compression_key = cks
+            .new_compression_private_key(TEST_COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128);
 
-            let (compression_key, decompression_key) =
-                cks.new_compression_decompression_keys(&private_compression_key);
+        let (compression_key, decompression_key) =
+            cks.new_compression_decompression_keys(&private_compression_key);
 
-            let builder = CompressedCiphertextListBuilder::new();
+        let builder = CompressedCiphertextListBuilder::new();
 
-            let compressed = builder.build(&compression_key);
+        let compressed = builder.build(&compression_key);
 
-            assert_eq!(compressed.len(), 0);
-            assert!(compressed
-                .get::<RadixCiphertext>(0, &decompression_key)
-                .unwrap()
-                .is_none())
-        }
+        assert_eq!(compressed.len(), 0);
+        assert!(compressed
+            .get::<RadixCiphertext>(0, &decompression_key)
+            .unwrap()
+            .is_none())
     }
 
     #[test]

tfhe/src/shortint/backward_compatibility/ciphertext/mod.rs

@@ -148,3 +148,8 @@ pub enum CompressedCiphertextListVersions {
 pub enum SquashedNoiseCiphertextVersions {
     V0(SquashedNoiseCiphertext),
 }
+
+#[derive(VersionsDispatch)]
+pub enum CompressedSquashedNoiseCiphertextListVersions {
+    V0(CompressedSquashedNoiseCiphertextList),
+}

tfhe/src/shortint/backward_compatibility/list_compression.rs

@@ -2,8 +2,9 @@ use tfhe_versionable::deprecation::{Deprecable, Deprecated};
 use tfhe_versionable::VersionsDispatch;
 
 use crate::shortint::list_compression::{
-    CompressedCompressionKey, CompressedDecompressionKey, CompressionKey, CompressionPrivateKeys,
-    DecompressionKey,
+    CompressedCompressionKey, CompressedDecompressionKey, CompressedNoiseSquashingCompressionKey,
+    CompressionKey, CompressionPrivateKeys, DecompressionKey, NoiseSquashingCompressionKey,
+    NoiseSquashingCompressionPrivateKey,
 };
 
 #[derive(VersionsDispatch)]
@@ -43,3 +44,18 @@ pub enum CompressedDecompressionKeyVersions {
 pub enum CompressionPrivateKeysVersions {
     V0(CompressionPrivateKeys),
 }
+
+#[derive(VersionsDispatch)]
+pub enum NoiseSquashingCompressionKeyVersions {
+    V0(NoiseSquashingCompressionKey),
+}
+
+#[derive(VersionsDispatch)]
+pub enum NoiseSquashingCompressionPrivateKeyVersions {
+    V0(NoiseSquashingCompressionPrivateKey),
+}
+
+#[derive(VersionsDispatch)]
+pub enum CompressedNoiseSquashingCompressionKeyVersions {
+    V0(CompressedNoiseSquashingCompressionKey),
+}

tfhe/src/shortint/backward_compatibility/parameters/list_compression.rs

@@ -1,8 +1,15 @@
 use tfhe_versionable::VersionsDispatch;
 
-use super::parameters::list_compression::CompressionParameters;
+use super::parameters::list_compression::{
+    CompressionParameters, NoiseSquashingCompressionParameters,
+};
 
 #[derive(VersionsDispatch)]
 pub enum CompressionParametersVersions {
     V0(CompressionParameters),
 }
+
+#[derive(VersionsDispatch)]
+pub enum NoiseSquashingCompressionParametersVersions {
+    V0(NoiseSquashingCompressionParameters),
+}

tfhe/src/shortint/ciphertext/compressed_ciphertext_list.rs

@@ -3,10 +3,14 @@ use tfhe_versionable::Versionize;
 use self::compressed_modulus_switched_glwe_ciphertext::CompressedModulusSwitchedGlweCiphertext;
 use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::prelude::*;
-use crate::shortint::backward_compatibility::ciphertext::CompressedCiphertextListVersions;
+use crate::shortint::backward_compatibility::ciphertext::{
+    CompressedCiphertextListVersions, CompressedSquashedNoiseCiphertextListVersions,
+};
 use crate::shortint::parameters::CompressedCiphertextConformanceParams;
 use crate::shortint::{CarryModulus, MessageModulus};
 
+use super::SquashedNoiseCiphertext;
+
 #[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize, Versionize)]
 #[versionize(CompressedCiphertextListVersions)]
 pub struct CompressedCiphertextList {
@@ -75,3 +79,61 @@ impl ParameterSetConformant for CompressedCiphertextList {
             && *pbs_order == params.pbs_order
     }
 }
+
+/// A compressed list of [`SquashedNoiseCiphertext`].
+#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize, Versionize)]
+#[versionize(CompressedSquashedNoiseCiphertextListVersions)]
+pub struct CompressedSquashedNoiseCiphertextList {
+    pub glwe_ciphertext_list: GlweCiphertextListOwned<u128>,
+    pub message_modulus: MessageModulus,
+    pub lwe_per_glwe: LweCiphertextCount,
+    pub count: CiphertextCount,
+}
+
+impl CompressedSquashedNoiseCiphertextList {
+    /// Unpack a single ciphertext from the list.
+    ///
+    /// Return an error if the index is greater than the size of the list.
+    ///
+    /// After unpacking, the individual ciphertexts must be decrypted with the
+    /// [`NoiseSquashingPrivateKey`] derived from the [`NoiseSquashingCompressionPrivateKey`] used
+    /// for compression.
+    ///
+    /// [`NoiseSquashingPrivateKey`]: crate::shortint::noise_squashing::NoiseSquashingPrivateKey
+    /// [`NoiseSquashingCompressionPrivateKey`]: crate::shortint::list_compression::NoiseSquashingCompressionPrivateKey
+    pub fn unpack(&self, index: usize) -> Result<SquashedNoiseCiphertext, crate::Error> {
+        if index >= self.count.0 {
+            return Err(crate::Error::new(format!(
+                "Tried getting index {index} for CompressedNoiseSquashedCiphertextList \
+                with {} elements, out of bound access.",
+                self.count.0
+            )));
+        }
+
+        let lwe_per_glwe = self.lwe_per_glwe.0;
+        let glwe_idx = index / lwe_per_glwe;
+
+        let glwe_dimension = self.glwe_ciphertext_list.glwe_size().to_glwe_dimension();
+        let polynomial_size = self.glwe_ciphertext_list.polynomial_size();
+        let ciphertext_modulus = self.glwe_ciphertext_list.ciphertext_modulus();
+
+        let lwe_size = glwe_dimension
+            .to_equivalent_lwe_dimension(polynomial_size)
+            .to_lwe_size();
+
+        let glwe = self.glwe_ciphertext_list.get(glwe_idx);
+
+        let monomial_degree = MonomialDegree(index % lwe_per_glwe);
+
+        let mut extracted_lwe =
+            SquashedNoiseCiphertext::new_zero(lwe_size, ciphertext_modulus, self.message_modulus);
+
+        extract_lwe_sample_from_glwe_ciphertext(
+            &glwe,
+            extracted_lwe.lwe_ciphertext_mut(),
+            monomial_degree,
+        );
+
+        Ok(extracted_lwe)
+    }
+}

tfhe/src/shortint/list_compression/compressed_server_keys.rs

@@ -1,5 +1,7 @@
+use super::server_keys::NoiseSquashingCompressionKeyConformanceParams;
 use super::{
     CompressionKey, CompressionKeyConformanceParams, CompressionPrivateKeys, DecompressionKey,
+    NoiseSquashingCompressionKey,
 };
 use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::fft_impl::fft64::crypto::bootstrap::LweBootstrapKeyConformanceParams;
@@ -12,6 +14,7 @@ use crate::core_crypto::prelude::{
 };
 use crate::shortint::backward_compatibility::list_compression::{
     CompressedCompressionKeyVersions, CompressedDecompressionKeyVersions,
+    CompressedNoiseSquashingCompressionKeyVersions,
 };
 use crate::shortint::client_key::ClientKey;
 use crate::shortint::engine::ShortintEngine;
@@ -179,3 +182,49 @@ impl ParameterSetConformant for CompressedDecompressionKey {
         blind_rotate_key.is_conformant(&params) && *lwe_per_glwe == parameter_set.lwe_per_glwe
     }
 }
+
+#[derive(Clone, Debug, Serialize, Deserialize, Versionize)]
+#[versionize(CompressedNoiseSquashingCompressionKeyVersions)]
+pub struct CompressedNoiseSquashingCompressionKey {
+    pub packing_key_switching_key: SeededLwePackingKeyswitchKey<Vec<u128>>,
+    pub lwe_per_glwe: LweCiphertextCount,
+}
+
+impl CompressedNoiseSquashingCompressionKey {
+    pub fn decompress(&self) -> NoiseSquashingCompressionKey {
+        let packing_key_switching_key = self
+            .packing_key_switching_key
+            .as_view()
+            .decompress_into_lwe_packing_keyswitch_key();
+
+        NoiseSquashingCompressionKey {
+            packing_key_switching_key,
+            lwe_per_glwe: self.lwe_per_glwe,
+        }
+    }
+}
+
+impl ParameterSetConformant for CompressedNoiseSquashingCompressionKey {
+    type ParameterSet = NoiseSquashingCompressionKeyConformanceParams;
+
+    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
+        let Self {
+            packing_key_switching_key,
+            lwe_per_glwe,
+        } = self;
+
+        let params = LwePackingKeyswitchKeyConformanceParams {
+            decomp_base_log: parameter_set.packing_ks_base_log,
+            decomp_level_count: parameter_set.packing_ks_level,
+            input_lwe_dimension: parameter_set
+                .uncompressed_glwe_dimension
+                .to_equivalent_lwe_dimension(parameter_set.uncompressed_polynomial_size),
+            output_glwe_size: parameter_set.packing_ks_glwe_dimension.to_glwe_size(),
+            output_polynomial_size: parameter_set.packing_ks_polynomial_size,
+            ciphertext_modulus: parameter_set.cipherext_modulus,
+        };
+
+        packing_key_switching_key.is_conformant(&params)
+            && *lwe_per_glwe == parameter_set.lwe_per_glwe
+    }
+}

tfhe/src/shortint/list_compression/mod.rs

@@ -1,8 +1,13 @@
 mod compressed_server_keys;
 mod compression;
+mod noise_squashing_compression;
 mod private_key;
 mod server_keys;
 
-pub use compressed_server_keys::{CompressedCompressionKey, CompressedDecompressionKey};
-pub use private_key::CompressionPrivateKeys;
-pub use server_keys::{CompressionKey, CompressionKeyConformanceParams, DecompressionKey};
+pub use compressed_server_keys::{
+    CompressedCompressionKey, CompressedDecompressionKey, CompressedNoiseSquashingCompressionKey,
+};
+pub use private_key::{CompressionPrivateKeys, NoiseSquashingCompressionPrivateKey};
+pub use server_keys::{
+    CompressionKey, CompressionKeyConformanceParams, DecompressionKey, NoiseSquashingCompressionKey,
+};