From 2779a29a34f29c161bccaf1c3bcd11e1f757f3d0 Mon Sep 17 00:00:00 2001
From: Zach Langley <zach@axiom.xyz>
Date: Tue, 28 Jan 2025 23:09:23 -0500
Subject: [PATCH] perf: More efficient touch_range

---
 crates/vm/src/system/memory/controller/interface.rs | 6 ++----
 crates/vm/src/system/memory/merkle/mod.rs           | 9 +++++++++
 crates/vm/src/system/memory/persistent.rs           | 8 ++++++++
 crates/vm/src/system/memory/volatile/mod.rs         | 4 ----
 4 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/crates/vm/src/system/memory/controller/interface.rs b/crates/vm/src/system/memory/controller/interface.rs
index 9665447108..6263e1fe21 100644
--- 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
index c0a74d32e2..5b09956a5f 100644
--- 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
index 3664579228..1ec12e62c8 100644
--- 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
index 70878ac683..de24bd831c 100644
--- 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> {