Avoid building Vec<(_, [F; CHUNK])>

openvm-org · Feb 5, 2025 · 4345385 · 4345385
1 parent 814cc14
commit 4345385
Showing 1 changed file with 12 additions and 29 deletions.
diff --git a/crates/vm/src/system/memory/tree/mod.rs b/crates/vm/src/system/memory/tree/mod.rs
@@ -66,7 +66,7 @@ impl<const CHUNK: usize, F: PrimeField32> MemoryNode<CHUNK, F> {
     }
 
     fn from_memory(
-        memory: &Vec<(u64, [F; CHUNK])>,
+        memory: &[(u64, F)],
         lookup_range: Range<usize>,
         height: usize,
         from: u64,
@@ -77,9 +77,11 @@ impl<const CHUNK: usize, F: PrimeField32> MemoryNode<CHUNK, F> {
             if lookup_range.is_empty() {
                 zero_leaf.clone()
             } else {
-                debug_assert_eq!(memory[lookup_range.start].0, from);
-                debug_assert_eq!(lookup_range.end - lookup_range.start, 1);
-                MemoryNode::new_leaf(hasher.hash(&memory[lookup_range.start].1))
+                let mut values = [F::ZERO; CHUNK];
+                for (index, value) in memory[lookup_range].iter() {
+                    values[(index % CHUNK as u64) as usize] = *value;
+                }
+                MemoryNode::new_leaf(hasher.hash(&values))
             }
         } else if lookup_range.is_empty() {
             let leaf_value = hasher.hash(&[F::ZERO; CHUNK]);
@@ -140,33 +142,14 @@ impl<const CHUNK: usize, F: PrimeField32> MemoryNode<CHUNK, F> {
     ) -> MemoryNode<CHUNK, F> {
         // Construct a Vec that includes the address space in the label calculation,
         // representing the entire memory tree.
-        let mut memory_partition: Vec<(u64, [F; CHUNK])> = Vec::new();
-        for ((address_space, pointer), value) in memory.items() {
-            if pointer as usize / CHUNK >= (1 << memory_dimensions.address_height) {
-                continue;
-            }
-            debug_assert!(pointer as usize / CHUNK < (1 << memory_dimensions.address_height));
-            debug_assert!(address_space >= memory_dimensions.as_offset);
-            debug_assert!(
-                address_space - memory_dimensions.as_offset < (1 << memory_dimensions.as_height)
-            );
-            let label = (address_space, pointer / CHUNK as u32);
-            let index = memory_dimensions.label_to_index(label);
-            if memory_partition.is_empty() || memory_partition.last().unwrap().0 != index {
-                memory_partition.push((index, [F::ZERO; CHUNK]));
-            }
-            let chunk = memory_partition.last_mut().unwrap();
-            chunk.1[(pointer % CHUNK as u32) as usize] = value;
-        }
-        debug_assert!(memory_partition.is_sorted_by_key(|(addr, _)| addr));
-        debug_assert!(
-            memory_partition.last().map_or(0, |(addr, _)| *addr)
-                < (1 << memory_dimensions.overall_height())
-        );
+        let memory_items = memory
+            .items()
+            .map(|(label, value)| (memory_dimensions.label_to_index(label), value))
+            .collect::<Vec<_>>();
         let zero_leaf = MemoryNode::new_leaf(hasher.hash(&[F::ZERO; CHUNK]));
         Self::from_memory(
-            &memory_partition,
-            0..memory_partition.len(),
+            &memory_items,
+            0..memory_items.len(),
             memory_dimensions.overall_height(),
             0,
             hasher,