Persist reverted account and storage slot lookups in JournaledState

crates/primitives/src/state.rs

@@ -41,6 +41,8 @@ bitflags! {
         /// used only for pre spurious dragon hardforks where existing and empty were two separate states.
         /// it became same state after EIP-161: State trie clearing
         const LoadedAsNotExisting = 0b0001000;
+        /// used to mark account as cold
+        const Cold = 0b0010000;
     }
 }
 
@@ -100,6 +102,21 @@ impl Account {
         self.status -= AccountStatus::Created;
     }
 
+    /// Mark account as cold.
+    pub fn mark_cold(&mut self) {
+        self.status |= AccountStatus::Cold;
+    }
+
+    /// Mark account as warm and return true if it was previously cold.
+    pub fn mark_warm(&mut self) -> bool {
+        if self.status.contains(AccountStatus::Cold) {
+            self.status -= AccountStatus::Cold;
+            true
+        } else {
+            false
+        }
+    }
+
     /// Is account loaded as not existing from database
     /// This is needed for pre spurious dragon hardforks where
     /// existing and empty were two separate states.
@@ -147,6 +164,8 @@ pub struct StorageSlot {
     pub previous_or_original_value: U256,
     /// When loaded with sload present value is set to original value
     pub present_value: U256,
+    /// Represents if the storage slot is cold.
+    pub is_cold: bool,
 }
 
 impl StorageSlot {
@@ -155,6 +174,7 @@ impl StorageSlot {
         Self {
             previous_or_original_value: original,
             present_value: original,
+            is_cold: false,
         }
     }
 
@@ -163,6 +183,7 @@ impl StorageSlot {
         Self {
             previous_or_original_value,
             present_value,
+            is_cold: false,
         }
     }
 
@@ -180,6 +201,16 @@ impl StorageSlot {
     pub fn present_value(&self) -> U256 {
         self.present_value
     }
+
+    /// Marks the storage slot as cold.
+    pub fn mark_cold(&mut self) {
+        self.is_cold = true;
+    }
+
+    /// Marks the storage slot as warm and returns a bool indicating if it was previously cold.
+    pub fn mark_warm(&mut self) -> bool {
+        core::mem::replace(&mut self.is_cold, false)
+    }
 }
 
 /// AccountInfo account information.
@@ -348,4 +379,24 @@ mod tests {
         assert!(account.is_touched());
         assert!(!account.is_selfdestructed());
     }
+
+    #[test]
+    fn account_is_cold() {
+        let mut account = Account::default();
+
+        // Account is not cold by default
+        assert!(!account.status.contains(crate::AccountStatus::Cold));
+
+        // When marking warm account as warm again, it should return false
+        assert!(!account.mark_warm());
+
+        // Mark account as cold
+        account.mark_cold();
+
+        // Account is cold
+        assert!(account.status.contains(crate::AccountStatus::Cold));
+
+        // When marking cold account as warm, it should return true
+        assert!(account.mark_warm());
+    }
 }

crates/revm/src/journaled_state.rs

@@ -322,7 +322,7 @@ impl JournaledState {
         for entry in journal_entries.into_iter().rev() {
             match entry {
                 JournalEntry::AccountLoaded { address } => {
-                    state.remove(&address);
+                    state.get_mut(&address).unwrap().mark_cold();
                 }
                 JournalEntry::AccountTouched { address } => {
                     if is_spurious_dragon_enabled && address == PRECOMPILE3 {
@@ -378,7 +378,7 @@ impl JournaledState {
                     if let Some(had_value) = had_value {
                         storage.get_mut(&key).unwrap().present_value = had_value;
                     } else {
-                        storage.remove(&key);
+                        storage.get_mut(&key).unwrap().mark_cold();
                     }
                 }
                 JournalEntry::TransientStorageChange {
@@ -555,8 +555,12 @@ impl JournaledState {
         address: Address,
         db: &mut DB,
     ) -> Result<(&mut Account, bool), EVMError<DB::Error>> {
-        Ok(match self.state.entry(address) {
-            Entry::Occupied(entry) => (entry.into_mut(), false),
+        let (value, is_cold) = match self.state.entry(address) {
+            Entry::Occupied(entry) => {
+                let account = entry.into_mut();
+                let is_cold = account.mark_warm();
+                (account, is_cold)
+            }
             Entry::Vacant(vac) => {
                 let account =
                     if let Some(account) = db.basic(address).map_err(EVMError::Database)? {
@@ -565,18 +569,22 @@ impl JournaledState {
                         Account::new_not_existing()
                     };
 
-                // journal loading of account. AccessList touch.
-                self.journal
-                    .last_mut()
-                    .unwrap()
-                    .push(JournalEntry::AccountLoaded { address });
-
                 // precompiles are warm loaded so we need to take that into account
                 let is_cold = !self.warm_preloaded_addresses.contains(&address);
 
                 (vac.insert(account), is_cold)
             }
-        })
+        };
+
+        // journal loading of cold account.
+        if is_cold {
+            self.journal
+                .last_mut()
+                .unwrap()
+                .push(JournalEntry::AccountLoaded { address });
+        }
+
+        Ok((value, is_cold))
     }
 
     /// Load account from database to JournaledState.
@@ -641,31 +649,39 @@ impl JournaledState {
         let account = self.state.get_mut(&address).unwrap();
         // only if account is created in this tx we can assume that storage is empty.
         let is_newly_created = account.is_created();
-        let load = match account.storage.entry(key) {
-            Entry::Occupied(occ) => (occ.get().present_value, false),
+        let (value, is_cold) = match account.storage.entry(key) {
+            Entry::Occupied(occ) => {
+                let slot = occ.into_mut();
+                let is_cold = slot.mark_warm();
+                (slot.present_value, is_cold)
+            }
             Entry::Vacant(vac) => {
                 // if storage was cleared, we don't need to ping db.
                 let value = if is_newly_created {
                     U256::ZERO
                 } else {
                     db.storage(address, key).map_err(EVMError::Database)?
                 };
-                // add it to journal as cold loaded.
-                self.journal
-                    .last_mut()
-                    .unwrap()
-                    .push(JournalEntry::StorageChange {
-                        address,
-                        key,
-                        had_value: None,
-                    });
 
                 vac.insert(StorageSlot::new(value));
 
                 (value, true)
             }
         };
-        Ok(load)
+
+        if is_cold {
+            // add it to journal as cold loaded.
+            self.journal
+                .last_mut()
+                .unwrap()
+                .push(JournalEntry::StorageChange {
+                    address,
+                    key,
+                    had_value: None,
+                });
+        }
+
+        Ok((value, is_cold))
     }
 
     /// Stores storage slot.