chore(zk): store inside the pke params the supported zk scheme

tfhe/docs/guides/zk-pok.md

@@ -10,71 +10,7 @@ You can enable this feature using the flag: `--features=zk-pok` when building **
 
 Using this feature is straightforward: during encryption, the client generates the proof, and the server validates it before conducting any homomorphic computations. The following example demonstrates how a client can encrypt and prove a ciphertext, and how a server can verify the ciphertext and compute it:
 
-```rust
-use rand::prelude::*;
-use tfhe::prelude::*;
-use tfhe::set_server_key;
-use tfhe::zk::{CompactPkeCrs, ZkComputeLoad};
-
-pub fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let mut rng = thread_rng();
-
-    let params = tfhe::shortint::parameters::PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64;
-    let config = tfhe::ConfigBuilder::with_custom_parameters(params);
-
-    let client_key = tfhe::ClientKey::generate(config.clone());
-    // This is done in an offline phase and the CRS is shared to all clients and the server
-    let crs = CompactPkeCrs::from_config(config.into(), 64).unwrap();
-    let server_key = tfhe::ServerKey::new(&client_key);
-    let public_key = tfhe::CompactPublicKey::try_new(&client_key).unwrap();
-    // This can be left empty, but if provided allows to tie the proof to arbitrary data
-    let metadata = [b'T', b'F', b'H', b'E', b'-', b'r', b's'];
-
-    let clear_a = rng.gen::<u64>();
-    let clear_b = rng.gen::<u64>();
-
-    let proven_compact_list = tfhe::ProvenCompactCiphertextList::builder(&public_key)
-        .push(clear_a)
-        .push(clear_b)
-        .build_with_proof_packed(&crs, &metadata, ZkComputeLoad::Proof)?;
-
-    // Server side
-    let result = {
-        set_server_key(server_key);
-
-        // Verify the ciphertexts
-        let expander = proven_compact_list.verify_and_expand(&crs, &public_key, &metadata)?;
-        let a: tfhe::FheUint64 = expander.get(0)?.unwrap();
-        let b: tfhe::FheUint64 = expander.get(1)?.unwrap();
-
-        a + b
-    };
-
-    // Back on the client side
-    let a_plus_b: u64 = result.decrypt(&client_key);
-    assert_eq!(a_plus_b, clear_a.wrapping_add(clear_b));
-
-    Ok(())
-}
-```
-
-Performance can be improved by setting `lto="fat"` in `Cargo.toml`
-```toml
-[profile.release]
-lto = "fat"
-```
-and by building the code for the native CPU architecture and in release mode, e.g. by calling `RUSTFLAGS="-C target-cpu=native" cargo run --release`.
-
-{% hint style="info" %}
-You can choose a more costly proof with `ZkComputeLoad::Proof`, which has a faster verification time.  Alternatively, you can select `ZkComputeLoad::Verify` for a faster proof and slower verification.
-{% endhint %}
-
-## Using dedicated Compact Public Key parameters
-
-### A first example
-You can use dedicated parameters for the compact public key encryption to reduce the size of encrypted data and speed up the zero-knowledge proof computation.
-
-This works essentially in the same way as before. Additionally, you need to indicate the dedicated parameters to use:
+Note that you need to use dedicated parameters for the compact public key encryption. This helps to reduce the size of encrypted data and speed up the zero-knowledge proof computation.
 
 ```rust
 use rand::prelude::*;
@@ -132,5 +68,16 @@ pub fn main() -> Result<(), Box<dyn std::error::Error>> {
 }
 ```
 
+Performance can be improved by setting `lto="fat"` in `Cargo.toml`
+```toml
+[profile.release]
+lto = "fat"
+```
+and by building the code for the native CPU architecture and in release mode, e.g. by calling `RUSTFLAGS="-C target-cpu=native" cargo run --release`.
+
+{% hint style="info" %}
+You can choose a more costly proof with `ZkComputeLoad::Proof`, which has a faster verification time.  Alternatively, you can select `ZkComputeLoad::Verify` for a faster proof and slower verification.
+{% endhint %}
+
 ## Benchmark 
 Please refer to the [Zero-knowledge proof benchmarks](../getting_started/benchmarks/zk_proof_benchmarks.md) for detailed performance benchmark results.

tfhe/src/c_api/shortint/parameters.rs

