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.

Author: jonathanpwang

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

@@ -8,6 +8,50 @@ use crate::{
     system::memory::{offline::INITIAL_TIMESTAMP, MemoryImage, RecordId},
 };
 
+/// API for guest memory conforming to OpenVM ISA
+pub trait GuestMemory {
+    /// Returns `[pointer:BLOCK_SIZE]_{address_space}`
+    ///
+    /// # Safety
+    /// The type `T` must be stack-allocated `repr(C)` or `repr(transparent)`,
+    /// and it must be the exact type used to represent a single memory cell in
+    /// address space `address_space`. For standard usage,
+    /// `T` is either `u8` or `F` where `F` is the base field of the ZK backend.
+    unsafe fn read<T: Copy, const BLOCK_SIZE: usize>(
+        &mut self, // &mut potentially for logs?
+        address_space: u32,
+        pointer: u32,
+    ) -> [T; BLOCK_SIZE];
+
+    /// Writes `values` to `[pointer:BLOCK_SIZE]_{address_space}`
+    ///
+    /// # Safety
+    /// See [`GuestMemory::read`].
+    unsafe fn write<T: Copy, const BLOCK_SIZE: usize>(
+        &mut self,
+        address_space: u32,
+        pointer: u32,
+        values: &[T; BLOCK_SIZE],
+    );
+
+    /// Writes `values` to `[pointer:BLOCK_SIZE]_{address_space}` and returns
+    /// the previous values.
+    ///
+    /// # Safety
+    /// See [`GuestMemory::read`].
+    #[inline(always)]
+    unsafe fn replace<T: Copy, const BLOCK_SIZE: usize>(
+        &mut self,
+        address_space: u32,
+        pointer: u32,
+        values: &[T; BLOCK_SIZE],
+    ) -> [T; BLOCK_SIZE] {
+        let prev = self.read(address_space, pointer);
+        self.write(address_space, pointer, values);
+        prev
+    }
+}
+
 // TO BE DELETED
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum MemoryLogEntry<T> {
@@ -80,7 +124,7 @@ impl Memory {
     ) -> (RecordId, [T; BLOCK_SIZE]) {
         debug_assert!(BLOCK_SIZE.is_power_of_two());
 
-        let prev_data = self.data.set_range((address_space, pointer), values);
+        let prev_data = self.data.replace(address_space, pointer, values);
 
         // self.log.push(MemoryLogEntry::Write {
         //     address_space,
@@ -113,7 +157,7 @@ impl Memory {
         //     len: N,
         // });
 
-        let values = self.data.get_range((address_space, pointer));
+        let values = self.data.read(address_space, pointer);
         self.timestamp += 1;
         (self.last_record_id(), values)
     }

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

@@ -5,6 +5,7 @@ use openvm_instructions::exe::SparseMemoryImage;
 use openvm_stark_backend::p3_field::PrimeField32;
 use serde::{Deserialize, Serialize};
 
+use super::online::GuestMemory;
 use crate::arch::MemoryConfig;
 
 /// (address_space, pointer)
@@ -72,6 +73,7 @@ impl<const PAGE_SIZE: usize> PagedVec<PAGE_SIZE> {
                 ptr::copy_nonoverlapping(page.as_ptr().add(offset), dst, len);
                 ptr::copy_nonoverlapping(new, page.as_mut_ptr().add(offset), len);
             } else {
+                assert_eq!(start_page + 1, end_page);
                 let offset = start % PAGE_SIZE;
                 let first_part = PAGE_SIZE - offset;
                 {
@@ -120,11 +122,41 @@ impl<const PAGE_SIZE: usize> PagedVec<PAGE_SIZE> {
         unsafe { result.assume_init() }
     }
 
+    /// # Panics
+    /// If `start..start + size_of<BLOCK>()` is out of bounds.
+    #[inline(always)]
+    pub fn set<BLOCK: Copy>(&mut self, start: usize, values: &BLOCK) {
+        let len = size_of::<BLOCK>();
+        let start_page = start / PAGE_SIZE;
+        let end_page = (start + len - 1) / PAGE_SIZE;
+        let src = values as *const _ as *const u8;
+        unsafe {
+            if start_page == end_page {
+                let offset = start % PAGE_SIZE;
+                let page = self.pages[start_page].get_or_insert_with(|| vec![0u8; PAGE_SIZE]);
+                ptr::copy_nonoverlapping(src, page.as_mut_ptr().add(offset), len);
+            } else {
+                assert_eq!(start_page + 1, end_page);
+                let offset = start % PAGE_SIZE;
+                let first_part = PAGE_SIZE - offset;
+                {
+                    let page = self.pages[start_page].get_or_insert_with(|| vec![0u8; PAGE_SIZE]);
+                    ptr::copy_nonoverlapping(src, page.as_mut_ptr().add(offset), first_part);
+                }
+                let second_part = len - first_part;
+                {
+                    let page = self.pages[end_page].get_or_insert_with(|| vec![0u8; PAGE_SIZE]);
+                    ptr::copy_nonoverlapping(src.add(first_part), page.as_mut_ptr(), second_part);
+                }
+            }
+        }
+    }
+
     /// memcpy of new `values` into pages, memcpy of old existing values into new returned value.
     /// # Panics
     /// If `from..from + size_of<BLOCK>()` is out of bounds.
     #[inline(always)]
-    pub fn set<BLOCK: Copy>(&mut self, from: usize, values: &BLOCK) -> BLOCK {
+    pub fn replace<BLOCK: Copy>(&mut self, from: usize, values: &BLOCK) -> BLOCK {
         // Create an uninitialized array for old values.
         let mut result: MaybeUninit<BLOCK> = MaybeUninit::uninit();
         self.set_range_generic(
@@ -278,7 +310,7 @@ impl<const PAGE_SIZE: usize> AddressMap<PAGE_SIZE> {
         );
         self.paged_vecs
             .get_unchecked_mut((addr_space - self.as_offset) as usize)
-            .set((ptr as usize) * size_of::<T>(), &data)
+            .replace((ptr as usize) * size_of::<T>(), &data)
     }
     pub fn is_empty(&self) -> bool {
         self.paged_vecs.iter().all(|page| page.is_empty())
@@ -302,11 +334,12 @@ impl<const PAGE_SIZE: usize> AddressMap<PAGE_SIZE> {
     }
 }
 
-impl<const PAGE_SIZE: usize> AddressMap<PAGE_SIZE> {
-    /// # Safety
-    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
-    /// - Assumes `addr_space` is within the configured memory and not out of bounds
-    pub unsafe fn get_range<T: Copy, const N: usize>(&self, (addr_space, ptr): Address) -> [T; N] {
+impl<const PAGE_SIZE: usize> GuestMemory for AddressMap<PAGE_SIZE> {
+    unsafe fn read<T: Copy, const BLOCK_SIZE: usize>(
+        &mut self,
+        addr_space: u32,
+        ptr: u32,
+    ) -> [T; BLOCK_SIZE] {
         debug_assert_eq!(
             size_of::<T>(),
             self.cell_size[(addr_space - self.as_offset) as usize]
@@ -316,22 +349,20 @@ impl<const PAGE_SIZE: usize> AddressMap<PAGE_SIZE> {
             .get((ptr as usize) * size_of::<T>())
     }
 
-    /// # Safety
-    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
-    /// - Assumes `addr_space` is within the configured memory and not out of bounds
-    pub unsafe fn set_range<T: Copy, const N: usize>(
+    unsafe fn write<T: Copy, const BLOCK_SIZE: usize>(
         &mut self,
-        (addr_space, ptr): Address,
-        values: &[T; N],
-    ) -> [T; N] {
+        addr_space: u32,
+        ptr: u32,
+        values: &[T; BLOCK_SIZE],
+    ) {
         debug_assert_eq!(
             size_of::<T>(),
             self.cell_size[(addr_space - self.as_offset) as usize],
             "addr_space={addr_space}"
         );
         self.paged_vecs
             .get_unchecked_mut((addr_space - self.as_offset) as usize)
-            .set((ptr as usize) * size_of::<T>(), values)
+            .set((ptr as usize) * size_of::<T>(), values);
     }
 }
 

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

@@ -206,7 +206,7 @@ mod tests {
     use super::{UserPublicValuesProof, PUBLIC_VALUES_ADDRESS_SPACE_OFFSET};
     use crate::{
         arch::{hasher::poseidon2::vm_poseidon2_hasher, SystemConfig},
-        system::memory::{paged_vec::AddressMap, tree::MemoryNode, CHUNK},
+        system::memory::{online::GuestMemory, paged_vec::AddressMap, tree::MemoryNode, CHUNK},
     };
 
     type F = BabyBear;
@@ -224,7 +224,7 @@ mod tests {
             1 << memory_dimensions.address_height,
         );
         unsafe {
-            memory.set_range::<F, 1>((pv_as, 15), &[F::ONE]);
+            memory.write::<F, 1>(pv_as, 15, &[F::ONE]);
         }
         let mut expected_pvs = F::zero_vec(num_public_values);
         expected_pvs[15] = F::ONE;