feat: Add estimate subcommand

cli/Cargo.toml

@@ -16,8 +16,10 @@ linux = ["risc0-zkvm/cuda"]
 anyhow = "1.0.86"
 atty = "0.2.14"
 bincode = "1.3.3"
+bonsol-interface.workspace = true
 bonsol-prover = { path = "../prover" }
 bonsol-sdk = { path = "../sdk" }
+bytemuck = "1.15.0"
 hex = "0.4.3"
 byte-unit = "4.0.19"
 bytes = "1.4.0"
@@ -34,7 +36,9 @@ reqwest = { version = "0.11.26", features = [
   "native-tls-vendored",
 ] }
 risc0-binfmt = { workspace = true }
-risc0-zkvm = { workspace = true, features = ["prove"] }
+risc0-zkvm = { workspace = true, default-features = false, features = ["prove", "std"] }
+risc0-zkvm-platform = { git = "https://github.com/anagrambuild/risc0", branch = "v1.0.1-bonsai-fix" }
+risc0-circuit-rv32im = { git = "https://github.com/anagrambuild/risc0", branch = "v1.0.1-bonsai-fix" }
 serde = { version = "1.0.197", features = ["derive"] }
 serde_json = "1.0.104"
 sha2 = "0.10.6"
@@ -46,4 +50,6 @@ tera = "1.17.1"
 thiserror = "1.0.65"
 tokio = { version = "1.38.0", features = ["full"] }
 
-bonsol-interface.workspace = true
+[dev-dependencies]
+assert_cmd = "2.0.16"
+predicates = "3.1.2"

cli/src/command.rs

@@ -227,6 +227,7 @@ pub enum Command {
         )]
         auto_confirm: bool,
     },
+    #[command(about = "Build a ZK program")]
     Build {
         #[arg(
             help = "The path to a ZK program folder containing a Cargo.toml",
@@ -235,6 +236,25 @@ pub enum Command {
         )]
         zk_program_path: String,
     },
