fix: record touched accounts in prestate default mode (#5138)

crates/revm/revm-inspectors/src/tracing/builder/geth.rs

@@ -13,7 +13,7 @@ use revm::{
     db::DatabaseRef,
     primitives::{AccountInfo, ResultAndState, KECCAK_EMPTY},
 };
-use std::collections::{BTreeMap, HashMap, VecDeque};
+use std::collections::{btree_map::Entry, BTreeMap, HashMap, VecDeque};
 
 /// A type for creating geth style traces
 #[derive(Clone, Debug)]
@@ -208,23 +208,66 @@ impl GethTraceBuilder {
         };
 
         let account_diffs = state.into_iter().map(|(addr, acc)| (*addr, acc));
-        let is_diff = prestate_config.is_diff_mode();
-        if !is_diff {
+
+        if prestate_config.is_default_mode() {
             let mut prestate = PreStateMode::default();
-            for (addr, changed_acc) in account_diffs {
-                let db_acc = db.basic_ref(addr)?.unwrap_or_default();
-                let code = load_account_code(&db_acc);
+            // in default mode we __only__ return the touched state
+            for node in self.nodes.iter() {
+                let addr = node.trace.address;
 
-                let mut pre_state =
-                    AccountState::from_account_info(db_acc.nonce, db_acc.balance, code);
+                let acc_state = match prestate.0.entry(addr) {
+                    Entry::Vacant(entry) => {
+                        let db_acc = db.basic_ref(addr)?.unwrap_or_default();
+                        let code = load_account_code(&db_acc);
+                        let acc_state =
+                            AccountState::from_account_info(db_acc.nonce, db_acc.balance, code);
+                        entry.insert(acc_state)
+                    }
+                    Entry::Occupied(entry) => entry.into_mut(),
+                };
 
-                // handle _touched_ storage slots
-                for (key, slot) in changed_acc.storage.iter() {
-                    pre_state.storage.insert((*key).into(), slot.previous_or_original_value.into());
+                for (key, value) in node.touched_slots() {
+                    match acc_state.storage.entry(key.into()) {
+                        Entry::Vacant(entry) => {
+                            entry.insert(value.into());
+                        }
+                        Entry::Occupied(_) => {
+                            // we've already recorded this slot
+                        }
+                    }
                 }
-
-                prestate.0.insert(addr, pre_state);
             }
+
+            // also need to check changed accounts for things like balance changes etc
+            for (addr, changed_acc) in account_diffs {
+                let acc_state = match prestate.0.entry(addr) {
+                    Entry::Vacant(entry) => {
+                        let db_acc = db.basic_ref(addr)?.unwrap_or_default();
+                        let code = load_account_code(&db_acc);
+                        let acc_state =
+                            AccountState::from_account_info(db_acc.nonce, db_acc.balance, code);
+                        entry.insert(acc_state)
+                    }
+                    Entry::Occupied(entry) => {
+                        // already recorded via touched accounts
+                        entry.into_mut()
+                    }
+                };
+
+                // in case we missed anything during the trace, we need to add the changed accounts
+                // storage
+                for (key, slot) in changed_acc.storage.iter() {
+                    match acc_state.storage.entry((*key).into()) {
+                        Entry::Vacant(entry) => {
+                            entry.insert(slot.previous_or_original_value.into());
+                        }
+                        Entry::Occupied(_) => {
+                            // we've already recorded this slot
+                        }
+                    }
+                }
+            }
+
             Ok(PreStateFrame::Default(prestate))
         } else {
             let mut state_diff = DiffMode::default();

crates/revm/revm-inspectors/src/tracing/types.rs

@@ -14,7 +14,7 @@ use revm::interpreter::{
     opcode, CallContext, CallScheme, CreateScheme, InstructionResult, OpCode, SharedMemory, Stack,
 };
 use serde::{Deserialize, Serialize};
-use std::collections::VecDeque;
+use std::collections::{BTreeMap, VecDeque};
 
 /// A unified representation of a call
 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
@@ -250,6 +250,28 @@ impl CallTraceNode {
         }
     }
 
+    /// Returns all storage slots touched by this trace and the value this storage.
+    ///
+    /// A touched slot is either a slot that was written to or read from.
+    ///
+    /// If the slot is accessed more than once, the result only includes the first time it was
+    /// accessed, in other words in only returns the original value of the slot.
+    pub(crate) fn touched_slots(&self) -> BTreeMap<U256, U256> {
+        let mut touched_slots = BTreeMap::new();
+        for change in self.trace.steps.iter().filter_map(|s| s.storage_change.as_ref()) {
+            match touched_slots.entry(change.key) {
+                std::collections::btree_map::Entry::Vacant(entry) => {
+                    entry.insert(change.value);
+                }
+                std::collections::btree_map::Entry::Occupied(_) => {
+                    // already touched
+                }
+            }
+        }
+
+        touched_slots
+    }
+
     /// Pushes all steps onto the stack in reverse order
     /// so that the first step is on top of the stack
     pub(crate) fn push_steps_on_stack<'a>(