chore(coprocessor): add Transaction::BatchInputProofs scenario

Author: goshawk-3

coprocessor/fhevm-engine/stress-test-generator/data/json/batch_input_proofs.json

@@ -0,0 +1,18 @@
+[
+    {
+        "transaction": "BatchInputProofs",
+        "kind": "Rate",
+        "variant": "NoCMUX",
+        "inputs": "NA",
+        "is_dependent": "Independent",
+        "contract_address": "0xa5880e99d86F081E8D3868A8C4732C8f65dfdB08",
+        "user_address": "0xa0534e99d86F081E8D3868A8C4732C8f65dfdB07",
+        "batch_size": 100,
+        "scenario": [
+            [
+                1.0,
+                1
+            ]
+        ]
+    }
+]

coprocessor/fhevm-engine/stress-test-generator/src/bin/stress_generator.rs

@@ -21,6 +21,7 @@ use std::{
 };
 use stress_test_generator::{
     args::parse_args, dex::dex_swap_claim_transaction, utils::new_transaction_id,
+    zk_gen::generate_and_insert_inputs_batch,
 };
 use stress_test_generator::{
     auction::batch_submit_encrypted_bids,
@@ -600,6 +601,19 @@ async fn generate_transaction(
     let ctx = &new_ctx;
 
     match scenario.transaction {
+        Transaction::BatchInputProofs => {
+            let batch_size = scenario.batch_size.unwrap_or(1);
+            generate_and_insert_inputs_batch(
+                ctx,
+                batch_size,
+                MAX_NUMBER_OF_BIDS as u8,
+                &scenario.contract_address,
+                &scenario.user_address,
+            )
+            .await?;
+
+            Ok((Handle::default(), Handle::default()))
+        }
         Transaction::BatchSubmitEncryptedBids => {
             let batch_size = min(MAX_NUMBER_OF_BIDS, scenario.batch_size.unwrap_or(1));
 
@@ -620,8 +634,7 @@ async fn generate_transaction(
             )
             .await?;
 
-            // TODO: how to make dependent if needed?
-            Ok((e_total_payment, e_total_payment)) // TODO: return meaningful second handle
+            Ok((e_total_payment, e_total_payment))
         }
         Transaction::BatchAllowHandles => {
             let mut handles = Vec::new();

coprocessor/fhevm-engine/stress-test-generator/src/utils.rs

@@ -101,6 +101,7 @@ pub enum Transaction {
     GenUsrDecHandles,
     BatchAllowHandles,
     BatchSubmitEncryptedBids,
+    BatchInputProofs,
 }
 
 #[derive(Debug, serde::Deserialize, serde::Serialize, Eq, PartialEq, Clone)]

coprocessor/fhevm-engine/stress-test-generator/src/zk_gen.rs

@@ -57,25 +57,28 @@ impl ZkData {
     }
 }
 
+#[allow(clippy::too_many_arguments)]
 async fn insert_proof(
     pool: &sqlx::PgPool,
     request_id: i64,
-    zk_pok: &[u8],
-    aux: &ZkData,
+    zk_pok: Vec<u8>,
+    aux: ZkData,
     db_notify_channel: &str,
     transaction_id: Handle,
+    retry_count: i32,
     _block_number: u64,
 ) -> Result<(), sqlx::Error> {
     //  Insert ZkPok into database
     sqlx::query(
-            "INSERT INTO verify_proofs (zk_proof_id, input, chain_id, contract_address, user_address, verified, transaction_id)
-            VALUES ($1, $2, $3, $4, $5, NULL, $6)" 
+            "INSERT INTO verify_proofs (zk_proof_id, input, chain_id, contract_address, user_address, verified, transaction_id, retry_count)
+            VALUES ($1, $2, $3, $4, $5, NULL, $6, $7)" 
         ).bind(request_id)
         .bind(zk_pok)
         .bind(aux.chain_id)
         .bind(aux.contract_address.clone())
         .bind(aux.user_address.clone())
         .bind(transaction_id.to_vec())
+        .bind(retry_count)
         .execute(pool).await?;
     sqlx::query("SELECT pg_notify($1, '')")
         .bind(db_notify_channel)
@@ -190,11 +193,12 @@ pub async fn generate_random_handle_vec(
     insert_proof(
         &pool,
         zk_id,
-        &zk_pok,
-        &zk_data,
+        zk_pok,
+        zk_data,
         &ctx.args.zkproof_notify_channel,
         transaction_id,
         0,
+        0,
     )
     .await?;
 
@@ -205,6 +209,67 @@ pub async fn generate_random_handle_vec(
     Ok(handles)
 }
 
+pub async fn generate_and_insert_inputs_batch(
+    ctx: &Context,
+    batch_size: usize,
+    inputs_count: u8,
+    contract_address: &String,
+    user_address: &String,
+) -> Result<(), Box<dyn std::error::Error>> {
+    assert!(inputs_count <= 254);
+    let ecfg = EnvConfig::new();
+    let pool = sqlx::postgres::PgPoolOptions::new()
+        .max_connections(20)
+        .connect(&ecfg.evgen_db_url)
+        .await
+        .unwrap();
+
+    let (pks, public_params) = query_and_save_pks(ecfg.tenant_id, &pool).await?;
+    let mut db_inserts = vec![];
+
+    // Generate a batch of zkpoks
+    for idx in 0..batch_size {
+        let transaction_id = next_random_handle(DEF_TYPE);
+
+        let zk_data = ZkData {
+            contract_address: contract_address.to_owned(),
+            user_address: user_address.to_owned(),
+            acl_contract_address: ecfg.acl_contract_address.clone(),
+            chain_id: ecfg.chain_id,
+        };
+        let aux_data = zk_data.to_owned().assemble()?;
+
+        let mut builder = tfhe::ProvenCompactCiphertextList::builder(&pks);
+        for _ in 0..inputs_count {
+            builder.push(rand::rng().random::<u64>());
+        }
+        let zk_id = ZK_PROOF_ID.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
+        info!(target: "tool", "zkpok id: {}, count = {:?}, seq_num = {} of {} , txn: {:?}, ", zk_id, inputs_count, idx, batch_size, transaction_id );
+        let the_list = builder
+            .build_with_proof_packed(&public_params, &aux_data, tfhe::zk::ZkComputeLoad::Proof)
+            .unwrap();
+
+        let zk_pok = fhevm_engine_common::utils::safe_serialize(&the_list);
+
+        // retry_count = 100 to ensure the txn-sender will delete it after first try
+        // If not deleted, txn-sender will report too many VerifyProofNotRequested errors
+        db_inserts.push(insert_proof(
+            &pool,
+            zk_id,
+            zk_pok.clone(),
+            zk_data,
+            &ctx.args.zkproof_notify_channel,
+            transaction_id,
+            100,
+            0,
+        ));
+    }
+
+    futures::future::try_join_all(db_inserts).await?;
+
+    Ok(())
+}
+
 pub async fn get_inputs_vector(
     ctx: &Context,
     in_type: Inputs,