openvm_circuit/arch/
state.rs

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
18/// Represents the core state of a VM.
19pub 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    /// The public values of the PublicValuesAir when it exists
26    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    /// `num_custom_pvs` should only be nonzero when the PublicValuesAir exists.
35    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
95/// Represents the full execution state of a VM during execution.
96/// The global state is generic in guest memory `MEM` and additional context `CTX`.
97/// The host state is execution context specific.
98// @dev: Do not confuse with `ExecutionState` struct.
99pub struct VmExecState<F, MEM, CTX> {
100    /// Core VM state
101    pub vm_state: VmState<F, MEM>,
102    /// Execution-specific fields
103    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    /// Runtime read operation for a block of memory
136    #[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    /// Runtime write operation for a block of memory
148    #[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        // SAFETY:
178        // - T is stack-allocated repr(C) or repr(transparent), usually u8 or F where F is the base
179        //   field
180        // - T is the exact memory cell type for this address space, satisfying the type requirement
181        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        // SAFETY:
192        // - T is stack-allocated repr(C) or repr(transparent), usually u8 or F where F is the base
193        //   field
194        // - T is the exact memory cell type for this address space, satisfying the type requirement
195        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        // SAFETY:
201        // - T is stack-allocated repr(C) or repr(transparent), usually u8 or F where F is the base
202        //   field
203        // - T is the exact memory cell type for this address space, satisfying the type requirement
204        // - panics if the slice is out of bounds
205        unsafe { self.memory.get_slice(addr_space, ptr, len) }
206    }
207}