feat(all): add server key conformance

zama-ai · Sep 27, 2024 · 43ceb70 · 43ceb70
1 parent 88164e5
commit 43ceb70
Show file tree

Hide file tree

Showing 22 changed files with 1,296 additions and 51 deletions.
diff --git a/tfhe/src/core_crypto/entities/lwe_keyswitch_key.rs b/tfhe/src/core_crypto/entities/lwe_keyswitch_key.rs
@@ -2,6 +2,7 @@
 
 use tfhe_versionable::Versionize;
 
+use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::backward_compatibility::entities::lwe_keyswitch_key::LweKeyswitchKeyVersions;
 use crate::core_crypto::commons::parameters::*;
 use crate::core_crypto::commons::traits::*;
@@ -187,7 +188,12 @@ impl<Scalar: UnsignedInteger, C: Container<Element = Scalar>> LweKeyswitchKey<C>
             "Got an empty container to create an LweKeyswitchKey"
         );
         assert!(
-            container.container_len() % (decomp_level_count.0 * output_lwe_size.0) == 0,
+            container.container_len()
+                % lwe_keyswitch_key_input_key_element_encrypted_size(
+                    decomp_level_count,
+                    output_lwe_size
+                )
+                == 0,
             "The provided container length is not valid. \
         It needs to be dividable by decomp_level_count * output_lwe_size: {}. \
         Got container length: {} and decomp_level_count: {decomp_level_count:?}, \
@@ -424,3 +430,36 @@ impl<Scalar: UnsignedInteger, C: ContainerMut<Element = Scalar>> ContiguousEntit
     where
         Self: 'this;
 }
+
+pub struct KeyswitchKeyConformanceParams {
+    pub decomp_base_log: DecompositionBaseLog,
+    pub decomp_level_count: DecompositionLevelCount,
+    pub output_lwe_size: LweSize,
+    pub input_lwe_dimension: LweDimension,
+    pub ciphertext_modulus: CiphertextModulus<u64>,
+}
+
+impl<C: Container<Element = u64>> ParameterSetConformant for LweKeyswitchKey<C> {
+    type ParameterSet = KeyswitchKeyConformanceParams;
+
+    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
+        let Self {
+            data,
+            decomp_base_log,
+            decomp_level_count,
+            output_lwe_size,
+            ciphertext_modulus,
+        } = self;
+
+        *ciphertext_modulus == parameter_set.ciphertext_modulus
+            && data.container_len()
+                == parameter_set.input_lwe_dimension.0
+                    * lwe_keyswitch_key_input_key_element_encrypted_size(
+                        parameter_set.decomp_level_count,
+                        parameter_set.output_lwe_size,
+                    )
+            && *decomp_base_log == parameter_set.decomp_base_log
+            && *decomp_level_count == parameter_set.decomp_level_count
+            && *output_lwe_size == parameter_set.output_lwe_size
+    }
+}
diff --git a/tfhe/src/core_crypto/entities/lwe_multi_bit_bootstrap_key.rs b/tfhe/src/core_crypto/entities/lwe_multi_bit_bootstrap_key.rs
@@ -1,5 +1,6 @@
 //! Module containing the definition of the [`LweMultiBitBootstrapKey`].
 
+use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::backward_compatibility::entities::lwe_multi_bit_bootstrap_key::{
     FourierLweMultiBitBootstrapKeyVersioned, FourierLweMultiBitBootstrapKeyVersionedOwned,
     LweMultiBitBootstrapKeyVersions,
@@ -690,3 +691,52 @@ impl FourierLweMultiBitBootstrapKeyOwned {
         }
     }
 }
