feat: typeless AddressMap with typed APIs (#1559)

Note: this PR is not targeting `main`.
I've used `TODO` and `TEMP` to mark places in code that will need to be
cleaned up before merging to `main`.

Beginning the refactor of online memory to allow different host types in
different address spaces.
Going to touch a lot of APIs.
Focusing on stabilizing APIs - currently this PR will not improve
performance.

Tests will not all pass because I have intentionally disabled some
logging required for trace generation.
Only execution tests will pass (or run the execute benchmark).


In future PR(s):

- [ ] make `Memory` trait for execution read/write API
- [ ] better handling of type conversions for memory image
- [ ] replace the underlying memory implementation with other
implementations like mmap

Towards INT-3743

Even with wasteful conversions, execution is faster:
Before: https://github.com/openvm-org/openvm/actions/runs/14318675080
After:
https://github.com/openvm-org/openvm/actions/runs/14371335248?pr=1559

Author: jonathanpwang

crates/toolchain/instructions/src/exe.rs

@@ -5,8 +5,9 @@ use serde::{Deserialize, Serialize};
 
 use crate::program::Program;
 
-/// Memory image is a map from (address space, address) to word.
-pub type MemoryImage<F> = BTreeMap<(u32, u32), F>;
+// TODO[jpw]: delete this
+/// Memory image is a map from (address space, address * size_of<CellType>) to u8.
+pub type SparseMemoryImage = BTreeMap<(u32, u32), u8>;
 /// Stores the starting address, end address, and name of a set of function.
 pub type FnBounds = BTreeMap<u32, FnBound>;
 
@@ -22,7 +23,7 @@ pub struct VmExe<F> {
     /// Start address of pc.
     pub pc_start: u32,
     /// Initial memory image.
-    pub init_memory: MemoryImage<F>,
+    pub init_memory: SparseMemoryImage,
     /// Starting + ending bounds for each function.
     pub fn_bounds: FnBounds,
 }
@@ -40,7 +41,7 @@ impl<F> VmExe<F> {
         self.pc_start = pc_start;
         self
     }
-    pub fn with_init_memory(mut self, init_memory: MemoryImage<F>) -> Self {
+    pub fn with_init_memory(mut self, init_memory: SparseMemoryImage) -> Self {
         self.init_memory = init_memory;
         self
     }

crates/toolchain/transpiler/src/util.rs

@@ -1,7 +1,7 @@
 use std::collections::BTreeMap;
 
 use openvm_instructions::{
-    exe::MemoryImage,
+    exe::SparseMemoryImage,
     instruction::Instruction,
     riscv::{RV32_MEMORY_AS, RV32_REGISTER_NUM_LIMBS},
     utils::isize_to_field,
@@ -163,17 +163,14 @@ pub fn nop<F: PrimeField32>() -> Instruction<F> {
     }
 }
 
-/// Converts our memory image (u32 -> [u8; 4]) into Vm memory image ((as, address) -> word)
-pub fn elf_memory_image_to_openvm_memory_image<F: PrimeField32>(
+/// Converts our memory image (u32 -> [u8; 4]) into Vm memory image ((as=2, address) -> byte)
+pub fn elf_memory_image_to_openvm_memory_image(
     memory_image: BTreeMap<u32, u32>,
-) -> MemoryImage<F> {
-    let mut result = MemoryImage::new();
+) -> SparseMemoryImage {
+    let mut result = SparseMemoryImage::new();
     for (addr, word) in memory_image {
         for (i, byte) in word.to_le_bytes().into_iter().enumerate() {
-            result.insert(
-                (RV32_MEMORY_AS, addr + i as u32),
-                F::from_canonical_u8(byte),
-            );
+            result.insert((RV32_MEMORY_AS, addr + i as u32), byte);
         }
     }
     result

crates/vm/src/arch/extensions.rs

@@ -786,7 +786,7 @@ impl<F: PrimeField32, E, P> VmChipComplex<F, E, P> {
         self.base.program_chip.set_program(program);
     }
 
-    pub(crate) fn set_initial_memory(&mut self, memory: MemoryImage<F>) {
+    pub(crate) fn set_initial_memory(&mut self, memory: MemoryImage) {
         self.base.memory_controller.set_initial_memory(memory);
     }
 

crates/vm/src/arch/segment.rs

@@ -145,7 +145,7 @@ where
 {
     pub chip_complex: VmChipComplex<F, VC::Executor, VC::Periphery>,
     /// Memory image after segment was executed. Not used in trace generation.
-    pub final_memory: Option<MemoryImage<F>>,
+    pub final_memory: Option<MemoryImage>,
 
     pub since_last_segment_check: usize,
     pub trace_height_constraints: Vec<LinearConstraint>,
@@ -168,7 +168,7 @@ impl<F: PrimeField32, VC: VmConfig<F>> ExecutionSegment<F, VC> {
         config: &VC,
         program: Program<F>,
         init_streams: Streams<F>,
-        initial_memory: Option<MemoryImage<F>>,
+        initial_memory: Option<MemoryImage>,
         trace_height_constraints: Vec<LinearConstraint>,
         #[allow(unused_variables)] fn_bounds: FnBounds,
     ) -> Self {

crates/vm/src/arch/vm.rs

@@ -47,7 +47,6 @@ pub enum GenerationError {
 }
 
 /// VM memory state for continuations.
-pub type VmMemoryState<F> = MemoryImage<F>;
 
 #[derive(Clone, Default, Debug)]
 pub struct Streams<F> {
@@ -95,19 +94,19 @@ pub enum ExitCode {
 pub struct VmExecutorResult<SC: StarkGenericConfig> {
     pub per_segment: Vec<ProofInput<SC>>,
     /// When VM is running on persistent mode, public values are stored in a special memory space.
-    pub final_memory: Option<VmMemoryState<Val<SC>>>,
+    pub final_memory: Option<MemoryImage>,
 }
 
 pub struct VmExecutorNextSegmentState<F: PrimeField32> {
-    pub memory: MemoryImage<F>,
+    pub memory: MemoryImage,
     pub input: Streams<F>,
     pub pc: u32,
     #[cfg(feature = "bench-metrics")]
     pub metrics: VmMetrics,
 }
 
 impl<F: PrimeField32> VmExecutorNextSegmentState<F> {
-    pub fn new(memory: MemoryImage<F>, input: impl Into<Streams<F>>, pc: u32) -> Self {
+    pub fn new(memory: MemoryImage, input: impl Into<Streams<F>>, pc: u32) -> Self {
         Self {
             memory,
             input: input.into(),
@@ -170,12 +169,13 @@ where
         let mem_config = self.config.system().memory_config;
         let exe = exe.into();
         let mut segment_results = vec![];
-        let memory = AddressMap::from_iter(
+        let memory = AddressMap::from_sparse(
             mem_config.as_offset,
             1 << mem_config.as_height,
             1 << mem_config.pointer_max_bits,
             exe.init_memory.clone(),
         );
+
         let pc = exe.pc_start;
         let mut state = VmExecutorNextSegmentState::new(memory, input, pc);
         let mut segment_idx = 0;
@@ -271,7 +271,7 @@ where
         &self,
         exe: impl Into<VmExe<F>>,
         input: impl Into<Streams<F>>,
-    ) -> Result<Option<VmMemoryState<F>>, ExecutionError> {
+    ) -> Result<Option<MemoryImage>, ExecutionError> {
         let mut last = None;
         self.execute_and_then(
             exe,
@@ -580,7 +580,7 @@ where
         &self,
         exe: impl Into<VmExe<F>>,
         input: impl Into<Streams<F>>,
-    ) -> Result<Option<VmMemoryState<F>>, ExecutionError> {
+    ) -> Result<Option<MemoryImage>, ExecutionError> {
         self.executor.execute(exe, input)
     }
 

crates/vm/src/system/memory/controller/interface.rs

@@ -13,7 +13,7 @@ pub enum MemoryInterface<F> {
     Persistent {
         boundary_chip: PersistentBoundaryChip<F, CHUNK>,
         merkle_chip: MemoryMerkleChip<CHUNK, F>,
-        initial_memory: MemoryImage<F>,
+        initial_memory: MemoryImage,
     },
 }
 

crates/vm/src/system/memory/controller/mod.rs

@@ -3,7 +3,6 @@ use std::{
     collections::BTreeMap,
     iter,
     marker::PhantomData,
-    mem,
     sync::{Arc, Mutex},
 };
 
@@ -62,7 +61,7 @@ pub const BOUNDARY_AIR_OFFSET: usize = 0;
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
 pub struct RecordId(pub usize);
 
-pub type MemoryImage<F> = AddressMap<F, PAGE_SIZE>;
+pub type MemoryImage = AddressMap<PAGE_SIZE>;
 
 #[repr(C)]
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -98,7 +97,7 @@ pub struct MemoryController<F> {
     // Store separately to avoid smart pointer reference each time
     range_checker_bus: VariableRangeCheckerBus,
     // addr_space -> Memory data structure
-    memory: Memory<F>,
+    memory: Memory,
     /// A reference to the `OfflineMemory`. Will be populated after `finalize()`.
     offline_memory: Arc<Mutex<OfflineMemory<F>>>,
     pub access_adapters: AccessAdapterInventory<F>,
@@ -314,7 +313,7 @@ impl<F: PrimeField32> MemoryController<F> {
         }
     }
 
-    pub fn memory_image(&self) -> &MemoryImage<F> {
+    pub fn memory_image(&self) -> &MemoryImage {
         &self.memory.data
     }
 
@@ -344,7 +343,7 @@ impl<F: PrimeField32> MemoryController<F> {
         }
     }
 
-    pub fn set_initial_memory(&mut self, memory: MemoryImage<F>) {
+    pub fn set_initial_memory(&mut self, memory: MemoryImage) {
         if self.timestamp() > INITIAL_TIMESTAMP + 1 {
             panic!("Cannot set initial memory after first timestamp");
         }
@@ -379,58 +378,67 @@ impl<F: PrimeField32> MemoryController<F> {
         (record_id, data)
     }
 
-    pub fn read<const N: usize>(&mut self, address_space: F, pointer: F) -> (RecordId, [F; N]) {
+    // TEMP[jpw]: Function is safe temporarily for refactoring
+    /// # 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.
+    pub fn read<T: Copy, const N: usize>(
+        &mut self,
+        address_space: F,
+        pointer: F,
+    ) -> (RecordId, [T; N]) {
         let address_space_u32 = address_space.as_canonical_u32();
         let ptr_u32 = pointer.as_canonical_u32();
         assert!(
             address_space == F::ZERO || ptr_u32 < (1 << self.mem_config.pointer_max_bits),
             "memory out of bounds: {ptr_u32:?}",
         );
 
-        let (record_id, values) = self.memory.read::<N>(address_space_u32, ptr_u32);
+        let (record_id, values) = unsafe { self.memory.read::<T, N>(address_space_u32, ptr_u32) };
 
         (record_id, values)
     }
 
     /// Reads a word directly from memory without updating internal state.
     ///
     /// Any value returned is unconstrained.
-    pub fn unsafe_read_cell(&self, addr_space: F, ptr: F) -> F {
-        self.unsafe_read::<1>(addr_space, ptr)[0]
+    pub fn unsafe_read_cell<T: Copy>(&self, addr_space: F, ptr: F) -> T {
+        self.unsafe_read::<T, 1>(addr_space, ptr)[0]
     }
 
     /// Reads a word directly from memory without updating internal state.
     ///
     /// Any value returned is unconstrained.
-    pub fn unsafe_read<const N: usize>(&self, addr_space: F, ptr: F) -> [F; N] {
+    pub fn unsafe_read<T: Copy, const N: usize>(&self, addr_space: F, ptr: F) -> [T; N] {
         let addr_space = addr_space.as_canonical_u32();
         let ptr = ptr.as_canonical_u32();
-        array::from_fn(|i| self.memory.get(addr_space, ptr + i as u32))
+        unsafe { array::from_fn(|i| self.memory.get::<T>(addr_space, ptr + i as u32)) }
     }
 
     /// Writes `data` to the given cell.
     ///
     /// Returns the `RecordId` and previous data.
-    pub fn write_cell(&mut self, address_space: F, pointer: F, data: F) -> (RecordId, F) {
-        let (record_id, [data]) = self.write(address_space, pointer, [data]);
+    pub fn write_cell<T: Copy>(&mut self, address_space: F, pointer: F, data: T) -> (RecordId, T) {
+        let (record_id, [data]) = self.write(address_space, pointer, &[data]);
         (record_id, data)
     }
 
-    pub fn write<const N: usize>(
+    pub fn write<T: Copy, const N: usize>(
         &mut self,
         address_space: F,
         pointer: F,
-        data: [F; N],
-    ) -> (RecordId, [F; N]) {
-        assert_ne!(address_space, F::ZERO);
+        data: &[T; N],
+    ) -> (RecordId, [T; N]) {
+        debug_assert_ne!(address_space, F::ZERO);
         let address_space_u32 = address_space.as_canonical_u32();
         let ptr_u32 = pointer.as_canonical_u32();
         assert!(
             ptr_u32 < (1 << self.mem_config.pointer_max_bits),
             "memory out of bounds: {ptr_u32:?}",
         );
 
-        self.memory.write(address_space_u32, ptr_u32, data)
+        unsafe { self.memory.write::<T, N>(address_space_u32, ptr_u32, data) }
     }
 
     pub fn aux_cols_factory(&self) -> MemoryAuxColsFactory<F> {
@@ -455,26 +463,27 @@ impl<F: PrimeField32> MemoryController<F> {
     }
 
     fn replay_access_log(&mut self) {
-        let log = mem::take(&mut self.memory.log);
-        if log.is_empty() {
-            // Online memory logs may be empty, but offline memory may be replayed from external sources.
-            // In these cases, we skip the calls to replay access logs because `set_log_capacity` would
-            // panic.
-            tracing::debug!("skipping replay_access_log");
-            return;
-        }
-
-        let mut offline_memory = self.offline_memory.lock().unwrap();
-        offline_memory.set_log_capacity(log.len());
-
-        for entry in log {
-            Self::replay_access(
-                entry,
-                &mut offline_memory,
-                &mut self.interface_chip,
-                &mut self.access_adapters,
-            );
-        }
+        unimplemented!();
+        // let log = mem::take(&mut self.memory.log);
+        // if log.is_empty() {
+        //     // Online memory logs may be empty, but offline memory may be replayed from external sources.
+        //     // In these cases, we skip the calls to replay access logs because `set_log_capacity` would
+        //     // panic.
+        //     tracing::debug!("skipping replay_access_log");
+        //     return;
+        // }
+
+        // let mut offline_memory = self.offline_memory.lock().unwrap();
+        // offline_memory.set_log_capacity(log.len());
+
+        // for entry in log {
+        //     Self::replay_access(
+        //         entry,
+        //         &mut offline_memory,
+        //         &mut self.interface_chip,
+        //         &mut self.access_adapters,
+        //     );
+        // }
     }
 
     /// Low-level API to replay a single memory access log entry and populate the [OfflineMemory], [MemoryInterface], and `AccessAdapterInventory`.
@@ -703,13 +712,13 @@ impl<F: PrimeField32> MemoryController<F> {
     pub fn offline_memory(&self) -> Arc<Mutex<OfflineMemory<F>>> {
         self.offline_memory.clone()
     }
-    pub fn get_memory_logs(&self) -> &Vec<MemoryLogEntry<F>> {
+    pub fn get_memory_logs(&self) -> &Vec<MemoryLogEntry<u8>> {
         &self.memory.log
     }
-    pub fn set_memory_logs(&mut self, logs: Vec<MemoryLogEntry<F>>) {
+    pub fn set_memory_logs(&mut self, logs: Vec<MemoryLogEntry<u8>>) {
         self.memory.log = logs;
     }
-    pub fn take_memory_logs(&mut self) -> Vec<MemoryLogEntry<F>> {
+    pub fn take_memory_logs(&mut self) -> Vec<MemoryLogEntry<u8>> {
         std::mem::take(&mut self.memory.log)
     }
 }
@@ -855,9 +864,9 @@ mod tests {
 
             if rng.gen_bool(0.5) {
                 let data = F::from_canonical_u32(rng.gen_range(0..1 << 30));
-                memory_controller.write(address_space, pointer, [data]);
+                memory_controller.write(address_space, pointer, &[data]);
             } else {
-                memory_controller.read::<1>(address_space, pointer);
+                memory_controller.read::<F, 1>(address_space, pointer);
             }
         }
         assert!(memory_controller