1use std::{
2 fmt::Debug,
3 ops::{Deref, DerefMut},
4};
5
6use openvm_instructions::exe::SparseMemoryImage;
7use rand::{rngs::StdRng, SeedableRng};
8use tracing::instrument;
9
10use super::{create_memory_image, ExecutionError, Streams};
11#[cfg(feature = "metrics")]
12use crate::metrics::VmMetrics;
13use crate::{
14 arch::{execution_mode::ExecutionCtxTrait, SystemConfig},
15 system::memory::online::GuestMemory,
16};
17
18pub struct VmState<F, MEM = GuestMemory> {
20 pub instret: u64,
21 pub pc: u32,
22 pub memory: MEM,
23 pub streams: Streams<F>,
24 pub rng: StdRng,
25 pub(crate) custom_pvs: Vec<Option<F>>,
27 #[cfg(feature = "metrics")]
28 pub metrics: VmMetrics,
29}
30
31pub(super) const DEFAULT_RNG_SEED: u64 = 0;
32
33impl<F: Clone, MEM> VmState<F, MEM> {
34 pub fn new(
36 instret: u64,
37 pc: u32,
38 memory: MEM,
39 streams: impl Into<Streams<F>>,
40 seed: u64,
41 num_custom_pvs: usize,
42 ) -> Self {
43 Self {
44 instret,
45 pc,
46 memory,
47 streams: streams.into(),
48 rng: StdRng::seed_from_u64(seed),
49 custom_pvs: vec![None; num_custom_pvs],
50 #[cfg(feature = "metrics")]
51 metrics: VmMetrics::default(),
52 }
53 }
54}
55
56impl<F: Clone> VmState<F, GuestMemory> {
57 #[instrument(name = "VmState::initial", level = "debug", skip_all)]
58 pub fn initial(
59 system_config: &SystemConfig,
60 init_memory: &SparseMemoryImage,
61 pc_start: u32,
62 inputs: impl Into<Streams<F>>,
63 ) -> Self {
64 let memory = create_memory_image(&system_config.memory_config, init_memory);
65 let num_custom_pvs = if system_config.has_public_values_chip() {
66 system_config.num_public_values
67 } else {
68 0
69 };
70 VmState::new(
71 0,
72 pc_start,
73 memory,
74 inputs.into(),
75 DEFAULT_RNG_SEED,
76 num_custom_pvs,
77 )
78 }
79
80 pub fn reset(
81 &mut self,
82 init_memory: &SparseMemoryImage,
83 pc_start: u32,
84 streams: impl Into<Streams<F>>,
85 ) {
86 self.instret = 0;
87 self.pc = pc_start;
88 self.memory.memory.fill_zero();
89 self.memory.memory.set_from_sparse(init_memory);
90 self.streams = streams.into();
91 self.rng = StdRng::seed_from_u64(DEFAULT_RNG_SEED);
92 }
93}
94
95pub struct VmExecState<F, MEM, CTX> {
100 pub vm_state: VmState<F, MEM>,
102 pub exit_code: Result<Option<u32>, ExecutionError>,
104 pub ctx: CTX,
105}
106
107impl<F, MEM, CTX> VmExecState<F, MEM, CTX> {
108 pub fn new(vm_state: VmState<F, MEM>, ctx: CTX) -> Self {
109 Self {
110 vm_state,
111 ctx,
112 exit_code: Ok(None),
113 }
114 }
115}
116
117impl<F, MEM, CTX> Deref for VmExecState<F, MEM, CTX> {
118 type Target = VmState<F, MEM>;
119
120 fn deref(&self) -> &Self::Target {
121 &self.vm_state
122 }
123}
124
125impl<F, MEM, CTX> DerefMut for VmExecState<F, MEM, CTX> {
126 fn deref_mut(&mut self) -> &mut Self::Target {
127 &mut self.vm_state
128 }
129}
130
131impl<F, CTX> VmExecState<F, GuestMemory, CTX>
132where
133 CTX: ExecutionCtxTrait,
134{
135 #[inline(always)]
137 pub fn vm_read<T: Copy + Debug, const BLOCK_SIZE: usize>(
138 &mut self,
139 addr_space: u32,
140 ptr: u32,
141 ) -> [T; BLOCK_SIZE] {
142 self.ctx
143 .on_memory_operation(addr_space, ptr, BLOCK_SIZE as u32);
144 self.host_read(addr_space, ptr)
145 }
146
147 #[inline(always)]
149 pub fn vm_write<T: Copy + Debug, const BLOCK_SIZE: usize>(
150 &mut self,
151 addr_space: u32,
152 ptr: u32,
153 data: &[T; BLOCK_SIZE],
154 ) {
155 self.ctx
156 .on_memory_operation(addr_space, ptr, BLOCK_SIZE as u32);
157 self.host_write(addr_space, ptr, data)
158 }
159
160 #[inline(always)]
161 pub fn vm_read_slice<T: Copy + Debug>(
162 &mut self,
163 addr_space: u32,
164 ptr: u32,
165 len: usize,
166 ) -> &[T] {
167 self.ctx.on_memory_operation(addr_space, ptr, len as u32);
168 self.host_read_slice(addr_space, ptr, len)
169 }
170
171 #[inline(always)]
172 pub fn host_read<T: Copy + Debug, const BLOCK_SIZE: usize>(
173 &self,
174 addr_space: u32,
175 ptr: u32,
176 ) -> [T; BLOCK_SIZE] {
177 unsafe { self.memory.read(addr_space, ptr) }
182 }
183
184 #[inline(always)]
185 pub fn host_write<T: Copy + Debug, const BLOCK_SIZE: usize>(
186 &mut self,
187 addr_space: u32,
188 ptr: u32,
189 data: &[T; BLOCK_SIZE],
190 ) {
191 unsafe { self.memory.write(addr_space, ptr, *data) }
196 }
197
198 #[inline(always)]
199 pub fn host_read_slice<T: Copy + Debug>(&self, addr_space: u32, ptr: u32, len: usize) -> &[T] {
200 unsafe { self.memory.get_slice(addr_space, ptr, len) }
206 }
207}