feat: Add estimate subcommand

cli/Cargo.toml

@@ -18,6 +18,7 @@ atty = "0.2.14"
 bincode = "1.3.3"
 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 +35,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"

cli/src/command.rs

@@ -1,3 +1,5 @@
+use std::path::PathBuf;
+
 use clap::{command, ArgGroup, Args, Parser, Subcommand, ValueEnum};
 
 #[derive(Parser, Debug)]
@@ -227,6 +229,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 +238,34 @@ pub enum Command {
         )]
         zk_program_path: String,
     },
+    #[command(about = "Estimate the execution cost of a ZK RISC0 program")]
+    Estimate {
+        #[arg(
+            help = "Specify the path to the RISC0 ELF",
+            long,
+            value_parser = |s: &str| {
+                if !PathBuf::from(s).exists() {
+                    anyhow::bail!("elf file path does not exist: '{s}'")
+                }
+                Ok(s.to_string())
+            }
+        )]
+        elf: String,
+
+        #[arg(
+            help = "Define the maximum number of cycles a single segment can take as a power of two, must be between 13 and 24 [default: 20usize]",
+            short = 's',
+            long
+        )]
+        segment_limit_po2: Option<usize>,
+
+        #[arg(
+            help = "Set the maximum number of cycles [default: 16777216u64]",
+            short = 'm',
+            long
+        )]
+        max_cycles: Option<u64>,
+    },
     Execute {
         #[arg(short = 'f', long)]
         execution_request_file: Option<String>,
@@ -296,6 +327,11 @@ pub enum ParsedCommand {
     Build {
         zk_program_path: String,
     },
+    Estimate {
+        elf: PathBuf,
+        segment_limit_po2: Option<usize>,
+        max_cycles: Option<u64>,
+    },
     Execute {
         execution_request_file: Option<String>,
 
@@ -351,6 +387,15 @@ impl TryFrom<Command> for ParsedCommand {
                 ),
             }),
             Command::Build { zk_program_path } => Ok(ParsedCommand::Build { zk_program_path }),
+            Command::Estimate {
+                elf,
+                segment_limit_po2,
+                max_cycles,
+            } => Ok(ParsedCommand::Estimate {
+                elf: PathBuf::from(elf),
+                segment_limit_po2,
+                max_cycles,
+            }),
             Command::Execute {
                 execution_request_file,
                 program_id,

cli/src/estimate.rs

@@ -0,0 +1,143 @@
+//! 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 indicatif::ProgressIterator;
+use risc0_binfmt::{MemoryImage, Program};
+use risc0_circuit_rv32im::prove::emu::{
+    exec::{execute, DEFAULT_SEGMENT_LIMIT_PO2},
+    testutil::DEFAULT_SESSION_LIMIT,
+};
+use risc0_zkvm::GUEST_MAX_MEM;
+use risc0_zkvm_platform::PAGE_SIZE;
+
+use self::emu_syscall::BasicSyscall;
+
+pub fn estimate<E: MkImage>(
+    elf: E,
+    segment_limit_po2: Option<usize>,
+    max_cycles: Option<u64>,
+) -> Result<()> {
+    let cycles: usize = get_cycle_count(elf, segment_limit_po2, max_cycles)?;
+    println!("number of cycles: {cycles}");
+
+    Ok(())
+}
+
+/// Get the total number of cycles by stepping through the ELF using emulation
+/// tools from the risc0_circuit_rv32im module.
+pub fn get_cycle_count<E: MkImage>(
+    elf: E,
+    segment_limit_po2: Option<usize>,
+    max_cycles: Option<u64>,
+) -> Result<usize> {
+    execute(
+        elf.mk_image()?,
+        segment_limit_po2.unwrap_or(DEFAULT_SEGMENT_LIMIT_PO2),
+        max_cycles.or(DEFAULT_SESSION_LIMIT),
+        &BasicSyscall::default(),
+        None,
+    )?
+    .segments
+    .iter()
+    .progress()
+    .try_fold(0, |acc, s| -> Result<usize> {
+        let trace = s.preflight()?;
+        let segment_cycles = trace.pre.cycles.len() + trace.body.cycles.len();
+        Ok(acc + segment_cycles)
+    })
+}
+
+/// 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)?;
+        dbg!(&program.entry);
+        // dbg!(&program.image);
+        program
+            .image
+            .keys()
+            .zip(program.image.values())
+            .for_each(|(k, v)| {
+                if v == &31 {
+                    eprintln!("found invalid guest address for key: {}", k);
+                }
+            });
+        MemoryImage::new(&program, PAGE_SIZE as u32)
+    }
+}
+
+pub mod emu_syscall {
+    //! The following is copied from risc0 emu test utils, likely this is okay for our use case since we only want the cycle count.
+    //! https://github.com/anagrambuild/risc0/blob/eb331d7ee30bc9ccf944bb1ea4835e60e21c25a2/risc0/circuit/rv32im/src/prove/emu/exec/tests.rs#L41
+
+    use std::cell::RefCell;
+
+    use anyhow::Result;
+    use risc0_circuit_rv32im::prove::emu::{
+        addr::ByteAddr,
+        exec::{Syscall, SyscallContext},
+    };
+    use risc0_zkvm_platform::syscall::reg_abi::{REG_A4, REG_A5};
+
+    #[derive(Default, Clone)]
+    pub struct BasicSyscallState {
+        syscall: String,
+        from_guest: Vec<u8>,
+        into_guest: Vec<u8>,
+    }
+
+    #[derive(Default)]
+    pub struct BasicSyscall {
+        state: RefCell<BasicSyscallState>,
+    }
+
+    impl Syscall for BasicSyscall {
+        fn syscall(
+            &self,
+            syscall: &str,
+            ctx: &mut dyn SyscallContext,
+            guest_buf: &mut [u32],
+        ) -> Result<(u32, u32)> {
+            self.state.borrow_mut().syscall = syscall.to_string();
+            let buf_ptr = ByteAddr(ctx.peek_register(REG_A4)?);
+            let buf_len = ctx.peek_register(REG_A5)?;
+            self.state.borrow_mut().from_guest = ctx.peek_region(buf_ptr, buf_len)?;
+            let guest_buf_bytes: &mut [u8] = bytemuck::cast_slice_mut(guest_buf);
+            let into_guest = &self.state.borrow().into_guest;
+            guest_buf_bytes[..into_guest.len()].clone_from_slice(into_guest);
+            Ok((0, 0))
+        }
+    }
+}
+
+#[cfg(test)]
+mod estimate_tests {
+    use anyhow::Result;
+    use risc0_binfmt::MemoryImage;
+    use risc0_circuit_rv32im::prove::emu::testutil::basic as basic_test_program;
+    use risc0_zkvm::PAGE_SIZE;
+
+    use super::MkImage;
+    use crate::estimate;
+
+    impl MkImage for MemoryImage {
+        fn mk_image(self) -> Result<MemoryImage> {
+            Ok(self)
+        }
+    }
+
+    #[test]
+    fn test_estimate() {
+        let program = basic_test_program();
+        let image = MemoryImage::new(&program, PAGE_SIZE as u32)
+            .expect("failed to create image from basic program");
+        let res = estimate::get_cycle_count(image, None, None);
+
+        assert_eq!(res.ok(), Some(15790));
+    }
+}

cli/src/main.rs

@@ -1,3 +1,4 @@
+use std::fs;
 use std::io::{self, Read};
 use std::path::Path;
 
@@ -13,6 +14,7 @@ use crate::error::BonsolCliError;
 
 mod build;
 mod deploy;
+mod estimate;
 mod execute;
 mod init;
 mod prove;
@@ -59,6 +61,11 @@ async fn main() -> anyhow::Result<()> {
             }
             deploy::deploy(rpc, keypair, deploy_args).await
         }
+        ParsedCommand::Estimate {
+            elf,
+            segment_limit_po2,
+            max_cycles,
+        } => estimate::estimate(fs::read(elf)?.as_slice(), segment_limit_po2, max_cycles),
         ParsedCommand::Execute {
             execution_request_file,
             program_id,