1 files changed, 42 insertions, 30 deletions
diff --git a/zkvms/risc0/host/src/main.rs b/zkvms/risc0/host/src/main.rs
index ec1eca3..4077a44 100644
--- a/zkvms/risc0/host/src/main.rs
+++ b/zkvms/risc0/host/src/main.rs
@@ -1,4 +1,5 @@
-use risc0_zkvm::{default_prover, default_executor, ExecutorEnv};
+use zkvms_host_io::{read_args, RunType::{ Execute, Prove, Verify }};
+use risc0_zkvm::{default_prover, default_executor, ExecutorEnv, Receipt};
 use risc0_zkp::core::digest::Digest;
 use hex::FromHex;
 
@@ -7,12 +8,10 @@ static HELLO_GUEST_ELF: &[u8] = include_bytes!("./guest");
 // https://github.com/risc0/risc0/blob/881e512732eca72849b2d0e263a1242aba3158af/risc0/build/src/lib.rs#L255
 static HELLO_GUEST_ID: &str = env!("GUEST_ID");
 
-fn main() {
-    let args: Vec<String> = std::env::args().collect();
-
-    let (graph, colors, coloring): (Vec<Vec<bool>>, u32, Vec<Vec<u32>>) = include!(env!("INPUTS"));
+type Input = (Vec<Vec<bool>>, u32, Vec<Vec<u32>>);
 
-    let env = ExecutorEnv::builder()
+fn build_env((graph, colors, coloring): &Input) -> ExecutorEnv {
+    ExecutorEnv::builder()
         .write(&graph)
         .unwrap()
         .write(&colors)
@@ -20,41 +19,54 @@ fn main() {
         .write(&coloring)
         .unwrap()
         .build()
-        .unwrap();
+        .unwrap()
+}
+
+fn prove(env: ExecutorEnv) -> Receipt {
+    default_prover()
+        .prove(env, HELLO_GUEST_ELF)
+        .expect("Error occured")
+        .receipt
+}
+
+fn journal(receipt: Receipt) -> bool {
+    receipt
+        .journal
+        .decode()
+        .unwrap()
+}
+
+fn main() {
+    let run_info = read_args();
+    let env = build_env(&run_info.input);
 
-    match args[1].as_str() {
-        "execute" => {
+    match run_info.run_type {
+        Execute => {
             let exec = default_executor();
-            let output = exec.execute(env, HELLO_GUEST_ELF).unwrap().receipt_claim.unwrap().output.value().unwrap();
+            let output = default_executor()
+                .execute(env, HELLO_GUEST_ELF)
+                .unwrap()
+                .receipt_claim
+                .unwrap()
+                .output
+                .value()
+                .unwrap();
 
             println!("{:#?}", output);
         },
-        "prove" => {
-            // Obtain the default prover.
-            let prover = default_prover();
-
-            // Produce a receipt by proving the specified ELF binary.
-            let receipt = prover.prove(env, HELLO_GUEST_ELF).expect("Error occured").receipt;
-
-            // Extract journal of receipt
-            let journal: bool = receipt.journal.decode().unwrap();
-            // Print, notice, after committing to a journal, the private input became public
-            println!("Output from journal: {:?}", journal);
+        Prove => {
+            let receipt = prove(env);
+            println!("Output from journal: {:?}", journal(receipt));
         },
-        "verify" => {
+        Verify => {
             // https://github.com/risc0/risc0/blob/881e512732eca72849b2d0e263a1242aba3158af/risc0/build/src/lib.rs#L197-L199
-            let guest_id: Digest = Digest::from_hex(HELLO_GUEST_ID).expect("");
-            // https://github.com/risc0/risc0/blob/881e512732eca72849b2d0e263a1242aba3158af/risc0/build/src/lib.rs#L278
-
-            let prover = default_prover();
-            let receipt = prover.prove(env, HELLO_GUEST_ELF).unwrap().receipt;
+            let guest_id: Digest = Digest::from_hex(HELLO_GUEST_ID).unwrap();
 
+            let receipt = prove(env);
             receipt.verify(guest_id).unwrap();
 
-            let journal: bool = receipt.journal.decode().unwrap();
-            println!("Output from verify: {:?}", journal);
+            println!("Output from verify: {:?}", journal(receipt));
         },
-        _ => println!("No arguments provided! Expected execute, prove or verify!"),
     }
 }