Skip to content

Commit ea0bdd9

Browse files
jonathanpwangjonathanpwang
authored
feat: GuestMemory trait (#1574)
Not merging to main Add `GuestMemory` trait and implement for `AddressMap`. We are moving more towards a trait based style to re-use code when different types of memory might be swapped out.
1 parent 225bb16 commit ea0bdd9

File tree

3 files changed

+94
-19
lines changed

3 files changed

+94
-19
lines changed

‎crates/vm/src/system/memory/online.rs‎

Lines changed: 46 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,50 @@ use crate::{
88
system::memory::{offline::INITIAL_TIMESTAMP, MemoryImage, RecordId},
99
};
1010

11+
/// API for guest memory conforming to OpenVM ISA
12+
pub trait GuestMemory {
13+
/// Returns `[pointer:BLOCK_SIZE]_{address_space}`
14+
///
15+
/// # Safety
16+
/// The type `T` must be stack-allocated `repr(C)` or `repr(transparent)`,
17+
/// and it must be the exact type used to represent a single memory cell in
18+
/// address space `address_space`. For standard usage,
19+
/// `T` is either `u8` or `F` where `F` is the base field of the ZK backend.
20+
unsafe fn read<T: Copy, const BLOCK_SIZE: usize>(
21+
&mut self, // &mut potentially for logs?
22+
address_space: u32,
23+
pointer: u32,
24+
) -> [T; BLOCK_SIZE];
25+
26+
/// Writes `values` to `[pointer:BLOCK_SIZE]_{address_space}`
27+
///
28+
/// # Safety
29+
/// See [`GuestMemory::read`].
30+
unsafe fn write<T: Copy, const BLOCK_SIZE: usize>(
31+
&mut self,
32+
address_space: u32,
33+
pointer: u32,
34+
values: &[T; BLOCK_SIZE],
35+
);
36+
37+
/// Writes `values` to `[pointer:BLOCK_SIZE]_{address_space}` and returns
38+
/// the previous values.
39+
///
40+
/// # Safety
41+
/// See [`GuestMemory::read`].
42+
#[inline(always)]
43+
unsafe fn replace<T: Copy, const BLOCK_SIZE: usize>(
44+
&mut self,
45+
address_space: u32,
46+
pointer: u32,
47+
values: &[T; BLOCK_SIZE],
48+
) -> [T; BLOCK_SIZE] {
49+
let prev = self.read(address_space, pointer);
50+
self.write(address_space, pointer, values);
51+
prev
52+
}
53+
}
54+
1155
// TO BE DELETED
1256
#[derive(Debug, Clone, Serialize, Deserialize)]
1357
pub enum MemoryLogEntry<T> {
@@ -78,7 +122,7 @@ impl Memory {
78122
) -> (RecordId, [T; BLOCK_SIZE]) {
79123
debug_assert!(BLOCK_SIZE.is_power_of_two());
80124

81-
let prev_data = self.data.set_range((address_space, pointer), values);
125+
let prev_data = self.data.replace(address_space, pointer, values);
82126

83127
// self.log.push(MemoryLogEntry::Write {
84128
// address_space,
@@ -111,7 +155,7 @@ impl Memory {
111155
// len: N,
112156
// });
113157

114-
let values = self.data.get_range((address_space, pointer));
158+
let values = self.data.read(address_space, pointer);
115159
self.timestamp += 1;
116160
(self.last_record_id(), values)
117161
}

‎crates/vm/src/system/memory/paged_vec.rs‎

Lines changed: 46 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use openvm_instructions::exe::SparseMemoryImage;
55
use openvm_stark_backend::p3_field::PrimeField32;
66
use serde::{Deserialize, Serialize};
77

8+
use super::online::GuestMemory;
89
use crate::arch::MemoryConfig;
910

1011
/// (address_space, pointer)
@@ -71,6 +72,7 @@ impl<const PAGE_SIZE: usize> PagedVec<PAGE_SIZE> {
7172
ptr::copy_nonoverlapping(page.as_ptr().add(offset), dst, len);
7273
ptr::copy_nonoverlapping(new, page.as_mut_ptr().add(offset), len);
7374
} else {
75+
assert_eq!(start_page + 1, end_page);
7476
let offset = start % PAGE_SIZE;
7577
let first_part = PAGE_SIZE - offset;
7678
{
@@ -119,11 +121,41 @@ impl<const PAGE_SIZE: usize> PagedVec<PAGE_SIZE> {
119121
unsafe { result.assume_init() }
120122
}
121123

124+
/// # Panics
125+
/// If `start..start + size_of<BLOCK>()` is out of bounds.
126+
#[inline(always)]
127+
pub fn set<BLOCK: Copy>(&mut self, start: usize, values: &BLOCK) {
128+
let len = size_of::<BLOCK>();
129+
let start_page = start / PAGE_SIZE;
130+
let end_page = (start + len - 1) / PAGE_SIZE;
131+
let src = values as *const _ as *const u8;
132+
unsafe {
133+
if start_page == end_page {
134+
let offset = start % PAGE_SIZE;
135+
let page = self.pages[start_page].get_or_insert_with(|| vec![0u8; PAGE_SIZE]);
136+
ptr::copy_nonoverlapping(src, page.as_mut_ptr().add(offset), len);
137+
} else {
138+
assert_eq!(start_page + 1, end_page);
139+
let offset = start % PAGE_SIZE;
140+
let first_part = PAGE_SIZE - offset;
141+
{
142+
let page = self.pages[start_page].get_or_insert_with(|| vec![0u8; PAGE_SIZE]);
143+
ptr::copy_nonoverlapping(src, page.as_mut_ptr().add(offset), first_part);
144+
}
145+
let second_part = len - first_part;
146+
{
147+
let page = self.pages[end_page].get_or_insert_with(|| vec![0u8; PAGE_SIZE]);
148+
ptr::copy_nonoverlapping(src.add(first_part), page.as_mut_ptr(), second_part);
149+
}
150+
}
151+
}
152+
}
153+
122154
/// memcpy of new `values` into pages, memcpy of old existing values into new returned value.
123155
/// # Panics
124156
/// If `from..from + size_of<BLOCK>()` is out of bounds.
125157
#[inline(always)]
126-
pub fn set<BLOCK: Copy>(&mut self, from: usize, values: &BLOCK) -> BLOCK {
158+
pub fn replace<BLOCK: Copy>(&mut self, from: usize, values: &BLOCK) -> BLOCK {
127159
// Create an uninitialized array for old values.
128160
let mut result: MaybeUninit<BLOCK> = MaybeUninit::uninit();
129161
self.set_range_generic(
@@ -275,7 +307,7 @@ impl<const PAGE_SIZE: usize> AddressMap<PAGE_SIZE> {
275307
);
276308
self.paged_vecs
277309
.get_unchecked_mut((addr_space - self.as_offset) as usize)
278-
.set((ptr as usize) * size_of::<T>(), &data)
310+
.replace((ptr as usize) * size_of::<T>(), &data)
279311
}
280312
pub fn is_empty(&self) -> bool {
281313
self.paged_vecs.iter().all(|page| page.is_empty())
@@ -299,11 +331,12 @@ impl<const PAGE_SIZE: usize> AddressMap<PAGE_SIZE> {
299331
}
300332
}
301333

302-
impl<const PAGE_SIZE: usize> AddressMap<PAGE_SIZE> {
303-
/// # Safety
304-
/// - `T` **must** be the correct type for a single memory cell for `addr_space`
305-
/// - Assumes `addr_space` is within the configured memory and not out of bounds
306-
pub unsafe fn get_range<T: Copy, const N: usize>(&self, (addr_space, ptr): Address) -> [T; N] {
334+
impl<const PAGE_SIZE: usize> GuestMemory for AddressMap<PAGE_SIZE> {
335+
unsafe fn read<T: Copy, const BLOCK_SIZE: usize>(
336+
&mut self,
337+
addr_space: u32,
338+
ptr: u32,
339+
) -> [T; BLOCK_SIZE] {
307340
debug_assert_eq!(
308341
size_of::<T>(),
309342
self.cell_size[(addr_space - self.as_offset) as usize]
@@ -313,22 +346,20 @@ impl<const PAGE_SIZE: usize> AddressMap<PAGE_SIZE> {
313346
.get((ptr as usize) * size_of::<T>())
314347
}
315348

316-
/// # Safety
317-
/// - `T` **must** be the correct type for a single memory cell for `addr_space`
318-
/// - Assumes `addr_space` is within the configured memory and not out of bounds
319-
pub unsafe fn set_range<T: Copy, const N: usize>(
349+
unsafe fn write<T: Copy, const BLOCK_SIZE: usize>(
320350
&mut self,
321-
(addr_space, ptr): Address,
322-
values: &[T; N],
323-
) -> [T; N] {
351+
addr_space: u32,
352+
ptr: u32,
353+
values: &[T; BLOCK_SIZE],
354+
) {
324355
debug_assert_eq!(
325356
size_of::<T>(),
326357
self.cell_size[(addr_space - self.as_offset) as usize],
327358
"addr_space={addr_space}"
328359
);
329360
self.paged_vecs
330361
.get_unchecked_mut((addr_space - self.as_offset) as usize)
331-
.set((ptr as usize) * size_of::<T>(), values)
362+
.set((ptr as usize) * size_of::<T>(), values);
332363
}
333364
}
334365

‎crates/vm/src/system/memory/tree/public_values.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -206,7 +206,7 @@ mod tests {
206206
use super::{UserPublicValuesProof, PUBLIC_VALUES_ADDRESS_SPACE_OFFSET};
207207
use crate::{
208208
arch::{hasher::poseidon2::vm_poseidon2_hasher, SystemConfig},
209-
system::memory::{paged_vec::AddressMap, tree::MemoryNode, CHUNK},
209+
system::memory::{online::GuestMemory, paged_vec::AddressMap, tree::MemoryNode, CHUNK},
210210
};
211211

212212
type F = BabyBear;
@@ -224,7 +224,7 @@ mod tests {
224224
1 << memory_dimensions.address_height,
225225
);
226226
unsafe {
227-
memory.set_range::<F, 1>((pv_as, 15), &[F::ONE]);
227+
memory.write::<F, 1>(pv_as, 15, &[F::ONE]);
228228
}
229229
let mut expected_pvs = F::zero_vec(num_public_values);
230230
expected_pvs[15] = F::ONE;

0 commit comments

Comments
 (0)