WIP: Implement prestateTracer (#3923)

Co-authored-by: Matthias Seitz <matthias.seitz@outlook.de>

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

@@ -4,8 +4,12 @@ use crate::tracing::{
     types::{CallTraceNode, CallTraceStepStackItem},
     TracingInspectorConfig,
 };
-use reth_primitives::{Address, Bytes, H256};
-use reth_rpc_types::trace::geth::*;
+use reth_primitives::{Address, Bytes, H256, U256};
+use reth_rpc_types::trace::geth::{
+    AccountState, CallConfig, CallFrame, DefaultFrame, DiffMode, GethDefaultTracingOptions,
+    PreStateConfig, PreStateFrame, PreStateMode, StructLog,
+};
+use revm::{db::DatabaseRef, primitives::ResultAndState};
 use std::collections::{BTreeMap, HashMap, VecDeque};
 
 /// A type for creating geth style traces
@@ -147,4 +151,75 @@ impl GethTraceBuilder {
             }
         }
     }
+
+    ///  Returns the accounts necessary for transaction execution.
+    ///
+    /// The prestate mode returns the accounts necessary to execute a given transaction.
+    /// diff_mode returns the differences between the transaction's pre and post-state.
+    ///
+    /// * `state` - The state post-transaction execution.
+    /// * `diff_mode` - if prestate is in diff or prestate mode.
+    /// * `db` - The database to fetch state pre-transaction execution.
+    pub fn geth_prestate_traces<DB>(
+        &self,
+        ResultAndState { state, .. }: &ResultAndState,
+        prestate_config: PreStateConfig,
+        db: DB,
+    ) -> Result<PreStateFrame, DB::Error>
+    where
+        DB: DatabaseRef,
+    {
+        let account_diffs: Vec<_> =
+            state.into_iter().map(|(addr, acc)| (*addr, &acc.info)).collect();
+
+        if prestate_config.is_diff_mode() {
+            let mut prestate = PreStateMode::default();
+            for (addr, _) in account_diffs {
+                let db_acc = db.basic(addr)?.unwrap_or_default();
+                prestate.0.insert(
+                    addr,
+                    AccountState {
+                        balance: Some(db_acc.balance),
+                        nonce: Some(U256::from(db_acc.nonce)),
+                        code: db_acc.code.as_ref().map(|code| Bytes::from(code.original_bytes())),
+                        storage: None,
+                    },
+                );
+            }
+            self.update_storage_from_trace(&mut prestate.0, false);
+            Ok(PreStateFrame::Default(prestate))
+        } else {
+            let mut state_diff = DiffMode::default();
+            for (addr, changed_acc) in account_diffs {
+                let db_acc = db.basic(addr)?.unwrap_or_default();
+                let pre_state = AccountState {
+                    balance: Some(db_acc.balance),
+                    nonce: Some(U256::from(db_acc.nonce)),
+                    code: db_acc.code.as_ref().map(|code| Bytes::from(code.original_bytes())),
+                    storage: None,
+                };
+                let post_state = AccountState {
+                    balance: Some(changed_acc.balance),
+                    nonce: Some(U256::from(changed_acc.nonce)),
+                    code: changed_acc.code.as_ref().map(|code| Bytes::from(code.original_bytes())),
+                    storage: None,
+                };
+                state_diff.pre.insert(addr, pre_state);
+                state_diff.post.insert(addr, post_state);
+            }
+            self.update_storage_from_trace(&mut state_diff.pre, false);
+            self.update_storage_from_trace(&mut state_diff.post, true);
+            Ok(PreStateFrame::Diff(state_diff))
+        }
+    }
+
+    fn update_storage_from_trace(
+        &self,
+        account_states: &mut BTreeMap<Address, AccountState>,
+        post_value: bool,
+    ) {
+        for node in self.nodes.iter() {
+            node.geth_update_account_storage(account_states, post_value);
+        }
+    }
 }

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

@@ -3,7 +3,7 @@
 use crate::tracing::{config::TraceStyle, utils::convert_memory};
 use reth_primitives::{abi::decode_revert_reason, bytes::Bytes, Address, H256, U256};
 use reth_rpc_types::trace::{
-    geth::{CallFrame, CallLogFrame, GethDefaultTracingOptions, StructLog},
+    geth::{AccountState, CallFrame, CallLogFrame, GethDefaultTracingOptions, StructLog},
     parity::{
         Action, ActionType, CallAction, CallOutput, CallType, ChangedType, CreateAction,
         CreateOutput, Delta, SelfdestructAction, StateDiff, TraceOutput, TransactionTrace,
@@ -13,7 +13,7 @@ use revm::interpreter::{
     opcode, CallContext, CallScheme, CreateScheme, InstructionResult, Memory, OpCode, Stack,
 };
 use serde::{Deserialize, Serialize};
-use std::collections::{btree_map::Entry, VecDeque};
+use std::collections::{btree_map::Entry, BTreeMap, VecDeque};
 
 /// A unified representation of a call
 #[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Serialize, Deserialize)]
@@ -443,6 +443,34 @@ impl CallTraceNode {
 
         call_frame
     }
+
+    /// Adds storage in-place to account state for all accounts that were touched in the trace
+    /// [CallTrace] execution.
+    ///
+    /// * `account_states` - the account map updated in place.
+    /// * `post_value` - if true, it adds storage values after trace transaction execution, if
+    ///   false, returns the storage values before trace execution.
+    pub(crate) fn geth_update_account_storage(
+        &self,
+        account_states: &mut BTreeMap<Address, AccountState>,
+        post_value: bool,
+    ) {
+        let addr = self.trace.address;
+        let acc_state = account_states.entry(addr).or_insert_with(AccountState::default);
+        for change in self.trace.steps.iter().filter_map(|s| s.storage_change) {
+            let StorageChange { key, value, had_value } = change;
+            let storage_map = acc_state.storage.get_or_insert_with(BTreeMap::new);
+            let value_to_insert = if post_value {
+                H256::from(value)
+            } else {
+                match had_value {
+                    Some(had_value) => H256::from(had_value),
+                    None => continue,
+                }
+            };
+            storage_map.insert(key.into(), value_to_insert);
+        }
+    }
 }
 
 pub(crate) struct CallTraceStepStackItem<'a> {