feat: cycle limit (#1027)

Author: jtguibas

core/src/runtime/context.rs

@@ -7,16 +7,23 @@ use super::{hookify, BoxedHook, HookEnv, HookRegistry, SubproofVerifier};
 #[derive(Clone, Default)]
 pub struct SP1Context<'a> {
     /// The registry of hooks invokable from inside SP1.
-    /// `None` denotes the default list of hooks.
+    ///
+    /// Note: `None` denotes the default list of hooks.
     pub hook_registry: Option<HookRegistry<'a>>,
+
+    /// The verifier for verifying subproofs.
     pub subproof_verifier: Option<Arc<dyn SubproofVerifier + 'a>>,
+
+    /// The maximum number of cpu cycles to use for execution.
+    pub max_cycles: Option<u64>,
 }
 
 #[derive(Clone, Default)]
 pub struct SP1ContextBuilder<'a> {
     no_default_hooks: bool,
     hook_registry_entries: Vec<(u32, BoxedHook<'a>)>,
     subproof_verifier: Option<Arc<dyn SubproofVerifier + 'a>>,
+    max_cycles: Option<u64>,
 }
 
 impl<'a> SP1Context<'a> {
@@ -52,9 +59,11 @@ impl<'a> SP1ContextBuilder<'a> {
                 HookRegistry { table }
             });
         let subproof_verifier = take(&mut self.subproof_verifier);
+        let cycle_limit = take(&mut self.max_cycles);
         SP1Context {
             hook_registry,
             subproof_verifier,
+            max_cycles: cycle_limit,
         }
     }
 
@@ -91,6 +100,12 @@ impl<'a> SP1ContextBuilder<'a> {
         self.subproof_verifier = Some(subproof_verifier);
         self
     }
+
+    /// Set the maximum number of cpu cycles to use for execution.
+    pub fn max_cycles(&mut self, max_cycles: u64) -> &mut Self {
+        self.max_cycles = Some(max_cycles);
+        self
+    }
 }
 
 #[cfg(test)]
@@ -104,9 +119,11 @@ mod tests {
         let SP1Context {
             hook_registry,
             subproof_verifier,
+            max_cycles: cycle_limit,
         } = SP1Context::builder().build();
         assert!(hook_registry.is_none());
         assert!(subproof_verifier.is_none());
+        assert!(cycle_limit.is_none());
     }
 
     #[test]

core/src/runtime/mod.rs

@@ -103,6 +103,9 @@ pub struct Runtime<'a> {
 
     /// Registry of hooks, to be invoked by writing to certain file descriptors.
     pub hook_registry: HookRegistry<'a>,
+
+    /// The maximum number of cpu cycles to use for execution.
+    pub max_cycles: Option<u64>,
 }
 
 #[derive(Error, Debug)]
@@ -115,6 +118,8 @@ pub enum ExecutionError {
     UnsupportedSyscall(u32),
     #[error("breakpoint encountered")]
     Breakpoint(),
+    #[error("exceeded cycle limit of {0}")]
+    ExceededCycleLimit(u64),
     #[error("got unimplemented as opcode")]
     Unimplemented(),
 }
@@ -176,6 +181,7 @@ impl<'a> Runtime<'a> {
             print_report: false,
             subproof_verifier,
             hook_registry,
+            max_cycles: context.max_cycles,
         }
     }
 
@@ -992,6 +998,13 @@ impl<'a> Runtime<'a> {
             self.state.channel = 0;
         }
 
+        // If the cycle limit is exceeded, return an error.
+        if let Some(max_cycles) = self.max_cycles {
+            if self.state.global_clk >= max_cycles {
+                return Err(ExecutionError::ExceededCycleLimit(max_cycles));
+            }
+        }
+
         Ok(self.state.pc.wrapping_sub(self.program.pc_base)
             >= (self.program.instructions.len() * 4) as u32)
     }
@@ -1105,7 +1118,9 @@ impl<'a> Runtime<'a> {
 
         // Ensure that all proofs and input bytes were read, otherwise warn the user.
         if self.state.proof_stream_ptr != self.state.proof_stream.len() {
-            panic!("Not all proofs were read. Proving will fail during recursion. Did you pass too many proofs in or forget to call verify_sp1_proof?");
+            panic!(
+                "Not all proofs were read. Proving will fail during recursion. Did you pass too many proofs in or forget to call verify_sp1_proof?"
+            );
         }
         if self.state.input_stream_ptr != self.state.input_stream.len() {
             log::warn!("Not all input bytes were read.");

sdk/src/action.rs

@@ -63,6 +63,14 @@ impl<'a> Execute<'a> {
         self.context_builder.without_default_hooks();
         self
     }
+
+    /// Set the maximum number of cpu cycles to use for execution.
+    ///
+    /// If the cycle limit is exceeded, execution will return [sp1_core::runtime::ExecutionError::ExceededCycleLimit].
+    pub fn max_cycles(mut self, max_cycles: u64) -> Self {
+        self.context_builder.max_cycles(max_cycles);
+        self
+    }
 }
 
 /// Builder to prepare and configure proving execution of a program on an input.
@@ -175,4 +183,12 @@ impl<'a> Prove<'a> {
         self.opts.reconstruct_commitments = value;
         self
     }
+
+    /// Set the maximum number of cpu cycles to use for execution.
+    ///
+    /// If the cycle limit is exceeded, execution will return [sp1_core::runtime::ExecutionError::ExceededCycleLimit].
+    pub fn cycle_limit(mut self, cycle_limit: u64) -> Self {
+        self.context_builder.max_cycles(cycle_limit);
+        self
+    }
 }

sdk/src/lib.rs

@@ -335,6 +335,17 @@ mod tests {
         client.execute(elf, stdin).run().unwrap();
     }
 
+    #[should_panic]
+    #[test]
+    fn test_cycle_limit_fail() {
+        utils::setup_logger();
+        let client = ProverClient::local();
+        let elf = include_bytes!("../../tests/panic/elf/riscv32im-succinct-zkvm-elf");
+        let mut stdin = SP1Stdin::new();
+        stdin.write(&10usize);
+        client.execute(elf, stdin).max_cycles(1).run().unwrap();
+    }
+
     #[test]
     fn test_e2e_prove_plonk() {
         utils::setup_logger();

sdk/src/network/prover.rs

@@ -162,6 +162,7 @@ fn warn_if_not_default(opts: &SP1ProverOpts, context: &SP1Context) {
     let SP1Context {
         hook_registry,
         subproof_verifier,
+        ..
     } = context;
     if hook_registry.is_some() {
         tracing::warn!(