feat(core, ffi, ios): initialize zkey before generate proof

mopro-core/src/middleware/circom/mod.rs

@@ -50,6 +50,16 @@ type CircuitInputs = HashMap<String, Vec<BigInt>>;
 
 // TODO: Split up this namespace a bit, right now quite a lot of things going on
 
+pub struct CircomState2 {
+    zkey: Option<(ProvingKey<Bn254>, ConstraintMatrices<Fr>)>,
+}
+
+impl Default for CircomState2 {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 pub struct CircomState {
     builder: Option<CircomBuilder<Bn254>>,
     circuit: Option<CircomCircuit<Bn254>>,
@@ -245,6 +255,85 @@ pub fn verify_proof2(
     Ok(proof_verified)
 }
 
+impl CircomState2 {
+    pub fn new() -> Self {
+        Self { zkey: None }
+    }
+
+    pub fn initialize(&mut self, zkey_path: &str) -> Result<(), MoproError> {
+        let zkey = load_arkzkey_from_file(zkey_path)?;
+        self.zkey = Some(zkey);
+        Ok(())
+    }
+
+    pub fn generate_proof(
+        &self,
+        inputs: CircuitInputs,
+    ) -> Result<(SerializableProof, SerializableInputs), MoproError> {
+        let mut rng = thread_rng();
+        let rng = &mut rng;
+
+        let r = ark_bn254::Fr::rand(rng);
+        let s = ark_bn254::Fr::rand(rng);
+
+        println!("Generating proof 2");
+
+        let now = std::time::Instant::now();
+        let full_assignment = witness_calculator()
+            .lock()
+            .expect("Failed to lock witness calculator")
+            .calculate_witness_element::<Bn254, _>(inputs, false)
+            .map_err(|e| MoproError::CircomError(e.to_string()))?;
+
+        println!("Witness generation took: {:.2?}", now.elapsed());
+
+        let now = std::time::Instant::now();
+        //let zkey = zkey();
+        let zkey = self.zkey.as_ref().ok_or(MoproError::CircomError(
+            "Zkey has not been set up".to_string(),
+        ))?;
+        println!("Loading arkzkey took: {:.2?}", now.elapsed());
+
+        let public_inputs = full_assignment.as_slice()[1..zkey.1.num_instance_variables].to_vec();
+
+        let now = std::time::Instant::now();
+        let ark_proof = Groth16::<_, CircomReduction>::create_proof_with_reduction_and_matrices(
+            &zkey.0,
+            r,
+            s,
+            &zkey.1,
+            zkey.1.num_instance_variables,
+            zkey.1.num_constraints,
+            full_assignment.as_slice(),
+        );
+
+        let proof = ark_proof.map_err(|e| MoproError::CircomError(e.to_string()))?;
+
+        println!("proof generation took: {:.2?}", now.elapsed());
+        Ok((SerializableProof(proof), SerializableInputs(public_inputs)))
+    }
+
+    pub fn verify_proof(
+        &self,
+        serialized_proof: SerializableProof,
+        serialized_inputs: SerializableInputs,
+    ) -> Result<bool, MoproError> {
+        let start = Instant::now();
+        let zkey = self.zkey.as_ref().ok_or(MoproError::CircomError(
+            "Zkey has not been set up".to_string(),
+        ))?;
+        let pvk = prepare_verifying_key(&zkey.0.vk);
+
+        let proof_verified =
+            GrothBn::verify_with_processed_vk(&pvk, &serialized_inputs.0, &serialized_proof.0)
+                .map_err(|e| MoproError::CircomError(e.to_string()))?;
+
+        let verification_duration = start.elapsed();
+        println!("Verification time 2: {:?}", verification_duration);
+        Ok(proof_verified)
+    }
+}
+
 impl CircomState {
     pub fn new() -> Self {
         Self {
@@ -509,6 +598,57 @@ mod tests {
         assert!(verify_res.unwrap()); // Verifying that the proof was indeed verified
     }
 
+    #[test]
+    fn test_setup_prove_verify_keccak_zkey() {
+        let zkey_path = "./examples/circom/keccak256/target/keccak256_256_test_final.arkzkey";
+        // Instantiate CircomState
+        let mut circom_state = CircomState2::new();
+
+        // Setup
+        let setup_res = circom_state.initialize(zkey_path);
+        assert!(setup_res.is_ok());
+
+        let _serialized_pk = setup_res.unwrap();
+
+        // Deserialize the proving key and inputs if necessary
+
+        // Prepare inputs
+        let input_vec = vec![
+            116, 101, 115, 116, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            0, 0, 0, 0, 0, 0,
+        ];
+
+        // Expected output
+        let expected_output_vec = vec![
+            37, 17, 98, 135, 161, 178, 88, 97, 125, 150, 143, 65, 228, 211, 170, 133, 153, 9, 88,
+            212, 4, 212, 175, 238, 249, 210, 214, 116, 170, 85, 45, 21,
+        ];
+
+        let inputs = bytes_to_circuit_inputs(&input_vec);
+        let serialized_outputs = bytes_to_circuit_outputs(&expected_output_vec);
+
+        // Proof generation
+        let generate_proof_res = circom_state.generate_proof(inputs);
+
+        // Check and print the error if there is one
+        if let Err(e) = &generate_proof_res {
+            println!("Error: {:?}", e);
+        }
+
+        assert!(generate_proof_res.is_ok());
+
+        let (serialized_proof, serialized_inputs) = generate_proof_res.unwrap();
+
+        // Check output
+        assert_eq!(serialized_inputs, serialized_outputs);
+
+        // Proof verification
+        let verify_res = circom_state.verify_proof(serialized_proof, serialized_inputs);
+        assert!(verify_res.is_ok());
+
+        assert!(verify_res.unwrap()); // Verifying that the proof was indeed verified
+    }
+
     #[test]
     fn test_setup_error() {
         // Arrange: Create a new CircomState instance
@@ -740,7 +880,6 @@ mod tests {
     #[ignore = "ignore for ci"]
     #[test]
     fn test_setup_prove_rsa2() {
-
         let zkey_path = "./examples/circom/rsa/target/main_final.arkzkey";
         // Prepare inputs
         let signature = [

mopro-ffi/src/lib.rs

@@ -66,6 +66,16 @@ impl From<mopro_core::MoproError> for FFIError {
     }
 }
 
+pub struct MoproCircom2 {
+    state: RwLock<circom::CircomState2>,
+}
+
+impl Default for MoproCircom2 {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 pub struct MoproCircom {
     state: RwLock<circom::CircomState>,
 }
@@ -130,10 +140,15 @@ pub fn generate_proof2(
     })
 }
 
-pub fn verify_proof2(zkey_path: String,proof: Vec<u8>, public_input: Vec<u8>) -> Result<bool, MoproError> {
+pub fn verify_proof2(
+    zkey_path: String,
+    proof: Vec<u8>,
+    public_input: Vec<u8>,
+) -> Result<bool, MoproError> {
     let deserialized_proof = circom::serialization::deserialize_proof(proof);
     let deserialized_public_input = circom::serialization::deserialize_inputs(public_input);
-    let is_valid = circom::verify_proof2(&zkey_path, deserialized_proof, deserialized_public_input)?;
+    let is_valid =
+        circom::verify_proof2(&zkey_path, deserialized_proof, deserialized_public_input)?;
     Ok(is_valid)
 }
 
@@ -220,6 +235,57 @@ impl MoproCircom {
     }
 }
 
+impl MoproCircom2 {
+    pub fn new() -> Self {
+        Self {
+            state: RwLock::new(circom::CircomState2::new()),
+        }
+    }
+
+    pub fn initialize(&self, zkey_path: String) -> Result<(), MoproError> {
+        let mut state_guard = self.state.write().unwrap();
+        state_guard.initialize(zkey_path.as_str())?;
+        Ok(())
+    }
+
+    //             inputs: circom::serialization::serialize_inputs(&inputs),
+
+    pub fn generate_proof(
+        &self,
+        inputs: HashMap<String, Vec<String>>,
+    ) -> Result<GenerateProofResult, MoproError> {
+        let mut state_guard = self.state.write().unwrap();
+
+        // Convert inputs to BigInt
+        let bigint_inputs = inputs
+            .into_iter()
+            .map(|(k, v)| {
+                (
+                    k,
+                    v.into_iter()
+                        .map(|i| BigInt::from_str(&i).unwrap())
+                        .collect(),
+                )
+            })
+            .collect();
+
+        let (proof, inputs) = state_guard.generate_proof(bigint_inputs)?;
+
+        Ok(GenerateProofResult {
+            proof: circom::serialization::serialize_proof(&proof),
+            inputs: circom::serialization::serialize_inputs(&inputs),
+        })
+    }
+
+    pub fn verify_proof(&self, proof: Vec<u8>, public_input: Vec<u8>) -> Result<bool, MoproError> {
+        let state_guard = self.state.read().unwrap();
+        let deserialized_proof = circom::serialization::deserialize_proof(proof);
+        let deserialized_public_input = circom::serialization::deserialize_inputs(public_input);
+        let is_valid = state_guard.verify_proof(deserialized_proof, deserialized_public_input)?;
+        Ok(is_valid)
+    }
+}
+
 #[cfg(feature = "gpu-benchmarks")]
 pub fn run_msm_benchmark(num_msm: Option<u32>) -> Result<BenchmarkResult, MoproError> {
     let benchmarks = gpu_explorations::run_msm_benchmark(num_msm).unwrap();

mopro-ffi/src/mopro.udl

@@ -72,3 +72,17 @@ interface MoproCircom {
   [Throws=MoproError]
   boolean verify_proof(bytes proof, bytes public_input);
 };
+
+interface MoproCircom2 {
+  constructor();
+
+  [Throws=MoproError]
+  void initialize(string zkey_path);
+
+  [Throws=MoproError]
+  GenerateProofResult generate_proof(record<string, sequence<string>> circuit_inputs);
+
+  [Throws=MoproError]
+  boolean verify_proof(bytes proof, bytes public_input);
+};
+

mopro-ffi/tests/bindings/test_mopro_keccak.swift

@@ -1,10 +1,11 @@
 import mopro
 import Foundation
 
-let moproCircom = MoproCircom()
+let moproCircom = MoproCircom2()
 
 let wasmPath = "./../../../../mopro-core/examples/circom/keccak256/target/keccak256_256_test_js/keccak256_256_test.wasm"
 let r1csPath = "./../../../../mopro-core/examples/circom/keccak256/target/keccak256_256_test.r1cs"
+let zkeyPath = "./../../../../mopro-core/examples/circom/keccak256/target/keccak256_256_test_final.arkzkey"
 
 // Helper function to convert bytes to bits
 func bytesToBits(bytes: [UInt8]) -> [String] {
@@ -45,8 +46,8 @@ func serializeOutputs(_ stringArray: [String]) -> [UInt8] {
 
 do {
     // Setup
-    let setupResult = try moproCircom.setup(wasmPath: wasmPath, r1csPath: r1csPath)
-    assert(!setupResult.provingKey.isEmpty, "Proving key should not be empty")
+    let setupResult = try moproCircom.initialize(zkeyPath: zkeyPath)
+    // assert(!setupResult.provingKey.isEmpty, "Proving key should not be empty")
 
     // Prepare inputs
     let inputVec: [UInt8] = [