feat: support limit argument for dry-run

crates/cli/src/app_builder.rs

@@ -161,6 +161,19 @@ impl ArgBuilder<Option<usize>> for PaddingArg {
     }
 }
 
+struct InstructionLimitArg;
+impl ArgBuilder<Option<usize>> for InstructionLimitArg {
+    fn builder() -> Arg<'static> {
+        arg!(--limit [LIMIT] "Terminate if the number of instructions exceeds [LIMIT]")
+            .value_parser(value_parser!(usize))
+            .multiple_values(false)
+    }
+
+    fn parse(matches: &ArgMatches) -> Option<usize> {
+        matches.get_one("limit").cloned()
+    }
+}
+
 fn setup_command() -> Command<'static> {
     let command = Command::new("setup")
         .about("Setup a new zkWasm circuit for provided Wasm image")
@@ -204,6 +217,7 @@ fn dry_run_command() -> Command<'static> {
         .arg(ContextInputsArg::builder())
         .arg(ContextOutputArg::builder())
         .arg(OutputDirArg::builder())
+        .arg(InstructionLimitArg::builder())
 }
 
 fn prove_command() -> Command<'static> {
@@ -280,6 +294,7 @@ impl From<&ArgMatches> for DryRunArg {
         DryRunArg {
             wasm_image: WasmImageArg::parse(val).unwrap(),
             running_arg: val.into(),
+            instruction_limit: InstructionLimitArg::parse(val),
         }
     }
 }

crates/cli/src/command.rs

