From c626e4f0ac87399a6ca12f7678e34be90ce4b0f8 Mon Sep 17 00:00:00 2001
From: Kamen Mladenov <kamen@syndamia.com>
Date: Fri, 7 Feb 2025 11:43:50 +0200
Subject: feat(zkvms/risc0): Add benchmarking capability

---
 zkvms/risc0/host/src/main.rs | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

(limited to 'zkvms')

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();
 
-- 
cgit v1.2.3