feat(zkvms/risc0): Add benchmarking capability

1 files changed, 7 insertions, 5 deletions

diff --git a/zkvms/risc0/host/src/main.rs b/zkvms/risc0/host/src/main.rs
index 9497360..624b3f2 100644
--- a/zkvms/risc0/host/src/main.rs
+++ b/zkvms/risc0/host/src/main.rs
@@ -1,4 +1,4 @@
-use zkvms_host_io::{Input, Output, foreach_input_field, read_args, RunType::{ Execute, Prove, Verify }};
+use zkvms_host_io::{Input, Output, foreach_input_field, benchmarkable, read_args, RunType::{ Execute, Prove, Verify }};
 use risc0_zkvm::{default_prover, default_executor, ExecutorEnv, Receipt};
 use risc0_zkp::core::digest::Digest;
 use hex::FromHex;
@@ -32,10 +32,10 @@ fn journal(receipt: Receipt) -> Output {
 
 fn main() {
     let run_info = read_args();
-    let env = build_env(&run_info.input);
 
     match run_info.run_type {
-        Execute => {
+        Execute => benchmarkable!{
+            let env = build_env(&run_info.input);
             let exec = default_executor();
             let output = default_executor()
                 .execute(env, HELLO_GUEST_ELF)
@@ -48,14 +48,16 @@ fn main() {
 
             println!("{:#?}", output);
         },
-        Prove => {
+        Prove => benchmarkable!{
+            let env = build_env(&run_info.input);
             let receipt = prove(env);
             println!("Output from journal: {:?}", journal(receipt));
         },
-        Verify => {
+        Verify => benchmarkable!{
             // https://github.com/risc0/risc0/blob/881e512732eca72849b2d0e263a1242aba3158af/risc0/build/src/lib.rs#L197-L199
             let guest_id: Digest = Digest::from_hex(HELLO_GUEST_ID).unwrap();
 
+            let env = build_env(&run_info.input);
             let receipt = prove(env);
             receipt.verify(guest_id).unwrap();