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

openvm_circuit/arch/
state.rs