+
+pub struct MultiBitBootstrapKeyConformanceParams {
+    pub decomp_base_log: DecompositionBaseLog,
+    pub decomp_level_count: DecompositionLevelCount,
+    pub input_lwe_dimension: LweDimension,
+    pub output_glwe_size: GlweSize,
+    pub polynomial_size: PolynomialSize,
+    pub grouping_factor: LweBskGroupingFactor,
+    pub ciphertext_modulus: CiphertextModulus<u64>,
+}
+
+impl<C: Container<Element = c64>> ParameterSetConformant for FourierLweMultiBitBootstrapKey<C> {
+    type ParameterSet = MultiBitBootstrapKeyConformanceParams;
+
+    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
+        let Self {
+            fourier:
+                FourierPolynomialList {
+                    data,
+                    polynomial_size,
+                },
+            input_lwe_dimension,
+            glwe_size,
+            decomposition_base_log,
+            decomposition_level_count,
+            grouping_factor,
+        } = self;
+
+        if input_lwe_dimension.0 % grouping_factor.0 != 0 {
+            return false;
+        }
+
+        data.container_len()
+            == lwe_multi_bit_bootstrap_key_size(
+                *input_lwe_dimension,
+                *glwe_size,
+                *polynomial_size,
+                *decomposition_level_count,
+                *grouping_factor,
+            )
+            .unwrap()
+            && *grouping_factor == parameter_set.grouping_factor
+            && *decomposition_base_log == parameter_set.decomp_base_log
+            && *decomposition_level_count == parameter_set.decomp_level_count
+            && *input_lwe_dimension == parameter_set.input_lwe_dimension
+            && *glwe_size == parameter_set.output_glwe_size
+            && *polynomial_size == parameter_set.polynomial_size
+    }
+}
diff --git a/tfhe/src/core_crypto/entities/lwe_packing_keyswitch_key.rs b/tfhe/src/core_crypto/entities/lwe_packing_keyswitch_key.rs
@@ -2,6 +2,7 @@
 
 use tfhe_versionable::Versionize;
 
+use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::backward_compatibility::entities::lwe_packing_keyswitch_key::LwePackingKeyswitchKeyVersions;
 use crate::core_crypto::commons::parameters::*;
 use crate::core_crypto::commons::traits::*;
@@ -396,3 +397,39 @@ impl<Scalar: UnsignedInteger, C: ContainerMut<Element = Scalar>> ContiguousEntit
     where
         Self: 'this;
 }
+
+pub struct PackingKeyswitchConformanceParams {
+    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>,
+}
+
+impl<C: Container<Element = u64>> ParameterSetConformant for LwePackingKeyswitchKey<C> {
+    type ParameterSet = PackingKeyswitchConformanceParams;
+
+    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
+        let Self {
+            data,
+            decomp_base_log,
+            decomp_level_count,
+            output_glwe_size,
+            output_polynomial_size,
+            ciphertext_modulus,
+        } = self;
+
+        data.container_len()
+            == lwe_packing_keyswitch_key_input_key_element_encrypted_size(
+                *decomp_level_count,
+                *output_glwe_size,
+                *output_polynomial_size,
+            ) * parameter_set.input_lwe_dimension.0
+            && *decomp_base_log == parameter_set.decomp_base_log
+            && *decomp_level_count == parameter_set.decomp_level_count
+            && *output_glwe_size == parameter_set.output_glwe_size
+            && *output_polynomial_size == parameter_set.output_polynomial_size
+            && *ciphertext_modulus == parameter_set.ciphertext_modulus
+    }
+}
diff --git a/tfhe/src/core_crypto/entities/seeded_ggsw_ciphertext_list.rs b/tfhe/src/core_crypto/entities/seeded_ggsw_ciphertext_list.rs
@@ -1,7 +1,6 @@
 //! Module containing the definition of the SeededGgswCiphertextList.
 
