Ns/feat/noise squashed compression

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

@@ -47,6 +47,8 @@ impl<Scalar: UnsignedInteger> PackedIntegers<Scalar> {
 
         let in_len = slice.len();
 
+        assert!(log_modulus <= Scalar::BITS);
+
         let number_bits_to_pack = in_len * log_modulus;
 
         let len = number_bits_to_pack.div_ceil(Scalar::BITS);

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/core_crypto/fft_impl/common.rs

@@ -18,6 +18,11 @@ pub fn modulus_switch<Scalar: UnsignedInteger>(
     input: Scalar,
     log_modulus: CiphertextModulusLog,
 ) -> Scalar {
+    assert!(log_modulus.0 <= Scalar::BITS);
+    if log_modulus.0 == Scalar::BITS {
+        return input;
+    }
+
     // Flooring output_to_floor is equivalent to rounding the input
     let output_to_floor = input.wrapping_add(Scalar::ONE << (Scalar::BITS - log_modulus.0 - 1));
 

tfhe/src/integer/ciphertext/compressed_ciphertext_list.rs

@@ -169,6 +169,29 @@ mod tests {
     const NB_OPERATOR_TESTS: usize = 10;
     const NUM_BLOCKS: usize = 32;
 
+    #[test]
+    fn test_empty_list_compression() {
+        let params = TEST_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M128.into();
+
+        let (cks, _) = gen_keys::<ShortintParameterSet>(params, IntegerKeyKind::Radix);
+
+        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 builder = CompressedCiphertextListBuilder::new();
+
+        let compressed = builder.build(&compression_key);
+
+        assert_eq!(compressed.len(), 0);
+        assert!(compressed
+            .get::<RadixCiphertext>(0, &decompression_key)
+            .unwrap()
+            .is_none())
+    }
+
     #[test]
     fn test_ciphertext_compression() {
         for (params, comp_params) in [

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

@@ -257,3 +257,14 @@ pub enum CompressedCiphertextListVersions {
 pub enum SquashedNoiseCiphertextVersions {
     V0(SquashedNoiseCiphertext),
 }
+
+#[derive(VersionsDispatch)]
+pub enum CompressedSquashedNoiseCiphertextListVersions {
+    V0(CompressedSquashedNoiseCiphertextList),
+}
+
+#[derive(VersionsDispatch)]
+#[allow(dead_code)]
+pub(crate) enum CompressedSquashedNoiseCiphertextListMetaVersions {
+    V0(CompressedSquashedNoiseCiphertextListMeta),
+}

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

@@ -1,7 +1,14 @@
-use crate::shortint::parameters::noise_squashing::NoiseSquashingParameters;
+use crate::shortint::parameters::noise_squashing::{
+    NoiseSquashingCompressionParameters, NoiseSquashingParameters,
+};
 use tfhe_versionable::VersionsDispatch;
 
 #[derive(VersionsDispatch)]
 pub enum NoiseSquashingParametersVersions {
     V0(NoiseSquashingParameters),
 }
+
+#[derive(VersionsDispatch)]
+pub enum NoiseSquashingCompressionParametersVersions {
+    V0(NoiseSquashingCompressionParameters),
+}

tfhe/src/shortint/ciphertext/compressed_ciphertext_list.rs

@@ -3,20 +3,29 @@ 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::parameters::{AtomicPatternKind, CompressedCiphertextConformanceParams};
-use crate::shortint::{CarryModulus, MessageModulus};
+use crate::error;
+use crate::shortint::backward_compatibility::ciphertext::{
+    CompressedCiphertextListVersions, CompressedSquashedNoiseCiphertextListMetaVersions,
+    CompressedSquashedNoiseCiphertextListVersions,
+};
+use crate::shortint::parameters::{
+    CompressedCiphertextConformanceParams, CompressedSquashedNoiseCiphertextConformanceParams,
+};
+use crate::shortint::{AtomicPatternKind, CarryModulus, MessageModulus};
+
+use super::SquashedNoiseCiphertext;
 
 #[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize, Versionize)]
 #[versionize(CompressedCiphertextListVersions)]
 pub struct CompressedCiphertextList {
-    pub modulus_switched_glwe_ciphertext_list: Vec<CompressedModulusSwitchedGlweCiphertext<u64>>,
-    pub ciphertext_modulus: CiphertextModulus<u64>,
-    pub message_modulus: MessageModulus,
-    pub carry_modulus: CarryModulus,
-    pub atomic_pattern: AtomicPatternKind,
-    pub lwe_per_glwe: LweCiphertextCount,
-    pub count: CiphertextCount,
+    pub(crate) modulus_switched_glwe_ciphertext_list:
+        Vec<CompressedModulusSwitchedGlweCiphertext<u64>>,
+    pub(crate) ciphertext_modulus: CiphertextModulus<u64>,
+    pub(crate) message_modulus: MessageModulus,
+    pub(crate) carry_modulus: CarryModulus,
+    pub(crate) atomic_pattern: AtomicPatternKind,
+    pub(crate) lwe_per_glwe: LweCiphertextCount,
+    pub(crate) count: CiphertextCount,
 }
 
 impl CompressedCiphertextList {
@@ -75,3 +84,124 @@ impl ParameterSetConformant for CompressedCiphertextList {
             && *atomic_pattern == params.atomic_pattern
     }
 }
+
+#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize, Versionize)]
+#[versionize(CompressedSquashedNoiseCiphertextListMetaVersions)]
+pub(crate) struct CompressedSquashedNoiseCiphertextListMeta {
+    pub(crate) message_modulus: MessageModulus,
+    pub(crate) carry_modulus: CarryModulus,
+    pub(crate) lwe_per_glwe: LweCiphertextCount,
+}
+
+/// A compressed list of [`SquashedNoiseCiphertext`].
+#[derive(Clone, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize, Versionize)]
+#[versionize(CompressedSquashedNoiseCiphertextListVersions)]
+pub struct CompressedSquashedNoiseCiphertextList {
+    pub(crate) glwe_ciphertext_list: Vec<CompressedModulusSwitchedGlweCiphertext<u128>>,
+    pub(crate) meta: Option<CompressedSquashedNoiseCiphertextListMeta>,
+}
+
+impl ParameterSetConformant for CompressedSquashedNoiseCiphertextList {
+    type ParameterSet = CompressedSquashedNoiseCiphertextConformanceParams;
+
+    fn is_conformant(&self, params: &CompressedSquashedNoiseCiphertextConformanceParams) -> bool {
+        let Self {
+            glwe_ciphertext_list,
+            meta,
+        } = self;
+
+        let len = glwe_ciphertext_list.len();
+
+        if len == 0 {
+            return true;
+        }
+
+        let Some(meta) = meta.as_ref() else {
+            return false;
+        };
+
+        let last_body_count = glwe_ciphertext_list.last().unwrap().bodies_count().0;
+
+        let count_is_ok = glwe_ciphertext_list[..len - 1]
+            .iter()
+            .all(|a| a.bodies_count() == params.lwe_per_glwe)
+            && last_body_count <= params.lwe_per_glwe.0;
+
+        count_is_ok
+            && glwe_ciphertext_list
+                .iter()
+                .all(|glwe| glwe.is_conformant(&params.ct_params))
+            && meta.lwe_per_glwe.0 <= params.ct_params.polynomial_size.0
+            && meta.lwe_per_glwe == params.lwe_per_glwe
+            && meta.message_modulus == params.message_modulus
+            && meta.carry_modulus == params.carry_modulus
+    }
+}
+
+impl CompressedSquashedNoiseCiphertextList {
+    pub fn len(&self) -> usize {
+        self.glwe_ciphertext_list
+            .iter()
+            .map(|comp_glwe| comp_glwe.bodies_count().0)
+            .sum()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+
+    /// 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> {
+        // Check this first to make sure we don't try to access the metadata if the list is empty
+        if index >= self.len() {
+            return Err(error!(
+                "Tried getting index {index} for CompressedSquashedNoiseCiphertextList \
+                with {} elements, out of bound access.",
+                self.len()
+            ));
+        }
+
+        let meta = self.meta.as_ref().ok_or_else(|| {
+            error!("Missing ciphertext metadata in CompressedSquashedNoiseCiphertextList")
+        })?;
+
+        let lwe_per_glwe = meta.lwe_per_glwe.0;
+        let glwe_idx = index / lwe_per_glwe;
+
+        let glwe = self.glwe_ciphertext_list[glwe_idx].extract();
+
+        let glwe_dimension = glwe.glwe_size().to_glwe_dimension();
+        let polynomial_size = glwe.polynomial_size();
+        let ciphertext_modulus = glwe.ciphertext_modulus();
+
+        let lwe_size = glwe_dimension
+            .to_equivalent_lwe_dimension(polynomial_size)
+            .to_lwe_size();
+
+        let monomial_degree = MonomialDegree(index % lwe_per_glwe);
+
+        let mut extracted_lwe = SquashedNoiseCiphertext::new_zero(
+            lwe_size,
+            ciphertext_modulus,
+            meta.message_modulus,
+            meta.carry_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::atomic_pattern::AtomicPatternClientKey;
 use crate::shortint::client_key::ClientKey;
@@ -183,3 +186,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
+    }
+}