add command line options, add test case

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"

cli/src/command.rs

@@ -251,6 +251,16 @@ pub enum Command {
             }
         )]
         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)]
@@ -315,6 +325,8 @@ pub enum ParsedCommand {
     },
     Estimate {
         elf: PathBuf,
+        segment_limit_po2: Option<usize>,
+        max_cycles: Option<u64>,
     },
     Execute {
         execution_request_file: Option<String>,
@@ -371,7 +383,7 @@ impl TryFrom<Command> for ParsedCommand {
                 ),
             }),
             Command::Build { zk_program_path } => Ok(ParsedCommand::Build { zk_program_path }),
-            Command::Estimate { elf } => Ok(ParsedCommand::Estimate { elf: PathBuf::from(elf) }),
+            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

@@ -6,32 +6,131 @@
 //! specific functions as sov-cycle-tracer does and the
 //! defaults made into command line args.
 
-use std::{fs, path::PathBuf};
-
 use anyhow::Result;
+use indicatif::{ProgressIterator, ProgressStyle};
+use risc0_binfmt::{MemoryImage, Program};
 use risc0_circuit_rv32im::prove::emu::{
-    exec::{execute_elf, DEFAULT_SEGMENT_LIMIT_PO2},
-    testutil::{NullSyscall, DEFAULT_SESSION_LIMIT},
+    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(())
+}
 
-pub fn estimate(elf: PathBuf) -> Result<()> {
-    // TODO: We probably want to be able to let the user decide some of this
-    let cycles: usize = execute_elf(
-        &fs::read(elf)?,
-        DEFAULT_SEGMENT_LIMIT_PO2,
-        DEFAULT_SESSION_LIMIT,
-        &NullSyscall::default(),
+/// 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)
-    })?;
+    })
+}
 
-    println!("number of cycles: {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)?;
+        MemoryImage::new(&program, PAGE_SIZE as u32)
+    }
+}
 
-    Ok(())
+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;
 
@@ -60,7 +61,11 @@ async fn main() -> anyhow::Result<()> {
             }
             deploy::deploy(rpc, keypair, deploy_args).await
         }
-        ParsedCommand::Estimate { elf } => estimate::estimate(elf),
+        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,