feat: E2 supports suspension & VmExecState supports clone

crates/sdk/tests/integration_test.rs

@@ -8,7 +8,10 @@ use eyre::Result;
 use openvm_build::GuestOptions;
 use openvm_circuit::{
     self,
-    arch::{instructions::exe::VmExe, ContinuationVmProof, ExecutionError, VirtualMachineError},
+    arch::{
+        instructions::exe::VmExe, ContinuationVmProof, ExecutionError, VirtualMachine,
+        VirtualMachineError, VmExecState,
+    },
     utils::test_system_config,
 };
 use openvm_continuations::verifier::{
@@ -19,15 +22,16 @@ use openvm_native_circuit::{execute_program_with_config, NativeConfig, NativeCpu
 use openvm_native_compiler::{conversion::CompilerOptions, prelude::*};
 use openvm_sdk::{
     codec::{Decode, Encode},
-    config::{AggregationConfig, AppConfig, SdkSystemConfig, SdkVmConfig},
+    config::{AggregationConfig, AppConfig, SdkSystemConfig, SdkVmBuilder, SdkVmConfig},
     prover::verify_app_proof,
-    Sdk, StdIn,
+    DefaultStarkEngine, Sdk, StdIn,
 };
 use openvm_stark_sdk::{
     config::{
         baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine},
         setup_tracing, FriParameters,
     },
+    engine::StarkFriEngine,
     openvm_stark_backend::p3_field::FieldAlgebra,
     p3_baby_bear::BabyBear,
 };
@@ -265,6 +269,60 @@ fn test_public_values_and_leaf_verification() -> eyre::Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_metered_execution_suspension() -> eyre::Result<()> {
+    setup_tracing();
+    let app_log_blowup = 1;
+    let app_config = small_test_app_config(app_log_blowup);
+    let exe = app_exe_for_test();
+
+    let engine = DefaultStarkEngine::new(app_config.app_fri_params.fri_params);
+    let (vm, _) = VirtualMachine::new_with_keygen(engine, SdkVmBuilder, app_config.app_vm_config)?;
+    let executor_idx_to_air_idx = vm.executor_idx_to_air_idx();
+    let interpreter = vm
+        .executor()
+        .metered_instance(&exe, &executor_idx_to_air_idx)?;
+    let metered_ctx = vm.build_metered_ctx(&exe).with_suspend_on_segment(true);
+    let vm_state = interpreter.create_initial_vm_state(vec![]);
+    let mut vm_exec_state = VmExecState::new(vm_state, metered_ctx);
+    vm_exec_state = interpreter.execute_metered_until_suspension(vm_exec_state)?;
+    assert!(
+        vm_exec_state.exit_code.is_ok(),
+        "Execution exits with an error: {:?}",
+        vm_exec_state.exit_code
+    );
+
+    // Benchmark VmExecState clone.
+    unsafe {
+        vm_exec_state.memory.write(2, 100, [16u8; 1]);
+    }
+    let start = std::time::Instant::now();
+    let mut state_vec = vec![];
+    let n = 10;
+    for _ in 0..n {
+        state_vec.push(vm_exec_state.try_clone()?);
+    }
+    let dur = start.elapsed();
+    println!("Vm State clone average time: {}us", dur.as_micros() / n);
+
+    assert!(
+        vm_exec_state.exit_code?.is_none(),
+        "Expect more than 1 segment but the execution terminates."
+    );
+    // Avoid compiler from optimizing out the VmExecState clone.
+    {
+        let values: Vec<_> = unsafe {
+            state_vec
+                .iter()
+                .map(|state| state.vm_state.memory.read::<u8, 1>(2, 100))
+                .collect()
+        };
+        println!("Values at address space 2, ptr: 100: {:?}", values);
+    }
+
+    Ok(())
+}
+
 #[cfg(feature = "evm-verify")]
 #[test]
 #[ignore = "slow"]

crates/vm/src/arch/execution_mode/metered/ctx.rs

@@ -1,5 +1,6 @@
 use std::num::NonZero;
 
+use getset::{Getters, Setters, WithSetters};
 use itertools::Itertools;
 use openvm_instructions::riscv::{RV32_IMM_AS, RV32_REGISTER_AS};
 
@@ -17,12 +18,14 @@ use crate::{
 
 pub const DEFAULT_PAGE_BITS: usize = 6;
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, Getters, Setters, WithSetters)]
 pub struct MeteredCtx<const PAGE_BITS: usize = DEFAULT_PAGE_BITS> {
     pub trace_heights: Vec<u32>,
     pub is_trace_height_constant: Vec<bool>,
     pub memory_ctx: MemoryCtx<PAGE_BITS>,
     pub segmentation_ctx: SegmentationCtx,
+    #[getset(get = "pub", set = "pub", set_with = "pub")]
+    suspend_on_segment: bool,
 }
 
 impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