@@ -176,6 +176,7 @@ pub struct ShortintCompactPublicKeyEncryptionParameters {
     // these parameters will always require casting, as they always require casting we add a field
     // for the casting parameters here.
     pub casting_parameters: ShortintCompactCiphertextListCastingParameters,
+    pub zk_scheme: SupportedCompactPkeZkScheme,
 }
 
 impl TryFrom<ShortintCompactPublicKeyEncryptionParameters>
@@ -196,6 +197,7 @@ impl TryFrom<ShortintCompactPublicKeyEncryptionParameters>
             )?,
             expansion_kind:
                 crate::shortint::parameters::CompactCiphertextListExpansionKind::RequiresCasting,
+            zk_scheme: c_params.zk_scheme,
         })
     }
 }
@@ -233,6 +235,7 @@ impl ShortintCompactPublicKeyEncryptionParameters {
             casting_parameters: ShortintCompactCiphertextListCastingParameters::convert(
                 casting_parameters,
             ),
+            zk_scheme: compact_pke_params.zk_scheme,
         }
     }
 }

tfhe/src/integer/ciphertext/compact_list.rs

@@ -1164,9 +1164,7 @@ mod tests {
             .checked_pow(num_blocks as u32)
             .unwrap();
 
-        let crs =
-            CompactPkeCrs::from_shortint_params_legacy_v1(pke_params, LweCiphertextCount(512))
-                .unwrap();
+        let crs = CompactPkeCrs::from_shortint_params(pke_params, LweCiphertextCount(512)).unwrap();
         let cks = ClientKey::new(fhe_params);
         let sk = ServerKey::new_radix_server_key(&cks);
         let compact_private_key = CompactPrivateKey::new(pke_params);

tfhe/src/js_on_wasm_api/shortint.rs

@@ -258,7 +258,8 @@ impl ShortintCompactPublicKeyEncryptionParameters {
             CarryModulus(carry_modulus),
             ciphertext_modulus,
             // These parameters always requires casting
-            crate::shortint::parameters::CompactCiphertextListExpansionKind::RequiresCasting
+            crate::shortint::parameters::CompactCiphertextListExpansionKind::RequiresCasting,
+            crate::shortint::parameters::SupportedCompactPkeZkScheme::ZkNotSupported
         ).map_err(into_js_error)?;
 
         let casting_parameters =

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

@@ -1,14 +1,47 @@
-use tfhe_versionable::VersionsDispatch;
+use std::convert::Infallible;
+
+use tfhe_versionable::{Upgrade, Version, VersionsDispatch};
 
 use super::parameters::compact_public_key_only::CompactPublicKeyEncryptionParameters;
-use super::parameters::CompactCiphertextListExpansionKind;
+use super::parameters::{
+    CompactCiphertextListExpansionKind, DynamicDistribution, SupportedCompactPkeZkScheme,
+};
+use super::prelude::*;
 
 #[derive(VersionsDispatch)]
 pub enum CompactCiphertextListExpansionKindVersions {
     V0(CompactCiphertextListExpansionKind),
 }
 
+#[derive(Version)]
+pub struct CompactPublicKeyEncryptionParametersV0 {
+    pub encryption_lwe_dimension: LweDimension,
+    pub encryption_noise_distribution: DynamicDistribution<u64>,
+    pub message_modulus: MessageModulus,
+    pub carry_modulus: CarryModulus,
+    pub ciphertext_modulus: CiphertextModulus,
+    pub expansion_kind: CompactCiphertextListExpansionKind,
+}
+
+impl Upgrade<CompactPublicKeyEncryptionParameters> for CompactPublicKeyEncryptionParametersV0 {
+    type Error = Infallible;
+
+    fn upgrade(self) -> Result<CompactPublicKeyEncryptionParameters, Self::Error> {
+        Ok(CompactPublicKeyEncryptionParameters {
+            encryption_lwe_dimension: self.encryption_lwe_dimension,
+            encryption_noise_distribution: self.encryption_noise_distribution,
+            message_modulus: self.message_modulus,
+            carry_modulus: self.carry_modulus,
+            ciphertext_modulus: self.ciphertext_modulus,
+            expansion_kind: self.expansion_kind,
+            // TFHE-rs v0.10 and before used only the v1 zk scheme
+            zk_scheme: SupportedCompactPkeZkScheme::V1,
+        })
+    }
+}
+
 #[derive(VersionsDispatch)]
 pub enum CompactPublicKeyEncryptionParametersVersions {
-    V0(CompactPublicKeyEncryptionParameters),
+    V0(CompactPublicKeyEncryptionParametersV0),
+    V1(CompactPublicKeyEncryptionParameters),
 }

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

@@ -4,7 +4,7 @@ pub mod list_compression;
 
 use tfhe_versionable::VersionsDispatch;
 
-use crate::shortint::parameters::ShortintParameterSetInner;
+use crate::shortint::parameters::*;
 use crate::shortint::*;
 
 #[derive(VersionsDispatch)]
@@ -47,3 +47,8 @@ pub enum MultiBitPBSParametersVersions {
 pub enum WopbsParametersVersions {
     V0(WopbsParameters),
 }
+
+#[derive(VersionsDispatch)]
+pub enum SupportedCompactPkeZkSchemeVersions {
+    V0(SupportedCompactPkeZkScheme),
+}

tfhe/src/shortint/ciphertext/zk.rs

@@ -7,21 +7,23 @@ use crate::shortint::ciphertext::CompactCiphertextList;
 use crate::shortint::parameters::{
     CarryModulus, CiphertextListConformanceParams, CiphertextModulus,
     CompactCiphertextListExpansionKind, CompactPublicKeyEncryptionParameters, LweDimension,
-    MessageModulus, ShortintCompactCiphertextListCastingMode,
+    MessageModulus, ShortintCompactCiphertextListCastingMode, SupportedCompactPkeZkScheme,
 };
 use crate::shortint::{Ciphertext, CompactPublicKey};
 use crate::zk::{
     CompactPkeCrs, CompactPkeProof, CompactPkeZkScheme, ZkMSBZeroPaddingBitCount,
     ZkVerificationOutcome,
 };
+
 use rayon::prelude::*;
 use serde::{Deserialize, Serialize};
 use tfhe_versionable::Versionize;
 
 impl CompactPkeCrs {
     /// Construct the CRS that corresponds to the given parameters
     ///
-    /// max_num_message is how many message a single proof can prove
+    /// max_num_message is how many message a single proof can prove.
+    /// The version of the zk scheme is based on the [`CompactPkeZkScheme`] value in the params.
     pub fn from_shortint_params<P, E>(
         params: P,
         max_num_message: LweCiphertextCount,
@@ -43,51 +45,29 @@ impl CompactPkeCrs {
         // 1 padding bit for the PBS
         // Note that if we want to we can prove carry bits are 0 should we need it
         crate::shortint::engine::ShortintEngine::with_thread_local_mut(|engine| {
-            Self::new(
-                size,
-                max_num_message,
-                noise_distribution,
-                params.ciphertext_modulus,
-                plaintext_modulus,
-                ZkMSBZeroPaddingBitCount(1),
-                &mut engine.random_generator,
-            )
-        })
-    }
-
-    /// Construct the CRS for the legacy V1 zk scheme that corresponds to the given parameters
-    ///
-    /// max_num_message is how many message a single proof can prove
-    pub fn from_shortint_params_legacy_v1<P, E>(
-        params: P,
-        max_num_message: LweCiphertextCount,
-    ) -> crate::Result<Self>
-    where
-        P: TryInto<CompactPublicKeyEncryptionParameters, Error = E>,
-        crate::Error: From<E>,
-    {
-        let params: CompactPublicKeyEncryptionParameters = params.try_into()?;
-        let (size, noise_distribution) = (
-            params.encryption_lwe_dimension,
-            params.encryption_noise_distribution,
-        );
-
-        let mut plaintext_modulus = params.message_modulus.0 * params.carry_modulus.0;
-        // Our plaintext modulus does not take into account the bit of padding
-        plaintext_modulus *= 2;
-
-        // 1 padding bit for the PBS
-        // Note that if we want to we can prove carry bits are 0 should we need it
-        crate::shortint::engine::ShortintEngine::with_thread_local_mut(|engine| {
-            Self::new_legacy_v1(
-                size,
-                max_num_message,
-                noise_distribution,
-                params.ciphertext_modulus,
-                plaintext_modulus,
-                ZkMSBZeroPaddingBitCount(1),
-                &mut engine.random_generator,
-            )
+            match params.zk_scheme {
+                SupportedCompactPkeZkScheme::V1 => Self::new_legacy_v1(
+                    size,
+                    max_num_message,
+                    noise_distribution,
+                    params.ciphertext_modulus,
+                    plaintext_modulus,
+                    ZkMSBZeroPaddingBitCount(1),
+                    &mut engine.random_generator,
+                ),
+                SupportedCompactPkeZkScheme::V2 => Self::new(
+                    size,
+                    max_num_message,
+                    noise_distribution,
+                    params.ciphertext_modulus,
+                    plaintext_modulus,
+                    ZkMSBZeroPaddingBitCount(1),
+                    &mut engine.random_generator,
+                ),
+                SupportedCompactPkeZkScheme::ZkNotSupported => {
+                    Err("Zk proof of encryption is not supported by the provided parameters".into())
+                }
+            }
         })
     }
 }

tfhe/src/shortint/parameters/compact_public_key_only/mod.rs

@@ -7,9 +7,10 @@ use crate::shortint::backward_compatibility::parameters::compact_public_key_only
 };
 use crate::shortint::parameters::{
     CarryModulus, ClassicPBSParameters, MessageModulus, MultiBitPBSParameters, PBSParameters,
-    ShortintParameterSet,
+    ShortintParameterSet, SupportedCompactPkeZkScheme,
 };
 use crate::shortint::KeySwitchingKeyView;
+
 use crate::Error;
 use serde::{Deserialize, Serialize};
 use tfhe_versionable::Versionize;
@@ -50,6 +51,8 @@ pub struct CompactPublicKeyEncryptionParameters {
     pub carry_modulus: CarryModulus,
     pub ciphertext_modulus: CiphertextModulus,
     pub expansion_kind: CompactCiphertextListExpansionKind,
+    // Version of the PKE zk scheme compatible with these parameters
+    pub zk_scheme: SupportedCompactPkeZkScheme,
 }
 
 impl CompactPublicKeyEncryptionParameters {
@@ -60,6 +63,7 @@ impl CompactPublicKeyEncryptionParameters {
         carry_modulus: CarryModulus,
         ciphertext_modulus: CiphertextModulus,
         output_ciphertext_kind: CompactCiphertextListExpansionKind,
+        zk_scheme: SupportedCompactPkeZkScheme,
     ) -> Result<Self, Error> {
         let parameters = Self {
             encryption_lwe_dimension,
@@ -68,6 +72,7 @@ impl CompactPublicKeyEncryptionParameters {
             carry_modulus,
             ciphertext_modulus,
             expansion_kind: output_ciphertext_kind,
+            zk_scheme,
         };
 
         if !parameters.is_valid() {
@@ -126,6 +131,8 @@ impl TryFrom<ShortintParameterSet> for CompactPublicKeyEncryptionParameters {
             carry_modulus,
             ciphertext_modulus,
             output_ciphertext_kind,
+            // Zk needs specific pke parameters
+            SupportedCompactPkeZkScheme::ZkNotSupported,
         )
     }
 }

tfhe/src/shortint/parameters/compact_public_key_only/p_fail_2_minus_64/ks_pbs.rs

@@ -3,7 +3,7 @@ use crate::core_crypto::commons::parameters::{
 };
 use crate::shortint::parameters::{
     CarryModulus, CompactCiphertextListExpansionKind, CompactPublicKeyEncryptionParameters,
-    MessageModulus,
+    MessageModulus, SupportedCompactPkeZkScheme,
 };
 
 /// This parameter set should be used when doing zk proof of public key encryption
@@ -19,6 +19,7 @@ pub const V0_11_PARAM_PKE_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV2:
     carry_modulus: CarryModulus(4),
     ciphertext_modulus: CiphertextModulus::new_native(),
     expansion_kind: CompactCiphertextListExpansionKind::RequiresCasting,
+    zk_scheme: SupportedCompactPkeZkScheme::V2,
 }
 .validate();
 
@@ -33,6 +34,7 @@ pub const V0_11_PARAM_PKE_TO_SMALL_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1:
     carry_modulus: CarryModulus(4),
     ciphertext_modulus: CiphertextModulus::new_native(),
     expansion_kind: CompactCiphertextListExpansionKind::RequiresCasting,
+    zk_scheme: SupportedCompactPkeZkScheme::V1,
 }
 .validate();
 
@@ -46,5 +48,6 @@ pub const V0_11_PARAM_PKE_TO_BIG_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64_ZKV1:
     carry_modulus: CarryModulus(4),
     ciphertext_modulus: CiphertextModulus::new_native(),
     expansion_kind: CompactCiphertextListExpansionKind::RequiresCasting,
+    zk_scheme: SupportedCompactPkeZkScheme::V1,
 }
 .validate();