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

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

@@ -1,4 +1,4 @@
-use reth_primitives::{serde_helper::num::from_int_or_hex_opt, Address, H256, U256};
+use reth_primitives::{serde_helper::num::from_int_or_hex_opt, Address, Bytes, H256, U256};
 use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 
@@ -29,7 +29,7 @@ pub struct AccountState {
     )]
     pub balance: Option<U256>,
     #[serde(default, skip_serializing_if = "Option::is_none")]
-    pub code: Option<String>,
+    pub code: Option<Bytes>,
     #[serde(
         default,
         deserialize_with = "from_int_or_hex_opt",
@@ -47,6 +47,12 @@ pub struct PreStateConfig {
     pub diff_mode: Option<bool>,
 }
 
+impl PreStateConfig {
+    pub fn is_diff_mode(&self) -> bool {
+        self.diff_mode.unwrap_or_default()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -86,4 +92,11 @@ mod tests {
             _ => unreachable!(),
         }
     }
+
+    #[test]
+    fn test_is_diff_mode() {
+        assert!(PreStateConfig { diff_mode: Some(true) }.is_diff_mode());
+        assert!(!PreStateConfig { diff_mode: Some(false) }.is_diff_mode());
+        assert!(!PreStateConfig { diff_mode: None }.is_diff_mode());
+    }
 }

crates/rpc/rpc/src/debug.rs

@@ -2,7 +2,8 @@ use crate::{
     eth::{
         error::{EthApiError, EthResult},
         revm_utils::{
-            clone_into_empty_db, inspect, replay_transactions_until, result_output, EvmOverrides,
+            clone_into_empty_db, inspect, inspect_and_return_db, replay_transactions_until,
+            result_output, EvmOverrides,
         },
         EthTransactions, TransactionSource,
     },
@@ -255,7 +256,26 @@ where
                         return Ok(frame)
                     }
                     GethDebugBuiltInTracerType::PreStateTracer => {
-                        Err(EthApiError::Unsupported("pre state tracer currently unsupported."))
+                        let prestate_config = tracer_config
+                            .into_pre_state_config()
+                            .map_err(|_| EthApiError::InvalidTracerConfig)?;
+                        let mut inspector = TracingInspector::new(
+                            TracingInspectorConfig::from_geth_config(&config),
+                        );
+
+                        let frame =
+                            self.inner
+                                .eth_api
+                                .spawn_with_call_at(call, at, overrides, move |db, env| {
+                                    let (res, _, db) =
+                                        inspect_and_return_db(db, env, &mut inspector)?;
+                                    let frame = inspector
+                                        .into_geth_builder()
+                                        .geth_prestate_traces(&res, prestate_config, &db)?;
+                                    Ok(frame)
+                                })
+                                .await?;
+                        return Ok(frame.into())
                     }
                     GethDebugBuiltInTracerType::NoopTracer => Ok(NoopFrame::default().into()),
                 },
@@ -354,7 +374,22 @@ where
                         return Ok((frame.into(), res.state))
                     }
                     GethDebugBuiltInTracerType::PreStateTracer => {
-                        Err(EthApiError::Unsupported("prestate tracer is unimplemented yet."))
+                        let prestate_config = tracer_config
+                            .into_pre_state_config()
+                            .map_err(|_| EthApiError::InvalidTracerConfig)?;
+
+                        let mut inspector = TracingInspector::new(
+                            TracingInspectorConfig::from_geth_config(&config),
+                        );
+                        let (res, _) = inspect(&mut *db, env, &mut inspector)?;
+
+                        let frame = inspector.into_geth_builder().geth_prestate_traces(
+                            &res,
+                            prestate_config,
+                            &*db,
+                        )?;
+
+                        return Ok((frame.into(), res.state))
                     }
                     GethDebugBuiltInTracerType::NoopTracer => {
                         Ok((NoopFrame::default().into(), Default::default()))