perf: More efficient touch_range

openvm-org · Jan 29, 2025 · 2779a29 · 2779a29
1 parent 442a9df
commit 2779a29
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 8 deletions.
diff --git a/crates/vm/src/system/memory/controller/interface.rs b/crates/vm/src/system/memory/controller/interface.rs
@@ -48,10 +48,8 @@ impl<F: PrimeField32> MemoryInterface<F> {
                 merkle_chip,
                 ..
             } => {
-                for offset in 0..len {
-                    boundary_chip.touch_address(addr_space, pointer + offset);
-                    merkle_chip.touch_address(addr_space, pointer + offset);
-                }
+                boundary_chip.touch_range(addr_space, pointer, len);
+                merkle_chip.touch_range(addr_space, pointer, len);
             }
         }
     }

diff --git a/crates/vm/src/system/memory/merkle/mod.rs b/crates/vm/src/system/memory/merkle/mod.rs
@@ -76,4 +76,13 @@ impl<const CHUNK: usize, F: PrimeField32> MemoryMerkleChip<CHUNK, F> {
             address / CHUNK as u32,
         );
     }
+
+    pub fn touch_range(&mut self, address_space: u32, address: u32, len: u32) {
+        let as_label = address_space - self.air.memory_dimensions.as_offset;
+        let first_address_label = address / CHUNK as u32;
+        let last_address_label = (address + len - 1) / CHUNK as u32;
+        for address_label in first_address_label..=last_address_label {
+            self.touch_node(0, as_label, address_label);
+        }
+    }
 }
diff --git a/crates/vm/src/system/memory/persistent.rs b/crates/vm/src/system/memory/persistent.rs
@@ -192,6 +192,14 @@ impl<const CHUNK: usize, F: PrimeField32> PersistentBoundaryChip<F, CHUNK> {
         self.touched_labels.touch(address_space, label);
     }
 
+    pub fn touch_range(&mut self, address_space: u32, pointer: u32, len: u32) {
+        let start_label = pointer / CHUNK as u32;
+        let end_label = (pointer + len - 1) / CHUNK as u32;
+        for label in start_label..=end_label {
+            self.touched_labels.touch(address_space, label);
+        }
+    }
+
     pub fn finalize(
         &mut self,
         initial_memory: &MemoryImage<F>,

diff --git a/crates/vm/src/system/memory/volatile/mod.rs b/crates/vm/src/system/memory/volatile/mod.rs
@@ -163,10 +163,6 @@ impl<F: Field> VolatileBoundaryChip<F> {
     pub fn touch_address(&mut self, addr_space: u32, pointer: u32) {
         self.touched_addresses.insert((addr_space, pointer));
     }
-
-    pub fn all_addresses(&self) -> Vec<(u32, u32)> {
-        self.touched_addresses.iter().cloned().collect()
-    }
 }
 
 impl<F: PrimeField32> VolatileBoundaryChip<F> {