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#[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#[repr(C)]
104pub struct VmExecState<F, MEM, CTX> {
105 pub vm_state: VmState<F, MEM>,
107 pub ctx: CTX,
108 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 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 #[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 #[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 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 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 unsafe { self.memory.get_slice(addr_space, ptr, len) }
237 }
238}