4 files changed, 52 insertions, 64 deletions
diff --git a/zkvms/sp1/Cargo.lock b/zkvms/sp1/Cargo.lock
index c3abeae..3b29c89 100644
--- a/zkvms/sp1/Cargo.lock
+++ b/zkvms/sp1/Cargo.lock
@@ -2012,9 +2012,9 @@ dependencies = [
  "hex",
  "serde",
  "serde_json",
- "sp1-helper",
  "sp1-sdk",
  "tracing",
+ "zkvms_host_io",
 ]
 
 [[package]]
@@ -2565,9 +2565,9 @@ dependencies = [
 
 [[package]]
 name = "log"
-version = "0.4.22"
+version = "0.4.25"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24"
+checksum = "04cbf5b083de1c7e0222a7a51dbfdba1cbe1c6ab0b15e29fff3f6c077fd9cd9f"
 
 [[package]]
 name = "lru"
@@ -4376,13 +4376,6 @@ dependencies = [
 ]
 
 [[package]]
-name = "sp1-helper"
-version = "3.4.0"
-dependencies = [
- "sp1-build",
-]
-
-[[package]]
 name = "sp1-primitives"
 version = "3.4.0"
 dependencies = [
@@ -5865,3 +5858,11 @@ version = "0.1.0"
 dependencies = [
  "guests_macro",
 ]
+
+[[package]]
+name = "zkvms_host_io"
+version = "0.1.0"
+dependencies = [
+ "clap",
+ "num-traits",
+]
diff --git a/zkvms/sp1/default.nix b/zkvms/sp1/default.nix
index 6baee1f..9b42e58 100644
--- a/zkvms/sp1/default.nix
+++ b/zkvms/sp1/default.nix
@@ -17,6 +17,7 @@ let
           ./.
           ../../guests
           ../../guests_macro
+          ../../zkvms_host_io
           ../../Vertices-010.in
       ]);
     };
diff --git a/zkvms/sp1/host/Cargo.toml b/zkvms/sp1/host/Cargo.toml
index 40722b6..6ab50a3 100644
--- a/zkvms/sp1/host/Cargo.toml
+++ b/zkvms/sp1/host/Cargo.toml
@@ -11,3 +11,5 @@ clap = { version = "4.0", features = ["derive", "env"] }
 tracing = "0.1.40"
 hex = "0.4.3"
 alloy-sol-types = "0.7.7"
+
+zkvms_host_io = { path = "../../../zkvms_host_io" }
diff --git a/zkvms/sp1/host/src/main.rs b/zkvms/sp1/host/src/main.rs
index bed7c29..28cc319 100644
--- a/zkvms/sp1/host/src/main.rs
+++ b/zkvms/sp1/host/src/main.rs
@@ -1,69 +1,53 @@
-use clap::Parser;
-use sp1_sdk::{ProverClient, SP1Stdin};
+use zkvms_host_io::{read_args, RunType::{ Execute, Prove, Verify }};
+use sp1_sdk::{ProverClient, EnvProver, SP1Stdin, SP1ProofWithPublicValues, SP1VerifyingKey};
 
 /// The ELF (executable and linkable format) file for the Succinct RISC-V zkVM.
 pub const FIBONACCI_ELF: &[u8] = include_bytes!("./guest");
 
-/// The arguments for the command.
-#[derive(Parser, Debug)]
-#[clap(author, version, about, long_about = None)]
-struct Args {
-    #[clap(long)]
-    execute: bool,
+type Input = (Vec<Vec<bool>>, u32, Vec<Vec<u32>>);
 
-    #[clap(long)]
-    prove: bool,
-}
-
-fn main() {
-    // Setup the logger.
-    sp1_sdk::utils::setup_logger();
-
-    // Parse the command line arguments.
-    let args = Args::parse();
-
-    if args.execute == args.prove {
-        eprintln!("Error: You must specify either --execute or --prove");
-        std::process::exit(1);
-    }
-
-    // Setup the prover client.
-    let client = ProverClient::new();
-
-    // Setup the inputs.
+fn build_stdin((graph, colors, coloring): &Input) -> SP1Stdin {
     let mut stdin = SP1Stdin::new();
-
-    let (graph, colors, coloring): (Vec<Vec<bool>>, u32, Vec<Vec<u32>>) = include!(env!("INPUTS"));
-
     stdin.write(&graph);
     stdin.write(&colors);
     stdin.write(&coloring);
+    stdin
+}
 
-    if args.execute {
-        // Execute the program
-        let (output, report) = client.execute(FIBONACCI_ELF, &stdin).run().unwrap();
-        println!("Program executed successfully.");
-
-        // Read the output.
-        let decoded = output;
-        println!("{:?}", decoded);
-
-        // Record the number of cycles executed.
-        println!("Number of cycles: {}", report.total_instruction_count());
-    } else {
-        // Setup the program for proving.
-        let (pk, vk) = client.setup(FIBONACCI_ELF);
+fn prove(client: &EnvProver, stdin: SP1Stdin) -> (SP1ProofWithPublicValues, SP1VerifyingKey) {
+    let (pk, vk) = client.setup(FIBONACCI_ELF);
+    let proof = client
+        .prove(&pk, &stdin)
+        .run()
+        .expect("failed to generate proof");
+    (proof, vk)
+}
 
-        // Generate the proof
-        let proof = client
-            .prove(&pk, &stdin)
-            .run()
-            .expect("failed to generate proof");
+fn main() {
+    let run_info = read_args();
+    let stdin = build_stdin(&run_info.input);
 
-        println!("Successfully generated proof!");
+    sp1_sdk::utils::setup_logger();
+    let client = ProverClient::new();
 
-        // Verify the proof.
-        client.verify(&proof, &vk).expect("failed to verify proof");
-        println!("Successfully verified proof!");
+    match run_info.run_type {
+        Execute => {
+            let (output, report) = client.execute(FIBONACCI_ELF, &stdin).run().unwrap();
+
+            println!("Program executed successfully.");
+            println!("{:?}", output);
+            println!("Number of cycles: {}", report.total_instruction_count());
+        },
+        Prove => {
+            let _ = prove(&client, stdin);
+            println!("Successfully generated proof!");
+        },
+        Verify => {
+            let (proof, vk) = prove(&client, stdin);
+            println!("Successfully generated proof!");
+
+            client.verify(&proof, &vk).expect("failed to verify proof");
+            println!("Successfully verified proof!");
+        },
     }
 }