feat(shortint): add try_from_lwe_encryption_key

tfhe/src/high_level_api/tests/mod.rs

@@ -2,10 +2,12 @@ mod tags_on_entities;
 
 use crate::high_level_api::prelude::*;
 use crate::high_level_api::{
-    generate_keys, ClientKey, ConfigBuilder, FheBool, FheUint256, FheUint8, PublicKey,
+    generate_keys, ClientKey, ConfigBuilder, FheBool, FheUint256, FheUint8, PublicKey, ServerKey,
 };
 use crate::integer::U256;
-use crate::{CompactPublicKey, CompressedPublicKey, CompressedServerKey};
+use crate::{
+    set_server_key, CompactPublicKey, CompressedPublicKey, CompressedServerKey, FheUint32, Tag,
+};
 use std::fmt::Debug;
 
 fn assert_that_public_key_encryption_is_decrypted_by_client_key<FheType, ClearType>(
@@ -154,3 +156,33 @@ fn test_serialize_deserialize_are_implemented() {
     can_be_deserialized(&csks);
     can_be_deserialized(&pksz);
 }
+
+#[test]
+fn test_try_from_single_lwe_encryption_key() {
+    let parameters = crate::shortint::prelude::PARAM_MESSAGE_2_CARRY_2_KS_PBS;
+    let lwe_sk = crate::shortint::engine::ShortintEngine::with_thread_local_mut(|engine| {
+        crate::core_crypto::algorithms::allocate_and_generate_new_binary_lwe_secret_key(
+            parameters
+                .glwe_dimension
+                .to_equivalent_lwe_dimension(parameters.polynomial_size),
+            &mut engine.secret_generator,
+        )
+    });
+
+    let shortint_key =
+        crate::shortint::ClientKey::try_from_lwe_encryption_key(lwe_sk, parameters).unwrap();
+    let client_key = ClientKey::from_raw_parts(shortint_key.into(), None, None, Tag::default());
+    let sks = ServerKey::new(&client_key);
+
+    let clear_a = 1344u32;
+    let clear_b = 5u32;
+
+    let encrypted_a = FheUint32::encrypt(clear_a, &client_key);
+    let encrypted_b = FheUint32::encrypt(clear_b, &client_key);
+
+    set_server_key(sks);
+
+    let encrypted_res_mul = &encrypted_a + &encrypted_b;
+    let clear_res: u32 = encrypted_res_mul.decrypt(&client_key);
+    assert_eq!(clear_res, clear_a + clear_b);
+}

tfhe/src/shortint/client_key/mod.rs

@@ -5,7 +5,10 @@ use tfhe_versionable::Versionize;
 
 use super::PBSOrder;
 use crate::core_crypto::entities::*;
-use crate::core_crypto::prelude::decrypt_lwe_ciphertext;
+use crate::core_crypto::prelude::{
+    allocate_and_generate_new_binary_glwe_secret_key,
+    allocate_and_generate_new_binary_lwe_secret_key, decrypt_lwe_ciphertext,
+};
 use crate::shortint::backward_compatibility::client_key::ClientKeyVersions;
 use crate::shortint::ciphertext::{Ciphertext, CompressedCiphertext};
 use crate::shortint::engine::ShortintEngine;
@@ -15,7 +18,7 @@ use crate::shortint::parameters::{
 use crate::shortint::CarryModulus;
 use secret_encryption_key::SecretEncryptionKeyView;
 use serde::{Deserialize, Serialize};
-use std::fmt::Debug;
+use std::fmt::{Debug, Display};
 
 /// A structure containing the client key, which must be kept secret.
 ///
@@ -92,6 +95,68 @@ impl ClientKey {
         }
     }
 
+    pub fn try_from_lwe_encryption_key<P>(
+        encryption_key: LweSecretKeyOwned<u64>,
+        parameters: P,
+    ) -> crate::Result<Self>
+    where
+        P: TryInto<ShortintParameterSet>,
+        <P as TryInto<ShortintParameterSet>>::Error: Display,
+    {
+        let parameters = parameters
+            .try_into()
+            .map_err(|err| crate::Error::new(format!("{err}")))?;
+
+        let expected_lwe_dimension = parameters.encryption_lwe_dimension();
+        if encryption_key.lwe_dimension() != expected_lwe_dimension {
+            return Err(
+                crate::Error::new(
+                    format!(
+                        "The given encryption key does not have the correct LweDimension, expected: {:?}, got: {:?}",
+                        encryption_key.lwe_dimension(),
+                        expected_lwe_dimension)));
+        }
+
+        // The key we got is the one used to encrypt,
+        // we have to generate the other key
+        match parameters.encryption_key_choice() {
+            EncryptionKeyChoice::Big => {
+                // We have to generate the small lwe key
+                let small_key = ShortintEngine::with_thread_local_mut(|engine| {
+                    allocate_and_generate_new_binary_lwe_secret_key(
+                        parameters.lwe_dimension(),
+                        &mut engine.secret_generator,
+                    )
+                });
+
+                Ok(Self {
+                    glwe_secret_key: GlweSecretKeyOwned::from_container(
+                        encryption_key.into_container(),
+                        parameters.polynomial_size(),
+                    ),
+                    lwe_secret_key: small_key,
+                    parameters,
+                })
+            }
+            EncryptionKeyChoice::Small => {
+                // We have to generate the big lwe key
+                let glwe_secret_key = ShortintEngine::with_thread_local_mut(|engine| {
+                    allocate_and_generate_new_binary_glwe_secret_key(
+                        parameters.glwe_dimension(),
+                        parameters.polynomial_size(),
+                        &mut engine.secret_generator,
+                    )
+                });
+
+                Ok(Self {
+                    glwe_secret_key,
+                    lwe_secret_key: encryption_key,
+                    parameters,
+                })
+            }
+        }
+    }
+
     /// Deconstruct a [`ClientKey`] into its constituents.
     ///
     /// # Example