1 files changed, 6 insertions, 15 deletions

diff --git a/zkvms/zkm/host/src/main.rs b/zkvms/zkm/host/src/main.rs
index f608746..3b8d068 100644
--- a/zkvms/zkm/host/src/main.rs
+++ b/zkvms/zkm/host/src/main.rs
@@ -11,16 +11,6 @@ use zkvms_host_io::{
     RunType::{Execute, Prove, Verify},
 };
 
-async fn setup(
-    client_cfg: &mut ClientCfg,
-) {
-    let generation = zkm_recursion::groth16_setup(&client_cfg.vk_path);
-
-    if let Err(e) = generation {
-        panic!("Failed setup! Error: {e}");
-    }
-}
-
 async fn get_proof(
     prover_client: &mut ProverClient,
     prover_input: &mut ProverInput,
@@ -45,7 +35,7 @@ async fn execute(prover_client: &mut ProverClient, prover_input: &mut ProverInpu
 async fn prove(
     prover_client: &mut ProverClient,
     prover_input: &mut ProverInput,
-    vk_path: &String,
+    key_path: &String,
     proof_results_path: &String,
 ) {
     let prover_result = get_proof(prover_client, prover_input).await;
@@ -67,11 +57,11 @@ async fn main() -> Result<()> {
     let elf_path = env::var("ELF_PATH").expect("ELF PATH is missing");
 
     let proof_results_path = run_info.env_or("PROOF_RESULTS_PATH", "/tmp/contracts");
-    let vk_path = run_info.env_or("VERIFYING_KEY_PATH", "/tmp/input");
+    let key_path = run_info.env_or("VERIFYING_KEY_PATH", "/tmp/input");
 
     let mut client_config = ClientCfg {
         zkm_prover_type: "local".to_string(),
-        vk_path: vk_path.to_owned(),
+        key_path: Some(key_path.to_owned()),
         ..Default::default()
     };
 
@@ -89,14 +79,15 @@ async fn main() -> Result<()> {
     let mut prover_input = ProverInput {
         elf: read(elf_path).unwrap(),
         execute_only: run_info.run_type == Execute,
+        snark_setup: run_info.env_then_or("SNARK_SETUP", |flag| flag.parse::<bool>().ok(), false),
         seg_size,
+        proof_results_path: proof_results_path.clone(),
         public_inputstream,
         private_inputstream,
         ..Default::default()
     };
 
     let start = Instant::now();
-    setup(&mut client_config).await;
 
     match run_info.run_type {
         // only excute the guest program without generating the proof.
@@ -105,7 +96,7 @@ async fn main() -> Result<()> {
         },
         // excute the guest program and generate the proof
         Prove => benchmarkable! {
-            prove(&mut prover_client, &mut prover_input, &vk_path, &proof_results_path).await;
+            prove(&mut prover_client, &mut prover_input, &key_path, &proof_results_path).await;
         },
         Verify => unreachable!(),
     }