-use tfhe_versionable::Versionize;
-
+use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::algorithms::*;
 use crate::core_crypto::backward_compatibility::entities::seeded_ggsw_ciphertext_list::SeededGgswCiphertextListVersions;
 use crate::core_crypto::commons::generators::{
@@ -13,6 +12,8 @@ use crate::core_crypto::commons::math::random::{
 use crate::core_crypto::commons::parameters::*;
 use crate::core_crypto::commons::traits::*;
 use crate::core_crypto::entities::*;
+use crate::core_crypto::fft_impl::fft64::crypto::bootstrap::BootstrapKeyConformanceParams;
+use tfhe_versionable::Versionize;
 
 /// A contiguous list containing
 /// [`seeded GGSW ciphertexts`](`crate::core_crypto::entities::SeededGgswCiphertext`).
@@ -466,3 +467,71 @@ impl<Scalar: UnsignedInteger, C: ContainerMut<Element = Scalar>> ContiguousEntit
     where
         Self: 'this;
 }
+
+pub struct GgswCiphertextListConformanceParameters {
+    pub len: usize,
+    pub glwe_size: GlweSize,
+    pub polynomial_size: PolynomialSize,
+    pub decomp_base_log: DecompositionBaseLog,
+    pub decomp_level_count: DecompositionLevelCount,
+    pub ciphertext_modulus: CiphertextModulus<u64>,
+}
+
+impl TryFrom<&MultiBitBootstrapKeyConformanceParams> for GgswCiphertextListConformanceParameters {
+    type Error = ();
+
+    fn try_from(value: &MultiBitBootstrapKeyConformanceParams) -> Result<Self, ()> {
+        if value.input_lwe_dimension.0 % value.grouping_factor.0 != 0 {
+            return Err(());
+        }
+
+        let group_count = value.input_lwe_dimension.0 % value.grouping_factor.0;
+
+        Ok(Self {
+            len: group_count * value.grouping_factor.ggsw_per_multi_bit_element().0,
+            glwe_size: value.output_glwe_size,
+            polynomial_size: value.polynomial_size,
+            decomp_base_log: value.decomp_base_log,
+            decomp_level_count: value.decomp_level_count,
+            ciphertext_modulus: value.ciphertext_modulus,
+        })
+    }
+}
+
+impl From<&BootstrapKeyConformanceParams> for GgswCiphertextListConformanceParameters {
+    fn from(value: &BootstrapKeyConformanceParams) -> Self {
+        Self {
+            len: value.input_lwe_dimension.0,
+            glwe_size: value.output_glwe_size,
+            polynomial_size: value.polynomial_size,
+            decomp_base_log: value.decomp_base_log,
+            decomp_level_count: value.decomp_level_count,
+            ciphertext_modulus: value.ciphertext_modulus,
+        }
+    }
+}
+
+impl<C: Container<Element = u64>> ParameterSetConformant for SeededGgswCiphertextList<C> {
+    type ParameterSet = GgswCiphertextListConformanceParameters;
+
+    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
+        let Self {
+            data,
+            glwe_size,
+            polynomial_size,
+            decomp_base_log,
+            decomp_level_count,
+            compression_seed: _,
+            ciphertext_modulus,
+        } = self;
+
+        data.container_len()
+            == parameter_set.len
+                * seeded_ggsw_ciphertext_size(*glwe_size, *polynomial_size, *decomp_level_count)
+            && *decomp_base_log == parameter_set.decomp_base_log
+            && *decomp_level_count == parameter_set.decomp_level_count
+            && *glwe_size == parameter_set.glwe_size
+            && *polynomial_size == parameter_set.polynomial_size
+            && *ciphertext_modulus == parameter_set.ciphertext_modulus
+    }
+}
diff --git a/tfhe/src/core_crypto/entities/seeded_lwe_bootstrap_key.rs b/tfhe/src/core_crypto/entities/seeded_lwe_bootstrap_key.rs
@@ -2,12 +2,14 @@
 
 use tfhe_versionable::Versionize;
 
+use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::algorithms::*;
 use crate::core_crypto::backward_compatibility::entities::seeded_lwe_bootstrap_key::SeededLweBootstrapKeyVersions;
 use crate::core_crypto::commons::math::random::{ActivatedRandomGenerator, CompressionSeed};
 use crate::core_crypto::commons::parameters::*;
 use crate::core_crypto::commons::traits::*;
 use crate::core_crypto::entities::*;
+use crate::core_crypto::fft_impl::fft64::crypto::bootstrap::BootstrapKeyConformanceParams;
 
 /// A [`seeded LWE bootstrap key`](`SeededLweBootstrapKey`).
 ///
@@ -320,3 +322,15 @@ impl<Scalar: UnsignedInteger> SeededLweBootstrapKeyOwned<Scalar> {
         }
     }
 }
