openvm_circuit/arch/
state.rs

1use std::{
2    fmt::Debug,
3    ops::{Deref, DerefMut},
4};
5
6use eyre::eyre;
7use getset::{CopyGetters, MutGetters};
8use openvm_instructions::exe::SparseMemoryImage;
9use rand::{rngs::StdRng, SeedableRng};
10use tracing::instrument;
11
12use super::{create_memory_image, ExecutionError, Streams};
13#[cfg(feature = "metrics")]
14use crate::metrics::VmMetrics;
15use crate::{
16    arch::{execution_mode::ExecutionCtxTrait, SystemConfig, VmStateMut},
17    system::memory::online::GuestMemory,
18};
19
20/// Represents the core state of a VM.
21#[repr(C)]
22#[derive(derive_new::new, CopyGetters, MutGetters, Clone)]
23pub struct VmState<F, MEM = GuestMemory> {
24    #[getset(get_copy = "pub", get_mut = "pub")]
25    pc: u32,
26    pub memory: MEM,
27    pub streams: Streams<F>,
28    pub rng: StdRng,
29    #[cfg(feature = "metrics")]
30    pub metrics: VmMetrics,
31}
32
33pub(super) const DEFAULT_RNG_SEED: u64 = 0;
34
35impl<F, MEM> VmState<F, MEM> {
36    #[inline(always)]
37    pub fn set_pc(&mut self, pc: u32) {
38        self.pc = pc;
39    }
40}
41
42impl<F: Clone, MEM> VmState<F, MEM> {
43    pub fn new_with_defaults(
44        pc: u32,
45        memory: MEM,
46        streams: impl Into<Streams<F>>,
47        seed: u64,
48    ) -> Self {
49        Self {
50            pc,
51            memory,
52            streams: streams.into(),
53            rng: StdRng::seed_from_u64(seed),
54            #[cfg(feature = "metrics")]
55            metrics: VmMetrics::default(),
56        }
57    }
58
59    #[inline(always)]
60    pub fn into_mut<'a, RA>(&'a mut self, ctx: &'a mut RA) -> VmStateMut<'a, F, MEM, RA> {
61        VmStateMut {
62            pc: &mut self.pc,
63            memory: &mut self.memory,
64            streams: &mut self.streams,
65            rng: &mut self.rng,
66            ctx,
67            #[cfg(feature = "metrics")]
68            metrics: &mut self.metrics,
69        }
70    }
71}
72
73impl<F: Clone> VmState<F, GuestMemory> {
74    #[instrument(name = "VmState::initial", level = "debug", skip_all)]
75    pub fn initial(
76        system_config: &SystemConfig,
77        init_memory: &SparseMemoryImage,
78        pc_start: u32,
79        inputs: impl Into<Streams<F>>,
80    ) -> Self {
81        let memory = create_memory_image(&system_config.memory_config, init_memory);
82        VmState::new_with_defaults(pc_start, memory, inputs.into(), DEFAULT_RNG_SEED)
83    }
84
85    pub fn reset(
86        &mut self,
87        init_memory: &SparseMemoryImage,
88        pc_start: u32,
89        streams: impl Into<Streams<F>>,
90    ) {
91        self.pc = pc_start;
92        self.memory.memory.fill_zero();
93        self.memory.memory.set_from_sparse(init_memory);
94        self.streams = streams.into();
95        self.rng = StdRng::seed_from_u64(DEFAULT_RNG_SEED);
96    }
97}
98
99/// Represents the full execution state of a VM during execution.
100/// The global state is generic in guest memory `MEM` and additional context `CTX`.
101/// The host state is execution context specific.
102// @dev: Do not confuse with `ExecutionState` struct.
103#[repr(C)]
104pub struct VmExecState<F, MEM, CTX> {
105    /// Core VM state
106    pub vm_state: VmState<F, MEM>,
107    pub ctx: CTX,
108    /// Execution-specific fields
109    pub exit_code: Result<Option<u32>, ExecutionError>,
110}
111
112impl<F, CTX: ExecutionCtxTrait> VmExecState<F, GuestMemory, CTX> {
113    #[inline(always)]
114    pub fn should_suspend(&mut self) -> bool {
115        CTX::should_suspend(self)
116    }
117}
118
119impl<F, MEM, CTX> VmExecState<F, MEM, CTX> {
120    pub fn new(vm_state: VmState<F, MEM>, ctx: CTX) -> Self {
121        Self {
122            vm_state,
123            ctx,
124            exit_code: Ok(None),
125        }
126    }
127
128    /// Try to clone VmExecState. Return an error if `exit_code` is an error because `ExecutionEror`
129    /// cannot be cloned.
130    pub fn try_clone(&self) -> eyre::Result<Self>
131    where
132        VmState<F, MEM>: Clone,
133        CTX: Clone,
134    {
135        if self.exit_code.is_err() {
136            return Err(eyre!(
137                "failed to clone VmExecState because exit_code is an error"
138            ));
139        }
140        Ok(Self {
141            vm_state: self.vm_state.clone(),
142            exit_code: Ok(*self.exit_code.as_ref().unwrap()),
143            ctx: self.ctx.clone(),
144        })
145    }
146}
147
148impl<F, MEM, CTX> Deref for VmExecState<F, MEM, CTX> {
149    type Target = VmState<F, MEM>;
150
151    fn deref(&self) -> &Self::Target {
152        &self.vm_state
153    }
154}
155
156impl<F, MEM, CTX> DerefMut for VmExecState<F, MEM, CTX> {
157    fn deref_mut(&mut self) -> &mut Self::Target {
158        &mut self.vm_state
159    }
160}
161
162impl<F, CTX> VmExecState<F, GuestMemory, CTX>
163where
164    CTX: ExecutionCtxTrait,
165{
166    /// Runtime read operation for a block of memory
167    #[inline(always)]
168    pub fn vm_read<T: Copy + Debug, const BLOCK_SIZE: usize>(
169        &mut self,
170        addr_space: u32,
171        ptr: u32,
172    ) -> [T; BLOCK_SIZE] {
173        self.ctx
174            .on_memory_operation(addr_space, ptr, BLOCK_SIZE as u32);
175        self.host_read(addr_space, ptr)
176    }
177
178    /// Runtime write operation for a block of memory
179    #[inline(always)]
180    pub fn vm_write<T: Copy + Debug, const BLOCK_SIZE: usize>(
181        &mut self,
182        addr_space: u32,
183        ptr: u32,
184        data: &[T; BLOCK_SIZE],
185    ) {
186        self.ctx
187            .on_memory_operation(addr_space, ptr, BLOCK_SIZE as u32);
188        self.host_write(addr_space, ptr, data)
189    }
190
191    #[inline(always)]
192    pub fn vm_read_slice<T: Copy + Debug>(
193        &mut self,
194        addr_space: u32,
195        ptr: u32,
196        len: usize,
197    ) -> &[T] {
198        self.ctx.on_memory_operation(addr_space, ptr, len as u32);
199        self.host_read_slice(addr_space, ptr, len)
200    }
201
202    #[inline(always)]
203    pub fn host_read<T: Copy + Debug, const BLOCK_SIZE: usize>(
204        &self,
205        addr_space: u32,
206        ptr: u32,
207    ) -> [T; BLOCK_SIZE] {
208        // SAFETY:
209        // - T is stack-allocated repr(C) or repr(transparent), usually u8 or F where F is the base
210        //   field
211        // - T is the exact memory cell type for this address space, satisfying the type requirement
212        unsafe { self.memory.read(addr_space, ptr) }
213    }
214
215    #[inline(always)]
216    pub fn host_write<T: Copy + Debug, const BLOCK_SIZE: usize>(
217        &mut self,
218        addr_space: u32,
219        ptr: u32,
220        data: &[T; BLOCK_SIZE],
221    ) {
222        // SAFETY:
223        // - T is stack-allocated repr(C) or repr(transparent), usually u8 or F where F is the base
224        //   field
225        // - T is the exact memory cell type for this address space, satisfying the type requirement
226        unsafe { self.memory.write(addr_space, ptr, *data) }
227    }
228
229    #[inline(always)]
230    pub fn host_read_slice<T: Copy + Debug>(&self, addr_space: u32, ptr: u32, len: usize) -> &[T] {
231        // SAFETY:
232        // - T is stack-allocated repr(C) or repr(transparent), usually u8 or F where F is the base
233        //   field
234        // - T is the exact memory cell type for this address space, satisfying the type requirement
235        // - panics if the slice is out of bounds
236        unsafe { self.memory.get_slice(addr_space, ptr, len) }
237    }
238}