Storage cache: Use HashMap, group by ContractId

.changes/changed/1010.md

@@ -0,0 +1 @@
+Storage cache: Use HashMap, group by ContractId

fuel-vm/src/interpreter.rs

@@ -14,10 +14,7 @@ use crate::{
     verification,
 };
 use alloc::{
-    collections::{
-        BTreeMap,
-        BTreeSet,
-    },
+    collections::BTreeSet,
     vec::Vec,
 };
 use core::{
@@ -144,7 +141,7 @@ pub struct Interpreter<M, S, Tx = (), Ecal = NotSupportedEcal, V = verification:
     /// overhead of accessing the storage. This is especially important for charging
     /// for new bytes written.
     storage_slot_cache:
-        alloc::collections::BTreeMap<(ContractId, Bytes32), Option<Vec<u8>>>,
+        hashbrown::HashMap<ContractId, hashbrown::HashMap<Bytes32, Option<Vec<u8>>>>,
 }
 
 /// Interpreter parameters
@@ -252,9 +249,10 @@ impl<M, S, Tx, Ecal, V> Interpreter<M, S, Tx, Ecal, V> {
     /// ## Warning
     /// This function is excempt from semver guarantees and may be
     /// removed or changed without a major version bump.
-    pub const fn bench_storage_slot_cache(
+    pub fn bench_storage_slot_cache(
         &self,
-    ) -> &BTreeMap<(ContractId, Bytes32), Option<Vec<u8>>> {
+    ) -> &hashbrown::HashMap<ContractId, hashbrown::HashMap<Bytes32, Option<Vec<u8>>>>
+    {
         &self.storage_slot_cache
     }
 
@@ -264,7 +262,8 @@ impl<M, S, Tx, Ecal, V> Interpreter<M, S, Tx, Ecal, V> {
     /// removed or changed without a major version bump.
     pub fn bench_storage_slot_cache_mut(
         &mut self,
-    ) -> &mut BTreeMap<(ContractId, Bytes32), Option<Vec<u8>>> {
+    ) -> &mut hashbrown::HashMap<ContractId, hashbrown::HashMap<Bytes32, Option<Vec<u8>>>>
+    {
         &mut self.storage_slot_cache
     }
 

fuel-vm/src/interpreter/diff/storage.rs

@@ -219,7 +219,6 @@ where
             if let Change::Storage(Previous(from)) = change {
                 match from {
                     StorageState::State(MappableState { key, value }) => {
-                        let cache_key = (*key.contract_id(), *key.state_key());
                         if let Some(value) = value {
                             StorageMutate::<ContractsState>::insert(
                                 &mut self.storage,
@@ -228,7 +227,9 @@ where
                             )
                             .unwrap();
                             self.storage_slot_cache
-                                .insert(cache_key, Some(value.as_ref().to_vec()));
+                                .entry(*key.contract_id())
+                                .or_default()
+                                .insert(*key.state_key(), Some(value.as_ref().to_vec()));
                         } else {
                             // The slot is absent in the target state; delete it.
                             let _ = StorageMutate::<ContractsState>::take(
@@ -240,9 +241,16 @@ where
                             //  Some(None)    – already a valid warm-absent hit; keep.
                             //  None          – cold; leave cold so the next read pays
                             //                  the correct storage_read_cold gas cost.
-                            if let Some(Some(_)) = self.storage_slot_cache.get(&cache_key)
-                            {
-                                self.storage_slot_cache.insert(cache_key, None);
+                            if matches!(
+                                self.storage_slot_cache
+                                    .get(key.contract_id())
+                                    .and_then(|m| m.get(key.state_key())),
+                                Some(Some(_))
+                            ) {
+                                self.storage_slot_cache
+                                    .entry(*key.contract_id())
+                                    .or_default()
+                                    .insert(*key.state_key(), None);
                             }
                         }
                     }

fuel-vm/src/interpreter/diff/tests.rs

@@ -320,7 +320,7 @@ fn reset_vm_state_updates_cache_for_modified_slot() {
 
     let contract_id = ContractId::default();
     let slot_key = Bytes32::default();
-    let cache_key = (contract_id, slot_key);
+
     let v1 = alloc::vec![0xAAu8; 32];
 
     // Insert v1 via the Record wrapper so the change is recorded.
@@ -333,7 +333,11 @@ fn reset_vm_state_updates_cache_for_modified_slot() {
 
     // Poison the cache with a stale value to simulate the pre-fix bug.
     let v_stale = alloc::vec![0xFFu8; 32];
-    latest.storage_slot_cache.insert(cache_key, Some(v_stale));
+    latest
+        .storage_slot_cache
+        .entry(contract_id)
+        .or_default()
+        .insert(slot_key, Some(v_stale));
 
     // Build the diff: storage_diff captures "slot was None → v1".
     // After converting to InitialVmState the change carries v1 as the target.
@@ -349,7 +353,13 @@ fn reset_vm_state_updates_cache_for_modified_slot() {
     latest.reset_vm_state(&diff);
 
     // The cache must now hold v1, not the stale value.
-    assert_eq!(latest.storage_slot_cache.get(&cache_key), Some(&Some(v1)));
+    assert_eq!(
+        latest
+            .storage_slot_cache
+            .get(&contract_id)
+            .and_then(|m| m.get(&slot_key)),
+        Some(&Some(v1))
+    );
 }
 
 /// When a storage slot that was *absent* in the initial state is marked as such by
@@ -363,7 +373,7 @@ fn reset_vm_state_removes_newly_created_slot_from_cache() {
 
     let contract_id = ContractId::default();
     let slot_key = Bytes32::default();
-    let cache_key = (contract_id, slot_key);
+
     let v1 = alloc::vec![0xBBu8; 32];
 
     // Pre-populate backing storage with v1 *before* wrapping in Record.
@@ -389,7 +399,11 @@ fn reset_vm_state_removes_newly_created_slot_from_cache() {
     .unwrap();
 
     // Poison the cache with a stale Some(v1) to simulate the pre-fix bug.
-    latest.storage_slot_cache.insert(cache_key, Some(v1));
+    latest
+        .storage_slot_cache
+        .entry(contract_id)
+        .or_default()
+        .insert(slot_key, Some(v1));
 
     // Build the diff: the recorded Take produces "slot v1 → absent".
     // After converting to InitialVmState the change carries None as the target.
@@ -406,7 +420,13 @@ fn reset_vm_state_removes_newly_created_slot_from_cache() {
 
     // The stale Some(v1) must be corrected to None (warm-absent), not left
     // as a stale present value.
-    assert_eq!(latest.storage_slot_cache.get(&cache_key), Some(&None));
+    assert_eq!(
+        latest
+            .storage_slot_cache
+            .get(&contract_id)
+            .and_then(|m| m.get(&slot_key)),
+        Some(&None)
+    );
 }
 
 /// A slot that was **cold** (never accessed, not in cache) before execution and
@@ -422,7 +442,7 @@ fn reset_vm_state_does_not_warm_absent_cold_slot() {
 
     let contract_id = ContractId::default();
     let slot_key = Bytes32::default();
-    let cache_key = (contract_id, slot_key);
+
     let v1 = alloc::vec![0xAAu8; 32];
 
     // Pre-populate backing storage with v1 before wrapping in Record,
@@ -451,7 +471,13 @@ fn reset_vm_state_does_not_warm_absent_cold_slot() {
     .unwrap();
 
     // Confirm: slot is cold before reset (no cache entry at all).
-    assert_eq!(latest.storage_slot_cache.get(&cache_key), None);
+    assert_eq!(
+        latest
+            .storage_slot_cache
+            .get(&contract_id)
+            .and_then(|m| m.get(&slot_key)),
+        None
+    );
 
     // Build the diff: the Take produces "slot v1 → absent".
     // After converting to InitialVmState, Previous(None) is stored for this slot.
@@ -470,7 +496,10 @@ fn reset_vm_state_does_not_warm_absent_cold_slot() {
     // reset_vm_state should leave it cold (no cache entry), so the next read
     // is charged storage_read_cold, not storage_read_hot.
     assert_eq!(
-        latest.storage_slot_cache.get(&cache_key),
+        latest
+            .storage_slot_cache
+            .get(&contract_id)
+            .and_then(|m| m.get(&slot_key)),
         None,
         "cold-absent slot must not be warmed by reset_vm_state"
     );
@@ -487,7 +516,7 @@ fn reset_vm_state_preserves_warm_absent_cache_entry() {
 
     let contract_id = ContractId::default();
     let slot_key = Bytes32::default();
-    let cache_key = (contract_id, slot_key);
+
     let v1 = alloc::vec![0xCCu8; 32];
 
     // Pre-populate backing storage with v1 before wrapping in Record.
@@ -513,7 +542,11 @@ fn reset_vm_state_preserves_warm_absent_cache_entry() {
     .unwrap();
 
     // Simulate the slot already being cached as absent (warm-absent) before reset.
-    latest.storage_slot_cache.insert(cache_key, None);
+    latest
+        .storage_slot_cache
+        .entry(contract_id)
+        .or_default()
+        .insert(slot_key, None);
 
     // Build the diff: the recorded Take produces "slot v1 → absent".
     let storage_diff: Diff<InitialVmState> = latest.storage_diff().into();
@@ -530,7 +563,10 @@ fn reset_vm_state_preserves_warm_absent_cache_entry() {
     // The warm-absent entry must be preserved: the slot is still absent and
     // the cache correctly reflects that, so the next read should be hot.
     assert_eq!(
-        latest.storage_slot_cache.get(&cache_key),
+        latest
+            .storage_slot_cache
+            .get(&contract_id)
+            .and_then(|m| m.get(&slot_key)),
         Some(&None),
         "warm-absent cache entry must be preserved by reset_vm_state"
     );

fuel-vm/src/interpreter/storage.rs

@@ -41,9 +41,11 @@ where
     where
         F: FnOnce(&mut MemoryInstance, Option<&[u8]>) -> R,
     {
-        let cache_key = (contract_id, key);
-
-        if let Some(v) = self.storage_slot_cache.get(&cache_key) {
+        if let Some(v) = self
+            .storage_slot_cache
+            .get(&contract_id)
+            .and_then(|m| m.get(&key))
+        {
             // Cache hit
             let gas_charge_units = v.as_ref().map(|data| data.len() as u64).unwrap_or(0);
             let r = f(self.memory.as_mut(), v.as_deref());
@@ -69,7 +71,10 @@ where
                 .map_err(PanicReason::from)?,
             gas_charge_units,
         )?;
-        self.storage_slot_cache.insert(cache_key, value);
+        self.storage_slot_cache
+            .entry(contract_id)
+            .or_default()
+            .insert(key, value);
         Ok(r)
     }
 
@@ -91,8 +96,11 @@ where
         contract_id: ContractId,
         key: Bytes32,
     ) -> Result<usize, RuntimeError<S::DataError>> {
-        let cache_key = (contract_id, key);
-        if let Some(v) = self.storage_slot_cache.get(&cache_key) {
+        if let Some(v) = self
+            .storage_slot_cache
+            .get(&contract_id)
+            .and_then(|m| m.get(&key))
+        {
             return Ok(v.as_ref().map(|d| d.len()).unwrap_or(0));
         }
         let value = StorageRead::<ContractsState>::read_alloc(
@@ -101,7 +109,10 @@ where
         )
         .map_err(RuntimeError::Storage)?;
         let len = value.as_ref().map(|d| d.len()).unwrap_or(0);
-        self.storage_slot_cache.insert(cache_key, value);
+        self.storage_slot_cache
+            .entry(contract_id)
+            .or_default()
+            .insert(key, value);
         Ok(len)
     }
 
@@ -116,12 +127,14 @@ where
         if (value.len() as u64) > max_size {
             return Err(RuntimeError::Recoverable(PanicReason::StorageOutOfBounds));
         }
-        let cache_key = (contract_id, key);
         self.storage
             .contract_state_insert(&contract_id, &key, &value)
             .map_err(RuntimeError::Storage)?;
         let gas_charge_units = value.len() as u64;
-        self.storage_slot_cache.insert(cache_key, Some(value));
+        self.storage_slot_cache
+            .entry(contract_id)
+            .or_default()
+            .insert(key, Some(value));
         self.dependent_gas_charge(
             self.gas_costs()
                 .storage_write()
@@ -179,8 +192,10 @@ where
             .map_err(RuntimeError::Storage)?;
         for key in key_range(key, range) {
             let key = key.ok_or(PanicReason::TooManySlots)?;
-            let cache_key = (contract_id, key);
-            self.storage_slot_cache.insert(cache_key, None);
+            self.storage_slot_cache
+                .entry(contract_id)
+                .or_default()
+                .insert(key, None);
         }
         Ok(())
     }