+    #[command(about = "Estimate the execution cost of a ZK RISC0 program")]
+    Estimate {
+        #[arg(
+            help = "The path to the program's manifest file (manifest.json)",
+            short = 'm',
+            long
+        )]
+        manifest_path: String,
+
+        #[arg(help = "The path to the program input file", short = 'i', long)]
+        input_file: Option<String>,
+
+        #[arg(
+            help = "Set the maximum number of cycles [default: 16777216u64]",
+            short = 'c',
+            long
+        )]
+        max_cycles: Option<u64>,
+    },
     Execute {
         #[arg(short = 'f', long)]
         execution_request_file: Option<String>,
@@ -296,6 +316,11 @@ pub enum ParsedCommand {
     Build {
         zk_program_path: String,
     },
+    Estimate {
+        manifest_path: String,
+        input_file: Option<String>,
+        max_cycles: Option<u64>,
+    },
     Execute {
         execution_request_file: Option<String>,
 
@@ -351,6 +376,15 @@ impl TryFrom<Command> for ParsedCommand {
                 ),
             }),
             Command::Build { zk_program_path } => Ok(ParsedCommand::Build { zk_program_path }),
+            Command::Estimate {
+                manifest_path,
+                input_file,
+                max_cycles,
+            } => Ok(ParsedCommand::Estimate {
+                manifest_path,
+                input_file,
+                max_cycles,
+            }),
             Command::Execute {
                 execution_request_file,
                 program_id,

cli/src/estimate.rs

@@ -0,0 +1,74 @@
+//! Bare bones upper bound estimator that uses the rv32im
+//! emulation utils for fast lookups in the opcode list
+//! to extract the cycle count from an elf.
+
+use anyhow::Result;
+use risc0_binfmt::{MemoryImage, Program};
+use risc0_zkvm::{ExecutorEnv, ExecutorImpl, Session, GUEST_MAX_MEM};
+use risc0_zkvm_platform::PAGE_SIZE;
+
+pub fn estimate<E: MkImage>(elf: E, env: ExecutorEnv) -> Result<()> {
+    let session = get_session(elf, env)?;
+    println!(
+        "User cycles: {}\nTotal cycles: {}\nSegments: {}",
+        session.user_cycles,
+        session.total_cycles,
+        session.segments.len()
+    );
+
+    Ok(())
+}
+
+/// Get the total number of cycles by stepping through the ELF using emulation
+/// tools from the risc0_circuit_rv32im module.
+pub fn get_session<E: MkImage>(elf: E, env: ExecutorEnv) -> Result<Session> {
+    Ok(ExecutorImpl::new(env, elf.mk_image()?)?.run()?)
+}
+
+/// Helper trait for loading an image from an elf.
+pub trait MkImage {
+    fn mk_image(self) -> Result<MemoryImage>;
+}
+impl<'a> MkImage for &'a [u8] {
+    fn mk_image(self) -> Result<MemoryImage> {
+        let program = Program::load_elf(self, GUEST_MAX_MEM as u32)?;
+        MemoryImage::new(&program, PAGE_SIZE as u32)
+    }
+}
+
+#[cfg(test)]
+mod estimate_tests {
+    use anyhow::Result;
+    use risc0_binfmt::MemoryImage;
+    use risc0_circuit_rv32im::prove::emu::{
+        exec::DEFAULT_SEGMENT_LIMIT_PO2,
+        testutil::{basic as basic_test_program, DEFAULT_SESSION_LIMIT},
+    };
+    use risc0_zkvm::{ExecutorEnv, PAGE_SIZE};
+
+    use super::MkImage;
+    use crate::estimate;
+
+    impl MkImage for MemoryImage {
+        fn mk_image(self) -> Result<MemoryImage> {
+            Ok(self)
+        }
+    }
+
+    #[test]
+    fn estimate_basic() {
+        let program = basic_test_program();
+        let mut env = &mut ExecutorEnv::builder();
+        env = env
+            .segment_limit_po2(DEFAULT_SEGMENT_LIMIT_PO2 as u32)
+            .session_limit(DEFAULT_SESSION_LIMIT);
+        let image = MemoryImage::new(&program, PAGE_SIZE as u32)
+            .expect("failed to create image from basic program");
+        let res = estimate::get_session(image, env.build().unwrap());
+
+        assert_eq!(
+            res.ok().and_then(|session| Some(session.total_cycles)),
+            Some(16384)
+        );
+    }
+}

cli/src/main.rs

@@ -1,22 +1,31 @@
+use std::fs;
 use std::io::{self, Read};
 use std::path::Path;
 
 use atty::Stream;
 use bonsol_sdk::BonsolClient;
 use clap::Parser;
+use common::{execute_get_inputs, ZkProgramManifest};
+use risc0_circuit_rv32im::prove::emu::exec::DEFAULT_SEGMENT_LIMIT_PO2;
+use risc0_circuit_rv32im::prove::emu::testutil::DEFAULT_SESSION_LIMIT;
+use risc0_zkvm::ExecutorEnv;
 use solana_sdk::signature::read_keypair_file;
 use solana_sdk::signer::Signer;
 
 use crate::command::{BonsolCli, ParsedBonsolCli, ParsedCommand};
 use crate::common::{sol_check, try_load_from_config};
-use crate::error::BonsolCliError;
+use crate::error::{BonsolCliError, ZkManifestError};
 
 mod build;
 mod deploy;
+mod estimate;
 mod execute;
 mod init;
 mod prove;
 
+#[cfg(test)]
+mod tests;
+
 pub mod command;
 pub mod common;
 pub(crate) mod error;
@@ -59,6 +68,42 @@ async fn main() -> anyhow::Result<()> {
             }
             deploy::deploy(rpc, keypair, deploy_args).await
         }
+        ParsedCommand::Estimate {
+            manifest_path,
+            input_file,
+            max_cycles,
+        } => {
+            let manifest_file = fs::File::open(Path::new(&manifest_path)).map_err(|err| {
+                BonsolCliError::ZkManifestError(ZkManifestError::FailedToOpen {
+                    manifest_path: manifest_path.clone(),
+                    err,
+                })
+            })?;
+            let manifest: ZkProgramManifest =
+                serde_json::from_reader(manifest_file).map_err(|err| {
+                    BonsolCliError::ZkManifestError(ZkManifestError::FailedDeserialization {
+                        manifest_path,
+                        err,
+                    })
+                })?;
+            let elf = fs::read(&manifest.binary_path).map_err(|err| {
+                BonsolCliError::ZkManifestError(ZkManifestError::FailedToLoadBinary {
+                    binary_path: manifest.binary_path.clone(),
+                    err,
+                })
+            })?;
+            let mut env = &mut ExecutorEnv::builder();
+            env = env
+                .segment_limit_po2(DEFAULT_SEGMENT_LIMIT_PO2 as u32)
+                .session_limit(max_cycles.or(DEFAULT_SESSION_LIMIT));
+
+            if input_file.is_some() {
+                let inputs = execute_get_inputs(input_file, None)?;
+                let inputs: Vec<&str> = inputs.iter().map(|i| i.data.as_str()).collect();
+                env = env.write(&inputs.as_slice())?;
+            }
+            estimate::estimate(elf.as_slice(), env.build()?)
+        }
         ParsedCommand::Execute {
             execution_request_file,
             program_id,

cli/src/tests.rs

@@ -0,0 +1,32 @@
+use std::path::Path;
+
+use assert_cmd::Command;
+
+mod estimate;
+
+pub(crate) fn bonsol_cmd() -> Command {
+    let mut cmd = Command::cargo_bin("bonsol").unwrap();
+    // the test directory must be the project root
+    cmd.current_dir(Path::new(env!("CARGO_MANIFEST_DIR")).parent().unwrap());
+    cmd
+}
+
+pub(crate) fn bonsol_build() -> Command {
+    let mut cmd = bonsol_cmd();
+    let keypair = cmd
+        .get_current_dir()
+        .unwrap()
+        .join("cli")
+        .join("src")
+        .join("tests")
+        .join("test_data")
+        .join("test_id.json");
+    cmd.args(&[
+        "--keypair",
+        keypair.to_str().unwrap(),
+        "--rpc-url",
+        "http://localhost:8899",
+    ])
+    .arg("build");
+    cmd
+}