1 files changed, 31 insertions, 10 deletions

diff --git a/zkvms/zkwasm/host/src/main.rs b/zkvms/zkwasm/host/src/main.rs
index 5ec7139..16caea1 100644
--- a/zkvms/zkwasm/host/src/main.rs
+++ b/zkvms/zkwasm/host/src/main.rs
@@ -3,9 +3,6 @@ use std::io::{self, Write};
 use std::process::{Command, Stdio};
 use regex::Regex;
 
-static K: &str = "19";
-static SCHEME: &str = "shplonk";
-
 fn build_input(input: &Input) -> String {
     let numreg: Regex = Regex::new("(?:^|[^A-Za-z])([0-9]+)").unwrap();
 
@@ -45,32 +42,56 @@ fn run(cmd: &mut Command) {
 fn main() {
     let run_info = read_args();
 
+    let k = run_info
+        .env_or(
+            "ZKWASM_K",
+            "19",
+        );
+
+    let scheme = run_info
+        .env_or(
+            "ZKWASM_SCHEME",
+            "shplonk",
+        );
+
     run(zkwasm_command("setup")
-        .arg("-k").arg(K)
-        .arg("--scheme").arg(SCHEME));
+        .arg("-k").arg(k)
+        .arg("--scheme").arg(scheme));
 
     let input = build_input(&run_info.input);
 
+    let output = run_info
+        .env_or(
+            "ZKWASM_OUTPUT",
+            "./output",
+        );
+
+    let params = run_info
+        .env_or(
+            "ZKWASM_PARAMS",
+            "./params",
+        );
+
     match run_info.run_type {
         Execute => {
             run(zkwasm_command("dry-run")
                 .arg("--private").arg(input)
-                .arg("--output").arg("./output"));
+                .arg("--output").arg(output));
         },
         Prove => {
             run(zkwasm_command("prove")
                 .arg("--private").arg(input)
-                .arg("--output").arg("./output"));
+                .arg("--output").arg(output));
         },
         Verify => {
             run(zkwasm_command("prove")
                 .arg("--private").arg(input)
-                .arg("--output").arg("./output"));
+                .arg("--output").arg(output.clone()));
 
             run(Command::new("zkwasm-cli")
-                .arg("--params").arg("./params")
+                .arg("--params").arg(params)
                 .arg("prog").arg("verify")
-                .arg("--output").arg("./output"));
+                .arg("--output").arg(output));
         },
     }
 }