@@ -75,6 +78,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
             is_trace_height_constant,
             memory_ctx,
             segmentation_ctx,
+            suspend_on_segment: false,
         };
         if !config.continuation_enabled {
             // force single segment
@@ -121,7 +125,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
     }
 
     #[inline(always)]
-    pub fn check_and_segment(&mut self, instret: u64, segment_check_insns: u64) {
+    pub fn check_and_segment(&mut self, instret: u64, segment_check_insns: u64) -> bool {
         let threshold = self
             .segmentation_ctx
             .instret_last_segment_check
@@ -131,7 +135,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
             "overflow in segment check threshold calculation"
         );
         if instret < threshold {
-            return;
+            return false;
         }
 
         self.memory_ctx
@@ -145,6 +149,7 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
         if did_segment {
             self.reset_segment();
         }
+        did_segment
     }
 
     #[allow(dead_code)]
@@ -200,12 +205,13 @@ impl<const PAGE_BITS: usize> ExecutionCtxTrait for MeteredCtx<PAGE_BITS> {
         segment_check_insns: u64,
         exec_state: &mut VmExecState<F, GuestMemory, Self>,
     ) -> bool {
-        // E2 always runs until termination. Here we use the function as a hook called every
-        // instruction.
+        // If `segment_suspend` is set, suspend when a segment is determined (but the VM state might
+        // be after the segment boundary because the segment happens in the previous checkpoint).
+        // Otherwise, execute until termination.
         exec_state
             .ctx
-            .check_and_segment(instret, segment_check_insns);
-        false
+            .check_and_segment(instret, segment_check_insns)
+            && exec_state.ctx.suspend_on_segment
     }
 
     #[inline(always)]

crates/vm/src/arch/interpreter.rs

@@ -201,6 +201,15 @@ where
         })
     }
 
+    pub fn create_initial_vm_state(&self, inputs: impl Into<Streams<F>>) -> VmState<F> {
+        VmState::initial(
+            &self.system_config,
+            &self.init_memory,
+            self.pc_start,
+            inputs,
+        )
+    }
+
     /// # Safety
     /// - This function assumes that the `pc` is within program bounds - this should be the case if
     ///   the pc is checked to be in bounds before jumping to it.
@@ -381,12 +390,7 @@ where
         inputs: impl Into<Streams<F>>,
         ctx: MeteredCtx,
     ) -> Result<(Vec<Segment>, VmState<F, GuestMemory>), ExecutionError> {
-        let vm_state = VmState::initial(
-            &self.system_config,
-            &self.init_memory,
-            self.pc_start,
-            inputs,
-        );
+        let vm_state = self.create_initial_vm_state(inputs);
         self.execute_metered_from_state(vm_state, ctx)
     }
 
