From bed330fecbce0f05664a43de0967ef4054ac936b Mon Sep 17 00:00:00 2001
From: Petar Ivanov <29689712+dartdart26@users.noreply.github.com>
Date: Sun, 25 Aug 2024 17:57:10 +0300
Subject: [PATCH] feat: use tfhe-rs parameters supporting compression

Also, add a small binary to generate keys in fhevm-engine-common.
---
 fhevm-engine/.gitignore                       |  1 +
 .../migrations/20240723111257_coprocessor.sql |  4 +--
 .../src/bin/generate_keys.rs                  | 15 ++++++++++
 fhevm-engine/fhevm-engine-common/src/lib.rs   | 30 +++++++++++++++----
 4 files changed, 42 insertions(+), 8 deletions(-)
 create mode 100644 fhevm-engine/fhevm-engine-common/src/bin/generate_keys.rs

diff --git a/fhevm-engine/.gitignore b/fhevm-engine/.gitignore
index eb5a316c..8f958c4e 100644
--- a/fhevm-engine/.gitignore
+++ b/fhevm-engine/.gitignore
@@ -1 +1,2 @@
 target
+**/fhevm-keys
diff --git a/fhevm-engine/coprocessor/migrations/20240723111257_coprocessor.sql b/fhevm-engine/coprocessor/migrations/20240723111257_coprocessor.sql
index 939311ac..3c1e4f23 100644
--- a/fhevm-engine/coprocessor/migrations/20240723111257_coprocessor.sql
+++ b/fhevm-engine/coprocessor/migrations/20240723111257_coprocessor.sql
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:dc48dac8be30fd63ca9e4c738330893f53cb2fb74798a9dc3030fe79297859b7
-size 248351563
+oid sha256:3533155bc13f8aeea5994e7f82a8f5b1100115f22cb29e6533474529c8b866f4
+size 530459451
diff --git a/fhevm-engine/fhevm-engine-common/src/bin/generate_keys.rs b/fhevm-engine/fhevm-engine-common/src/bin/generate_keys.rs
new file mode 100644
index 00000000..5cb7235b
--- /dev/null
+++ b/fhevm-engine/fhevm-engine-common/src/bin/generate_keys.rs
@@ -0,0 +1,15 @@
+use fhevm_engine_common::generate_fhe_keys;
+
+fn main() {
+    let output_dir = "fhevm-keys";
+    println!("Generating keys...");
+    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"), keys.client_key).unwrap();
+    println!("Creating file {output_dir}/pks");
+    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"), keys.server_key).unwrap();
+}
diff --git a/fhevm-engine/fhevm-engine-common/src/lib.rs b/fhevm-engine/fhevm-engine-common/src/lib.rs
index 4f87f403..54181609 100644
--- a/fhevm-engine/fhevm-engine-common/src/lib.rs
+++ b/fhevm-engine/fhevm-engine-common/src/lib.rs
@@ -1,5 +1,14 @@
-pub mod types;
+use tfhe::{
+    generate_keys,
+    shortint::parameters::{
+        list_compression::COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
+        PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64,
+    },
+    CompactPublicKey, ConfigBuilder,
+};
+
 pub mod tfhe_ops;
+pub mod types;
 
 pub struct FhevmKeys {
     pub server_key: Vec<u8>,
@@ -8,10 +17,19 @@ pub struct FhevmKeys {
 }
 
 pub fn generate_fhe_keys() -> FhevmKeys {
-    let (client_key, server_key) = tfhe::generate_keys(tfhe::ConfigBuilder::default().build());
-    let compact_key = tfhe::CompactPublicKey::new(&client_key);
+    let config =
+        ConfigBuilder::with_custom_parameters(PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64)
+            .enable_compression(COMP_PARAM_MESSAGE_2_CARRY_2_KS_PBS_TUNIFORM_2M64)
+            .build();
+    let (client_key, server_key) = generate_keys(config);
+    let public_key = CompactPublicKey::new(&client_key);
+
     let client_key = bincode::serialize(&client_key).unwrap();
     let server_key = bincode::serialize(&server_key).unwrap();
-    let compact_public_key = bincode::serialize(&compact_key).unwrap();
-    FhevmKeys { server_key, client_key, compact_public_key }
-}
\ No newline at end of file
+    let compact_public_key = bincode::serialize(&public_key).unwrap();
+    FhevmKeys {
+        server_key,
+        client_key,
+        compact_public_key,
+    }
+}