[feat] Make online and offline memory use vectors of pages instead of hashmaps

crates/sdk/src/commit.rs

@@ -6,7 +6,10 @@ use openvm_circuit::{
         instructions::exe::VmExe,
         VmConfig,
     },
-    system::{memory::tree::MemoryNode, program::trace::VmCommittedExe},
+    system::{
+        memory::{paged_vec::AddressMap, tree::MemoryNode},
+        program::trace::VmCommittedExe,
+    },
 };
 use openvm_native_compiler::{conversion::CompilerOptions, ir::DIGEST_SIZE};
 use openvm_stark_backend::{config::StarkGenericConfig, p3_field::PrimeField32};
@@ -60,9 +63,15 @@ impl AppExecutionCommit<F> {
             .commit
             .into();
 
+        let mem_config = app_vm_config.system().memory_config;
         let init_memory_commit = MemoryNode::tree_from_memory(
             memory_dimensions,
-            &app_exe.exe.init_memory.clone().into_iter().collect(),
+            &AddressMap::from_iter(
+                mem_config.as_offset,
+                1 << mem_config.as_height,
+                1 << mem_config.pointer_max_bits,
+                app_exe.exe.init_memory.clone(),
+            ),
             &hasher,
         )
         .hash();

crates/vm/src/arch/config.rs

@@ -60,7 +60,7 @@ pub struct MemoryConfig {
 
 impl Default for MemoryConfig {
     fn default() -> Self {
-        Self::new(29, 1, 29, 29, 17, 64, 1 << 24)
+        Self::new(3, 1, 29, 29, 17, 64, 1 << 24)
     }
 }
 

crates/vm/src/arch/vm.rs

@@ -20,7 +20,7 @@ use crate::{
     arch::segment::ExecutionSegment,
     system::{
         connector::{VmConnectorPvs, DEFAULT_SUSPEND_EXIT_CODE},
-        memory::{merkle::MemoryMerklePvs, MemoryImage, CHUNK},
+        memory::{merkle::MemoryMerklePvs, paged_vec::AddressMap, MemoryImage, CHUNK},
         program::trace::VmCommittedExe,
     },
 };
@@ -113,11 +113,17 @@ where
         let exe = exe.into();
         let streams = input.into();
         let mut segments = vec![];
+        let mem_config = self.config.system().memory_config;
         let mut segment = ExecutionSegment::new(
             &self.config,
             exe.program.clone(),
             streams,
-            Some(exe.init_memory.into_iter().collect()),
+            Some(AddressMap::from_iter(
+                mem_config.as_offset,
+                1 << mem_config.as_height,
+                1 << mem_config.pointer_max_bits,
+                exe.init_memory.clone(),
+            )),
             exe.fn_bounds.clone(),
         );
         if let Some(overridden_heights) = self.overridden_heights.as_ref() {

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

@@ -25,11 +25,14 @@ use openvm_stark_backend::{
     rap::AnyRap,
     Chip, ChipUsageGetter,
 };
-use rustc_hash::FxHashMap;
 use serde::{Deserialize, Serialize};
 
 use self::interface::MemoryInterface;
-use super::{merkle::DirectCompressionBus, volatile::VolatileBoundaryChip};
+use super::{
+    merkle::DirectCompressionBus,
+    paged_vec::{AddressMap, PAGE_SIZE},
+    volatile::VolatileBoundaryChip,
+};
 use crate::{
     arch::{hasher::HasherChip, MemoryConfig},
     system::memory::{
@@ -41,7 +44,7 @@ use crate::{
             MemoryBridge, MemoryBus, MemoryReadAuxCols, MemoryReadOrImmediateAuxCols,
             MemoryWriteAuxCols, AUX_LEN,
         },
-        online::{Address, Memory, MemoryLogEntry},
+        online::{Memory, MemoryLogEntry},
         persistent::PersistentBoundaryChip,
         tree::MemoryNode,
     },
@@ -59,7 +62,7 @@ pub const BOUNDARY_AIR_OFFSET: usize = 0;
 #[derive(Debug, Clone, Copy, Serialize, Deserialize)]
 pub struct RecordId(pub usize);
 
-pub type MemoryImage<F> = FxHashMap<Address, F>;
+pub type MemoryImage<F> = AddressMap<F, PAGE_SIZE>;
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct TimestampedValues<T, const N: usize> {
@@ -229,7 +232,7 @@ impl<F: PrimeField32> MemoryController<F> {
         range_checker: SharedVariableRangeCheckerChip,
     ) -> Self {
         let range_checker_bus = range_checker.bus();
-        let initial_memory = MemoryImage::default();
+        let initial_memory = AddressMap::from_mem_config(&mem_config);
         Self {
             memory_bus,
             mem_config,
@@ -241,13 +244,14 @@ impl<F: PrimeField32> MemoryController<F> {
                     range_checker.clone(),
                 ),
             },
-            memory: Memory::new(mem_config.access_capacity),
+            memory: Memory::new(&mem_config),
             offline_memory: Arc::new(Mutex::new(OfflineMemory::new(
                 initial_memory,
                 1,
                 memory_bus,
                 range_checker.clone(),
                 mem_config.clk_max_bits,
+                mem_config,
             ))),
             access_adapters: AccessAdapterInventory::new(
                 range_checker.clone(),
@@ -285,19 +289,20 @@ impl<F: PrimeField32> MemoryController<F> {
                 compression_bus,
             ),
             merkle_chip: MemoryMerkleChip::new(memory_dims, merkle_bus, compression_bus),
-            initial_memory: MemoryImage::default(),
+            initial_memory: AddressMap::from_mem_config(&mem_config),
         };
         Self {
             memory_bus,
             mem_config,
             interface_chip,
-            memory: Memory::new(0), // it is expected that the memory will be set later
+            memory: Memory::new(&mem_config), // it is expected that the memory will be set later
             offline_memory: Arc::new(Mutex::new(OfflineMemory::new(
-                MemoryImage::default(),
+                AddressMap::from_mem_config(&mem_config),
                 CHUNK,
                 memory_bus,
                 range_checker.clone(),
                 mem_config.clk_max_bits,
+                mem_config,
             ))),
             access_adapters: AccessAdapterInventory::new(
                 range_checker.clone(),
@@ -346,14 +351,14 @@ impl<F: PrimeField32> MemoryController<F> {
             panic!("Cannot set initial memory after first timestamp");
         }
         let mut offline_memory = self.offline_memory.lock().unwrap();
-        offline_memory.set_initial_memory(memory.clone());
+        offline_memory.set_initial_memory(memory.clone(), self.mem_config);
 
         self.memory = Memory::from_image(memory.clone(), self.mem_config.access_capacity);
 
         match &mut self.interface_chip {
             MemoryInterface::Volatile { .. } => {
                 assert!(
-                    memory.is_empty(),
+                    memory.items().collect::<Vec<_>>().is_empty(),
                     "Cannot set initial memory for volatile memory"
                 );
             }

crates/vm/src/system/memory/merkle/tests/mod.rs

@@ -24,6 +24,7 @@ use crate::{
             columns::MemoryMerkleCols, tests::util::HashTestChip, MemoryDimensions,
             MemoryMerkleBus, MemoryMerkleChip,
         },
+        paged_vec::{AddressMap, PAGE_SIZE},
         tree::MemoryNode,
         Equipartition, MemoryImage,
     },
@@ -48,21 +49,24 @@ fn test<const CHUNK: usize>(
     let merkle_bus = MemoryMerkleBus(MEMORY_MERKLE_BUS);
 
     // checking validity of test data
-    for (&(address_space, pointer), value) in final_memory {
+    for ((address_space, pointer), value) in final_memory.items() {
         let label = pointer / CHUNK as u32;
         assert!(address_space - as_offset < (1 << as_height));
-        assert!(label < (1 << address_height));
-        if initial_memory.get(&(address_space, pointer)) != Some(value) {
+        assert!(pointer < ((CHUNK << address_height).div_ceil(PAGE_SIZE) * PAGE_SIZE) as u32);
+        if initial_memory.get(&(address_space, pointer)).cloned() != Some(value) {
             assert!(touched_labels.contains(&(address_space, label)));
         }
     }
-    for key in initial_memory.keys() {
-        assert!(final_memory.contains_key(key));
+    for key in initial_memory.items().map(|(key, _)| key) {
+        assert!(final_memory.get(&key).is_some());
     }
     for &(address_space, label) in touched_labels.iter() {
         let mut contains_some_key = false;
         for i in 0..CHUNK {
-            if final_memory.contains_key(&(address_space, label * CHUNK as u32 + i as u32)) {
+            if final_memory
+                .get(&(address_space, label * CHUNK as u32 + i as u32))
+                .is_some()
+            {
                 contains_some_key = true;
                 break;
             }
@@ -173,12 +177,12 @@ fn memory_to_partition<F: Default + Copy, const N: usize>(
     memory: &MemoryImage<F>,
 ) -> Equipartition<F, N> {
     let mut memory_partition = Equipartition::new();
-    for (&(address_space, pointer), value) in memory {
+    for ((address_space, pointer), value) in memory.items() {
         let label = (address_space, pointer / N as u32);
         let chunk = memory_partition
             .entry(label)
             .or_insert_with(|| [F::default(); N]);
-        chunk[(pointer % N as u32) as usize] = *value;
+        chunk[(pointer % N as u32) as usize] = value;
     }
     memory_partition
 }
@@ -192,8 +196,8 @@ fn random_test<const CHUNK: usize>(
     let mut rng = create_seeded_rng();
     let mut next_u32 = || rng.next_u64() as u32;
 
-    let mut initial_memory = MemoryImage::default();
-    let mut final_memory = MemoryImage::default();
+    let mut initial_memory = AddressMap::new(1, 2, CHUNK << height);
+    let mut final_memory = AddressMap::new(1, 2, CHUNK << height);
     let mut seen = HashSet::new();
     let mut touched_labels = BTreeSet::new();
 
@@ -210,15 +214,15 @@ fn random_test<const CHUNK: usize>(
             if is_initial && num_initial_addresses != 0 {
                 num_initial_addresses -= 1;
                 let value = BabyBear::from_canonical_u32(next_u32() % max_value);
-                initial_memory.insert((address_space, pointer), value);
-                final_memory.insert((address_space, pointer), value);
+                initial_memory.insert(&(address_space, pointer), value);
+                final_memory.insert(&(address_space, pointer), value);
             }
             if is_touched && num_touched_addresses != 0 {
                 num_touched_addresses -= 1;
                 touched_labels.insert((address_space, label));
                 if value_changes || !is_initial {
                     let value = BabyBear::from_canonical_u32(next_u32() % max_value);
-                    final_memory.insert((address_space, pointer), value);
+                    final_memory.insert(&(address_space, pointer), value);
                 }
             }
         }
@@ -260,7 +264,11 @@ fn expand_test_no_accesses() {
     };
     let mut hash_test_chip = HashTestChip::new();
 
-    let memory = MemoryImage::default();
+    let memory = AddressMap::new(
+        memory_dimensions.as_offset,
+        1 << memory_dimensions.as_height,
+        1 << memory_dimensions.address_height,
+    );
     let tree = MemoryNode::<DEFAULT_CHUNK, _>::tree_from_memory(
         memory_dimensions,
         &memory,
@@ -293,7 +301,11 @@ fn expand_test_negative() {
 
     let mut hash_test_chip = HashTestChip::new();
 
-    let memory = MemoryImage::default();
+    let memory = AddressMap::new(
+        memory_dimensions.as_offset,
+        1 << memory_dimensions.as_height,
+        1 << memory_dimensions.address_height,
+    );
     let tree = MemoryNode::<DEFAULT_CHUNK, _>::tree_from_memory(
         memory_dimensions,
         &memory,

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

@@ -6,6 +6,7 @@ pub mod merkle;
 mod offline;
 pub mod offline_checker;
 pub mod online;
+pub mod paged_vec;
 mod persistent;
 #[cfg(test)]
 mod tests;