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>(

crates/rpc/rpc-types/src/eth/trace/geth/pre_state.rs

@@ -196,6 +196,11 @@ impl AccountChangeKind {
     }
 }
 
+/// The config for the prestate tracer.
+///
+/// If `diffMode` is set to true, the response frame includes all the account and storage diffs for
+/// the transaction. If it's missing or set to false it only returns the accounts and storage
+/// necessary to execute the transaction.
 #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct PreStateConfig {
@@ -204,9 +209,17 @@ pub struct PreStateConfig {
 }
 
 impl PreStateConfig {
+    /// Returns true if this trace was requested with diffmode.
+    #[inline]
     pub fn is_diff_mode(&self) -> bool {
         self.diff_mode.unwrap_or_default()
     }
+
+    /// Is default mode if diff_mode is not set
+    #[inline]
+    pub fn is_default_mode(&self) -> bool {
+        !self.is_diff_mode()
+    }
 }
 
 #[cfg(test)]

crates/rpc/rpc/src/debug.rs

@@ -263,7 +263,10 @@ where
                             .into_pre_state_config()
                             .map_err(|_| EthApiError::InvalidTracerConfig)?;
                         let mut inspector = TracingInspector::new(
-                            TracingInspectorConfig::from_geth_config(&config),
+                            TracingInspectorConfig::from_geth_config(&config)
+                                // if in default mode, we need to return all touched storages, for
+                                // which we need to record steps and statediff
+                                .set_steps_and_state_diffs(prestate_config.is_default_mode()),
                         );
 
                         let frame =
@@ -490,7 +493,10 @@ where
                             .map_err(|_| EthApiError::InvalidTracerConfig)?;
 
                         let mut inspector = TracingInspector::new(
-                            TracingInspectorConfig::from_geth_config(&config),
+                            TracingInspectorConfig::from_geth_config(&config)
+                                // if in default mode, we need to return all touched storages, for
+                                // which we need to record steps and statediff
+                                .set_steps_and_state_diffs(prestate_config.is_default_mode()),
                         );
                         let (res, _) = inspect(&mut *db, env, &mut inspector)?;