Merge pull request #8 from zama-ai/davidk/move-tfhe-ops-into-separate…

…-crate

build: move tfhe ops into separate crate

fhevm-engine/coprocessor/Cargo.toml

@@ -20,11 +20,11 @@ regex = "1.10.5"
 lazy_static = "1.5.0"
 clap.workspace = true
 lru = "0.12.3"
-bincode = "1.3.3"
 hex = "0.4"
-strum = { version = "0.26", features = ["derive"] }
 bigdecimal = "0.4"
 fhevm-engine-common = { path = "../fhevm-engine-common" }
+strum = { version = "0.26", features = ["derive"] }
+bincode = "1.3.3"
 
 [dev-dependencies]
 testcontainers = "0.21"

fhevm-engine/coprocessor/src/main.rs

@@ -1,28 +1,24 @@
+use fhevm_engine_common::generate_fhe_keys;
 use tokio::task::JoinSet;
 
 mod cli;
 mod db_queries;
 mod server;
 #[cfg(test)]
 mod tests;
-mod tfhe_ops;
 mod tfhe_worker;
 mod types;
 mod utils;
 
 fn main() {
-
-    // TODO: remove, just to make sure it works
-    let _ = fhevm_engine_common::add(5, 5);
-
     let args = crate::cli::parse_args();
     assert!(
         args.work_items_batch_size < args.tenant_key_cache_size,
         "Work items batch size must be less than tenant key cache size"
     );
 
     if args.generate_fhe_keys {
-        generate_fhe_keys();
+        generate_dump_fhe_keys();
     } else {
         start_runtime(args, None);
     }
@@ -85,20 +81,16 @@ async fn async_main(
     Ok(())
 }
 
-fn generate_fhe_keys() {
+fn generate_dump_fhe_keys() {
     let output_dir = "fhevm-keys";
     println!("Generating keys...");
-    let (client_key, server_key) = tfhe::generate_keys(tfhe::ConfigBuilder::default().build());
-    let compact_key = tfhe::CompactPublicKey::new(&client_key);
-    let client_key = bincode::serialize(&client_key).unwrap();
-    let server_key = bincode::serialize(&server_key).unwrap();
-    let compact_key = bincode::serialize(&compact_key).unwrap();
+    let keys = generate_fhe_keys();
     println!("Creating directory {output_dir}");
     std::fs::create_dir_all(output_dir).unwrap();
     println!("Creating file {output_dir}/cks");
-    std::fs::write(format!("{output_dir}/cks"), client_key).unwrap();
+    std::fs::write(format!("{output_dir}/cks"), keys.client_key).unwrap();
     println!("Creating file {output_dir}/pks");
-    std::fs::write(format!("{output_dir}/pks"), compact_key).unwrap();
+    std::fs::write(format!("{output_dir}/pks"), keys.compact_public_key).unwrap();
     println!("Creating file {output_dir}/sks");
-    std::fs::write(format!("{output_dir}/sks"), server_key).unwrap();
+    std::fs::write(format!("{output_dir}/sks"), keys.server_key).unwrap();
 }

fhevm-engine/coprocessor/src/server.rs

@@ -1,8 +1,7 @@
 use crate::db_queries::{check_if_api_key_is_valid, check_if_ciphertexts_exist_in_db};
 use crate::server::coprocessor::GenericResponse;
-use crate::tfhe_ops::{
-    self, check_fhe_operand_types, current_ciphertext_version, debug_trivial_encrypt_be_bytes,
-};
+use fhevm_engine_common::tfhe_ops::{check_fhe_operand_types, current_ciphertext_version, debug_trivial_encrypt_be_bytes, deserialize_fhe_ciphertext};
+use fhevm_engine_common::types::FhevmError;
 use crate::types::CoprocessorError;
 use crate::utils::sort_computations_by_dependencies;
 use coprocessor::async_computation_input::Input;
@@ -167,7 +166,7 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
             let mut decrypted: Vec<DebugDecryptResponseSingle> = Vec::with_capacity(cts.len());
             for ct in cts {
                 let deserialized =
-                    tfhe_ops::deserialize_fhe_ciphertext(ct.ciphertext_type, &ct.ciphertext)
+                    deserialize_fhe_ciphertext(ct.ciphertext_type, &ct.ciphertext)
                         .unwrap();
                 decrypted.push(DebugDecryptResponseSingle {
                     output_type: ct.ciphertext_type as i32,
@@ -205,7 +204,7 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
             let ciphertext_type: i16 = i_ct
                 .ciphertext_type
                 .try_into()
-                .map_err(|_e| CoprocessorError::UnknownFheType(i_ct.ciphertext_type))?;
+                .map_err(|_e| CoprocessorError::FhevmError(FhevmError::UnknownFheType(i_ct.ciphertext_type)))?;
             let _ = sqlx::query!("
               INSERT INTO ciphertexts(tenant_id, handle, ciphertext, ciphertext_version, ciphertext_type)
               VALUES($1, $2, $3, $4, $5)
@@ -254,6 +253,7 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
             let mut handle_types = Vec::with_capacity(comp.inputs.len());
             let mut is_computation_scalar = false;
             let mut this_comp_inputs: Vec<Vec<u8>> = Vec::with_capacity(comp.inputs.len());
+            let mut is_scalar_op_vec: Vec<bool> = Vec::with_capacity(comp.inputs.len());
             for (idx, ih) in comp.inputs.iter().enumerate() {
                 if let Some(input) = &ih.input {
                     match input {
@@ -263,27 +263,30 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
                                 .expect("this must be found if operand is non scalar");
                             handle_types.push(*ct_type);
                             this_comp_inputs.push(ih.clone());
+                            is_scalar_op_vec.push(false);
                         }
                         Input::Scalar(sc) => {
                             is_computation_scalar = true;
                             handle_types.push(-1);
                             this_comp_inputs.push(sc.clone());
+                            is_scalar_op_vec.push(true);
                             assert!(idx == 1, "we should have checked earlier that only second operand can be scalar");
                         }
                     }
                 }
             }
 
-            computations_inputs.push(this_comp_inputs);
-            are_comps_scalar.push(is_computation_scalar);
             // check before we insert computation that it has
             // to succeed according to the type system
             let output_type = check_fhe_operand_types(
                 comp.operation,
                 &handle_types,
-                is_computation_scalar,
-                &comp.inputs,
-            )?;
+                &this_comp_inputs,
+                &is_scalar_op_vec,
+            ).map_err(|e| CoprocessorError::FhevmError(e))?;
+
+            computations_inputs.push(this_comp_inputs);
+            are_comps_scalar.push(is_computation_scalar);
             // fill in types with output handles that are computed as we go
             assert!(ct_types
                 .insert(comp.output_handle.clone(), output_type)
@@ -300,7 +303,7 @@ impl coprocessor::fhevm_coprocessor_server::FhevmCoprocessor for CoprocessorServ
             let fhe_operation: i16 = comp
                 .operation
                 .try_into()
-                .map_err(|_| CoprocessorError::UnknownFheOperation(comp.operation))?;
+                .map_err(|_| CoprocessorError::FhevmError(FhevmError::UnknownFheOperation(comp.operation)))?;
             let res = query!(
                 "
                     INSERT INTO computations(tenant_id, output_handle, dependencies, fhe_operation, is_completed, is_scalar)

fhevm-engine/coprocessor/src/tests/operators.rs

@@ -7,14 +7,12 @@ use crate::tests::utils::wait_until_all_ciphertexts_computed;
 use crate::{
     server::coprocessor::{async_computation_input::Input, AsyncComputationInput},
     tests::utils::{default_api_key, setup_test_app},
-    tfhe_ops::{
-        does_fhe_operation_support_both_encrypted_operands, does_fhe_operation_support_scalar,
-    },
-    types::{FheOperationType, SupportedFheOperations},
 };
 use bigdecimal::num_bigint::BigInt;
-use std::{ops::Not, str::FromStr};
+use fhevm_engine_common::tfhe_ops::{does_fhe_operation_support_both_encrypted_operands, does_fhe_operation_support_scalar};
+use fhevm_engine_common::types::{FheOperationType, SupportedFheOperations};
 use strum::IntoEnumIterator;
+use std::{ops::Not, str::FromStr};
 use tonic::metadata::MetadataValue;
 
 struct BinaryOperatorTestCase {

fhevm-engine/coprocessor/src/tfhe_worker.rs

@@ -1,7 +1,8 @@
-use crate::tfhe_ops::{
+use fhevm_engine_common::tfhe_ops::{
     current_ciphertext_version, deserialize_fhe_ciphertext, perform_fhe_operation,
 };
-use crate::types::{SupportedFheCiphertexts, TfheTenantKeys};
+use crate::types::TfheTenantKeys;
+use fhevm_engine_common::types::SupportedFheCiphertexts;
 use sqlx::{postgres::PgListener, query, Acquire};
 use std::{
     cell::Cell,
@@ -223,13 +224,19 @@ async fn tfhe_worker_cycle(
                         } else {
                             deserialized_cts.push(
                                 deserialize_fhe_ciphertext(*ct_type, ct_bytes.as_slice())
-                                    .map_err(|e| (e, w.tenant_id, w.output_handle.clone()))?,
+                                    .map_err(|e| {
+                                        let err: Box<dyn std::error::Error + Send + Sync> = Box::new(e);
+                                        (err, w.tenant_id, w.output_handle.clone())
+                                    })?,
                             );
                         }
                     }
 
                     let res = perform_fhe_operation(w.fhe_operation, &deserialized_cts)
-                        .map_err(|e| (e, w.tenant_id, w.output_handle.clone()))?;
+                        .map_err(|e| {
+                            let err: Box<dyn std::error::Error + Send + Sync> = Box::new(e);
+                            (err, w.tenant_id, w.output_handle.clone())
+                        })?;
                     let (db_type, db_bytes) = res.serialize();
 
                     Ok((w, db_type, db_bytes))