fix(zkvms/risc0): Benchmark only the specific function instead of everything during the process

1 files changed, 23 insertions, 15 deletions

diff --git a/zkvms/risc0/host/src/main.rs b/zkvms/risc0/host/src/main.rs
index 6428d84..d9d2af0 100644
--- a/zkvms/risc0/host/src/main.rs
+++ b/zkvms/risc0/host/src/main.rs
@@ -35,34 +35,42 @@ fn main() {
     let run_info = read_args();
 
     match run_info.run_type {
-        Execute => benchmarkable!{
+        Execute => {
             let env = build_env(&run_info.input);
             let exec = default_executor();
-            let output = default_executor()
-                .execute(env, HELLO_GUEST_ELF)
-                .unwrap()
-                .receipt_claim
-                .unwrap()
-                .output
-                .value()
-                .unwrap();
 
-            println!("{:#?}", output);
+            benchmarkable! {
+                let output = default_executor()
+                    .execute(env, HELLO_GUEST_ELF)
+                    .unwrap()
+                    .receipt_claim
+                    .unwrap()
+                    .output
+                    .value()
+                    .unwrap();
+                println!("{:#?}", output);
+            }
         },
-        Prove => benchmarkable!{
+        Prove => {
             let env = build_env(&run_info.input);
-            let receipt = prove(env);
+
+            benchmarkable! {
+                let receipt = prove(env);
+            }
+
             println!("Output from journal: {:?}", journal(receipt));
         },
-        Verify => benchmarkable!{
+        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).unwrap();
 
             let env = build_env(&run_info.input);
             let receipt = prove(env);
-            receipt.verify(guest_id).unwrap();
 
-            println!("Output from verify: {:?}", journal(receipt));
+            benchmarkable! {
+                receipt.verify(guest_id).unwrap();
+                println!("Output from verify: {:?}", journal(receipt));
+            }
         },
     }
 }