@@ -251,6 +251,7 @@ pub(crate) struct RunningArg {
 pub(crate) struct DryRunArg {
     pub(crate) wasm_image: PathBuf,
     pub(crate) running_arg: RunningArg,
+    pub(crate) instruction_limit: Option<usize>,
 }
 
 /// Execute the Wasm image and generate a proof.

crates/cli/src/config.rs

@@ -190,12 +190,13 @@ impl Config {
         output_dir: &Path,
         arg: ExecutionArg,
         context_output_filename: Option<String>,
+        instruction_limit: Option<usize>,
     ) -> Result<()> {
         let module = self.read_wasm_image(wasm_image)?;
 
         let env = env_builder.create_env(self.k, arg);
 
-        let mut monitor = StatisticMonitor::new(&self.phantom_functions, &env);
+        let mut monitor = StatisticMonitor::new(&self.phantom_functions, &env, instruction_limit);
 
         let result = {
             let loader = ZkWasmLoader::new(self.k, env)?;

crates/cli/src/main.rs

@@ -84,6 +84,7 @@ fn main() -> Result<()> {
                     context_inputs,
                 },
                 arg.running_arg.context_output,
+                arg.instruction_limit,
             )?;
         }
         Subcommands::Prove(arg) => {

crates/zkwasm/src/runtime/monitor/plugins/statistic.rs

@@ -7,20 +7,28 @@ use wasmi::isa::Keep;
 use wasmi::monitor::Monitor;
 use wasmi::runner::InstructionOutcome;
 use wasmi::FuncRef;
+use wasmi::Trap;
+use wasmi::TrapCode;
 
 use super::phantom::PhantomHelper;
 use crate::runtime::monitor::Observer;
 
 pub struct StatisticPlugin {
     phantom_helper: PhantomHelper,
     observer: Rc<RefCell<Observer>>,
+    instruction_limit: Option<usize>,
 }
 
 impl StatisticPlugin {
-    pub fn new(phantom_regex: &[String], wasm_input: FuncRef) -> Self {
+    pub fn new(
+        phantom_regex: &[String],
+        wasm_input: FuncRef,
+        instruction_limit: Option<usize>,
+    ) -> Self {
         Self {
             phantom_helper: PhantomHelper::new(phantom_regex, wasm_input),
             observer: Rc::new(RefCell::new(Observer::default())),
+            instruction_limit,
         }
     }
 
@@ -52,10 +60,16 @@ impl Monitor for StatisticPlugin {
         _function_context: &wasmi::runner::FunctionContext,
         _instruction: &wasmi::isa::Instruction,
         outcome: &wasmi::runner::InstructionOutcome,
-    ) {
+    ) -> Result<(), Trap> {
         self.observer.borrow_mut().counter +=
             !self.phantom_helper.is_in_phantom_function() as usize;
 
+        if let Some(instruction_limit) = self.instruction_limit {
+            if self.observer.borrow_mut().counter > instruction_limit {
+                return Err(Trap::Code(TrapCode::InstructionExceedsLimit));
+            }
+        }
+
         match outcome {
             InstructionOutcome::ExecuteCall(func_ref) => {
                 if let FuncInstanceInternal::Internal { index, .. } = func_ref.as_internal() {
@@ -89,5 +103,7 @@ impl Monitor for StatisticPlugin {
             }
             _ => {}
         }
+
+        Ok(())
     }
 }

crates/zkwasm/src/runtime/monitor/plugins/table/mod.rs

@@ -37,6 +37,7 @@ use wasmi::Error;
 use wasmi::FuncRef;
 use wasmi::RuntimeValue;
 use wasmi::Signature;
+use wasmi::Trap;
 use wasmi::DEFAULT_VALUE_STACK_LIMIT;
 
 use crate::circuits::compute_slice_capability;
@@ -566,7 +567,7 @@ impl Monitor for TablePlugin {
         function_context: &FunctionContext,
         instruction: &Instruction,
         outcome: &InstructionOutcome,
-    ) {
+    ) -> Result<(), Trap> {
         if !self.phantom_helper.is_in_phantom_function() {
             let current_event = self.unresolved_event.take();
 
@@ -662,6 +663,8 @@ impl Monitor for TablePlugin {
             }
             _ => {}
         }
+
+        Ok(())
     }
 
     fn invoke_call_host_post_hook(&mut self, return_value: Option<RuntimeValue>) {

crates/zkwasm/src/runtime/monitor/statistic_monitor.rs

@@ -11,6 +11,7 @@ use wasmi::ModuleImportResolver;
 use wasmi::ModuleRef;
 use wasmi::RuntimeValue;
 use wasmi::Signature;
+use wasmi::Trap;
 use wasmi::ValueType;
 
 use crate::runtime::host::host_env::HostEnv;
@@ -24,7 +25,7 @@ pub struct StatisticMonitor {
 }
 
 impl StatisticMonitor {
-    pub fn new(phantom_regex: &[String], env: &HostEnv) -> Self {
+    pub fn new(phantom_regex: &[String], env: &HostEnv, instruction_limit: Option<usize>) -> Self {
         let wasm_input = env
             .resolve_func(
                 "wasm_input",
@@ -33,7 +34,7 @@ impl StatisticMonitor {
             .expect("Failed to resolve wasm_input function, please make sure it is imported in the wasm image.");
 
         Self {
-            statistic_plugin: StatisticPlugin::new(phantom_regex, wasm_input),
+            statistic_plugin: StatisticPlugin::new(phantom_regex, wasm_input, instruction_limit),
         }
     }
 }
@@ -78,7 +79,7 @@ impl Monitor for StatisticMonitor {
         function_context: &FunctionContext,
         instruction: &Instruction,
         outcome: &InstructionOutcome,
-    ) {
+    ) -> Result<(), Trap> {
         self.statistic_plugin.invoke_instruction_post_hook(
             fid,
             iid,
@@ -88,7 +89,9 @@ impl Monitor for StatisticMonitor {
             function_context,
             instruction,
             outcome,
-        );
+        )?;
+
+        Ok(())
     }
 
     fn invoke_call_host_post_hook(&mut self, return_value: Option<RuntimeValue>) {

crates/zkwasm/src/runtime/monitor/table_monitor.rs

@@ -14,6 +14,7 @@ use wasmi::ModuleImportResolver;
 use wasmi::ModuleRef;
 use wasmi::RuntimeValue;
 use wasmi::Signature;
+use wasmi::Trap;
 use wasmi::ValueType;
 
 use crate::runtime::host::host_env::HostEnv;
@@ -45,7 +46,7 @@ impl TableMonitor {
                 wasm_input.clone(),
                 backend,
             ),
-            statistic_plugin: StatisticPlugin::new(phantom_regex, wasm_input),
+            statistic_plugin: StatisticPlugin::new(phantom_regex, wasm_input, None),
         }
     }
 
@@ -103,7 +104,7 @@ impl Monitor for TableMonitor {
         function_context: &FunctionContext,
         instruction: &Instruction,
         outcome: &InstructionOutcome,
-    ) {
+    ) -> Result<(), Trap> {
         self.table_plugin.invoke_instruction_post_hook(
             fid,
             iid,
@@ -113,7 +114,7 @@ impl Monitor for TableMonitor {
             function_context,
             instruction,
             outcome,
-        );
+        )?;
         self.statistic_plugin.invoke_instruction_post_hook(
             fid,
             iid,
@@ -123,7 +124,9 @@ impl Monitor for TableMonitor {
             function_context,
             instruction,
             outcome,
-        );
+        )?;
+
+        Ok(())
     }
 
     fn invoke_call_host_post_hook(&mut self, return_value: Option<RuntimeValue>) {