Support args (#29)

host-program/src/bin/zkm-prove.rs

@@ -48,6 +48,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 private_inputstream: "".into(),
                 seg_size: 0,
                 execute_only: false,
+                args: "".into(),
             }
         }
     };
@@ -164,6 +165,7 @@ fn set_sha2_rust_input(seg_size_u: u32, execute_only_b: bool) -> anyhow::Result<
         private_inputstream: pri_buf,
         seg_size: seg_size_u,
         execute_only: execute_only_b,
+        args: "".into(),
     };
 
     Ok(input)
@@ -236,6 +238,7 @@ fn set_sha2_go_input(seg_size_u: u32, execute_only_b: bool) -> anyhow::Result<Pr
         private_inputstream: "".into(),
         seg_size: seg_size_u,
         execute_only: execute_only_b,
+        args: "".into(),
     };
 
     Ok(input)
@@ -249,6 +252,7 @@ fn set_mem_alloc_vec_input(seg_size_u: u32, execute_only_b: bool) -> anyhow::Res
         private_inputstream: "".into(),
         seg_size: seg_size_u,
         execute_only: execute_only_b,
+        args: "".into(),
     };
 
     Ok(input)
@@ -263,6 +267,7 @@ fn set_revme_input(seg_size_u: u32, execute_only_b: bool) -> anyhow::Result<Prov
         private_inputstream: "".into(),
         seg_size: seg_size_u,
         execute_only: execute_only_b,
+        args: "".into(),
     };
 
     Ok(input)

sdk/src/local/stark.rs

@@ -13,9 +13,14 @@ pub fn prove_stark(
     let seg_size = input.seg_size as usize;
     let file = ElfBytes::<AnyEndian>::minimal_parse(input.elf.as_slice())
         .expect("opening elf file failed");
+    let mut args: Vec<&str> = input.args.split_whitespace().collect();
+    if args.len() > 2 {
+        args.truncate(2);
+    }
+    log::info!("args: {:?}", args);
     let mut state = State::load_elf(&file);
     state.patch_elf(&file);
-    state.patch_stack(vec![]);
+    state.patch_stack(args);
 
     state.input_stream.push(input.public_inputstream.clone());
     state.input_stream.push(input.private_inputstream.clone());

sdk/src/network/prover.rs

@@ -98,6 +98,7 @@ impl Prover for NetworkProver {
             public_input_stream: input.public_inputstream.clone(),
             private_input_stream: input.private_inputstream.clone(),
             execute_only: input.execute_only,
+            args: input.args.clone(),
             ..Default::default()
         };
         self.sign_ecdsa(&mut request).await;

sdk/src/prover.rs

@@ -10,6 +10,7 @@ pub struct ProverInput {
     pub private_inputstream: Vec<u8>,
     pub seg_size: u32,
     pub execute_only: bool,
+    pub args: String,
 }
 
 #[derive(Debug, Default, Deserialize, Serialize, Clone)]