@@ -405,6 +409,44 @@ where
     ) -> Result<(Vec<Segment>, VmState<F, GuestMemory>), ExecutionError> {
         let mut exec_state = VmExecState::new(from_state, ctx);
 
+        loop {
+            exec_state = self.execute_metered_until_suspension(exec_state)?;
+            // The execution has terminated.
+            if exec_state.exit_code.is_ok() && exec_state.exit_code.as_ref().unwrap().is_some() {
+                break;
+            }
+            if exec_state.exit_code.is_err() {
+                return Err(exec_state.exit_code.unwrap_err());
+            }
+        }
+        check_termination(exec_state.exit_code)?;
+        let VmExecState { vm_state, ctx, .. } = exec_state;
+        Ok((ctx.into_segments(), vm_state))
+    }
+    /// Executes a metered virtual machine operation starting from a given execution state until
+    /// suspension.
+    ///
+    /// This function resumes and continues execution of a guest virtual machine until either it:
+    /// - Hits a suspension trigger (e.g. out of gas or a specific halt condition). ATTENTION: when
+    ///   a suspension is triggered, the VM state is not at the boundary of the last segment.
+    ///   Instead, the VM state is slightly after the segment boundary.
+    /// - Completes its run based on the instructions or context provided.
+    ///
+    /// # Parameters
+    /// - `self`: The reference to the current executor or VM context.
+    /// - `exec_state`: A mutable `VmExecState<F, GuestMemory, MeteredCtx>` which represents the
+    ///   execution state of the virtual machine, including its program counter (`pc`), instruction
+    ///   retirement (`instret`), and execution context (`MeteredCtx`).
+    ///
+    /// # Returns
+    /// - `Ok(VmExecState<F, GuestMemory, MeteredCtx>)`: The execution state after suspension or
+    ///   normal completion.
+    /// - `Err(ExecutionError)`: If there is an error during execution, such as an invalid state or
+    ///   run-time error.
+    pub fn execute_metered_until_suspension(
+        &self,
+        mut exec_state: VmExecState<F, GuestMemory, MeteredCtx>,
+    ) -> Result<VmExecState<F, GuestMemory, MeteredCtx>, ExecutionError> {
         let instret = exec_state.instret();
         let pc = exec_state.pc();
         let segmentation_check_insns = exec_state.ctx.segmentation_ctx.segment_check_insns;
@@ -418,9 +460,7 @@ where
             exec_state,
             MeteredCtx
         );
-        check_termination(exec_state.exit_code)?;
-        let VmExecState { vm_state, ctx, .. } = exec_state;
-        Ok((ctx.into_segments(), vm_state))
+        Ok(exec_state)
     }
 }
 
@@ -437,12 +477,7 @@ where
         inputs: impl Into<Streams<F>>,
         ctx: MeteredCostCtx,
     ) -> Result<(u64, VmState<F, GuestMemory>), ExecutionError> {
-        let vm_state = VmState::initial(
-            &self.system_config,
-            &self.init_memory,
-            self.pc_start,
-            inputs,
-        );
+        let vm_state = self.create_initial_vm_state(inputs);
         self.execute_metered_cost_from_state(vm_state, ctx)
     }
 

crates/vm/src/arch/state.rs

@@ -3,6 +3,7 @@ use std::{
     ops::{Deref, DerefMut},
 };
 
+use eyre::eyre;
 use getset::{CopyGetters, MutGetters};
 use openvm_instructions::exe::SparseMemoryImage;
 use rand::{rngs::StdRng, SeedableRng};
@@ -17,7 +18,7 @@ use crate::{
 };
 
 /// Represents the core state of a VM.
-#[derive(derive_new::new, CopyGetters, MutGetters)]
+#[derive(derive_new::new, CopyGetters, MutGetters, Clone)]
 pub struct VmState<F, MEM = GuestMemory> {
     #[getset(get_copy = "pub", get_mut = "pub")]
     instret: u64,
@@ -136,6 +137,25 @@ impl<F, MEM, CTX> VmExecState<F, MEM, CTX> {
             exit_code: Ok(None),
         }
     }
+
+    /// Try to clone VmExecState. Return an error if `exit_code` is an error because `ExecutionEror`
+    /// cannot be cloned.
+    pub fn try_clone(&self) -> eyre::Result<Self>
+    where
+        VmState<F, MEM>: Clone,
+        CTX: Clone,
+    {
+        if self.exit_code.is_err() {
+            return Err(eyre!(
+                "failed to clone VmExecState because exit_code is an error"
+            ));
+        }
+        Ok(Self {
+            vm_state: self.vm_state.clone(),
+            exit_code: Ok(*self.exit_code.as_ref().unwrap()),
+            ctx: self.ctx.clone(),
+        })
+    }
 }
 
 impl<F, MEM, CTX> Deref for VmExecState<F, MEM, CTX> {