+
+impl<C: Container<Element = u64>> ParameterSetConformant for SeededLweBootstrapKey<C> {
+    type ParameterSet = BootstrapKeyConformanceParams;
+
+    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
+        let Self { ggsw_list } = self;
+
+        let params = parameter_set.into();
+
+        ggsw_list.is_conformant(&params)
+    }
+}
diff --git a/tfhe/src/core_crypto/entities/seeded_lwe_keyswitch_key.rs b/tfhe/src/core_crypto/entities/seeded_lwe_keyswitch_key.rs
@@ -1,13 +1,13 @@
 //! Module containing the definition of the [`SeededLweKeyswitchKey`].
 
-use tfhe_versionable::Versionize;
-
+use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::algorithms::*;
 use crate::core_crypto::backward_compatibility::entities::seeded_lwe_keyswitch_key::SeededLweKeyswitchKeyVersions;
 use crate::core_crypto::commons::math::random::{ActivatedRandomGenerator, CompressionSeed};
 use crate::core_crypto::commons::parameters::*;
 use crate::core_crypto::commons::traits::*;
 use crate::core_crypto::entities::*;
+use tfhe_versionable::Versionize;
 
 /// A [`seeded LWE keyswitch key`](`SeededLweKeyswitchKey`).
 #[derive(Clone, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize, Versionize)]
@@ -420,3 +420,28 @@ impl<Scalar: UnsignedInteger, C: ContainerMut<Element = Scalar>> ContiguousEntit
     where
         Self: 'this;
 }
+
+impl<C: Container<Element = u64>> ParameterSetConformant for SeededLweKeyswitchKey<C> {
+    type ParameterSet = KeyswitchKeyConformanceParams;
+
+    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
+        let Self {
+            data,
+            decomp_base_log,
+            decomp_level_count,
+            output_lwe_size,
+            ciphertext_modulus,
+            compression_seed: _,
+        } = self;
+
+        *ciphertext_modulus == parameter_set.ciphertext_modulus
+            && data.container_len()
+                == parameter_set.input_lwe_dimension.0
+                    * seeded_lwe_keyswitch_key_input_key_element_encrypted_size(
+                        parameter_set.decomp_level_count,
+                    )
+            && *decomp_base_log == parameter_set.decomp_base_log
+            && *decomp_level_count == parameter_set.decomp_level_count
+            && *output_lwe_size == parameter_set.output_lwe_size
+    }
+}
diff --git a/tfhe/src/core_crypto/entities/seeded_lwe_multi_bit_bootstrap_key.rs b/tfhe/src/core_crypto/entities/seeded_lwe_multi_bit_bootstrap_key.rs
@@ -1,7 +1,6 @@
 //! Module containing the definition of the SeededLweBootstrapKey.
 
-use tfhe_versionable::Versionize;
-
+use crate::conformance::ParameterSetConformant;
 use crate::core_crypto::algorithms::*;
 use crate::core_crypto::backward_compatibility::entities::seeded_lwe_multi_bit_bootstrap_key::SeededLweMultiBitBootstrapKeyVersions;
 use crate::core_crypto::commons::generators::{
@@ -13,6 +12,7 @@ use crate::core_crypto::commons::math::random::{
 use crate::core_crypto::commons::parameters::*;
 use crate::core_crypto::commons::traits::*;
 use crate::core_crypto::entities::*;
+use tfhe_versionable::Versionize;
 
 /// A [`seeded LWE multi bit bootstrap key`](`SeededLweMultiBitBootstrapKey`).
 ///
@@ -445,3 +445,18 @@ impl<Scalar: UnsignedInteger> SeededLweMultiBitBootstrapKeyOwned<Scalar> {
         }
     }
 }
+
+impl<C: Container<Element = u64>> ParameterSetConformant for SeededLweMultiBitBootstrapKey<C> {
+    type ParameterSet = MultiBitBootstrapKeyConformanceParams;
+
+    fn is_conformant(&self, parameter_set: &Self::ParameterSet) -> bool {
+        let Self {
+            ggsw_list,
+            grouping_factor,
+        } = self;
+
+        let params = parameter_set.try_into().unwrap();
+
+        ggsw_list.is_conformant(&params) && *grouping_factor == parameter_set.grouping_factor
+    }
+}