diff --git a/crates/revm/revm-inspectors/src/tracing/builder/geth.rs b/crates/revm/revm-inspectors/src/tracing/builder/geth.rs index d1ee600a36..d40cc032a7 100644 --- a/crates/revm/revm-inspectors/src/tracing/builder/geth.rs +++ b/crates/revm/revm-inspectors/src/tracing/builder/geth.rs @@ -6,10 +6,13 @@ use crate::tracing::{ }; use reth_primitives::{Address, Bytes, B256, U256}; use reth_rpc_types::trace::geth::{ - AccountState, CallConfig, CallFrame, DefaultFrame, DiffMode, GethDefaultTracingOptions, - PreStateConfig, PreStateFrame, PreStateMode, StructLog, + AccountChangeKind, AccountState, CallConfig, CallFrame, DefaultFrame, DiffMode, DiffStateKind, + GethDefaultTracingOptions, PreStateConfig, PreStateFrame, PreStateMode, StructLog, +}; +use revm::{ + db::DatabaseRef, + primitives::{ResultAndState, KECCAK_EMPTY}, }; -use revm::{db::DatabaseRef, primitives::ResultAndState}; use std::collections::{BTreeMap, HashMap, VecDeque}; /// A type for creating geth style traces @@ -185,10 +188,9 @@ impl GethTraceBuilder { where DB: DatabaseRef, { - let account_diffs: Vec<_> = - state.into_iter().map(|(addr, acc)| (*addr, &acc.info)).collect(); - - if !prestate_config.is_diff_mode() { + let account_diffs = state.into_iter().map(|(addr, acc)| (*addr, acc)); + let is_diff = prestate_config.is_diff_mode(); + if !is_diff { let mut prestate = PreStateMode::default(); for (addr, _) in account_diffs { let db_acc = db.basic(addr)?.unwrap_or_default(); @@ -202,42 +204,117 @@ impl GethTraceBuilder { ), ); } - self.update_storage_from_trace(&mut prestate.0, false); + self.update_storage_from_trace_prestate_mode(&mut prestate.0, DiffStateKind::Pre); Ok(PreStateFrame::Default(prestate)) } else { let mut state_diff = DiffMode::default(); + let mut change_types = HashMap::with_capacity(account_diffs.len()); for (addr, changed_acc) in account_diffs { let db_acc = db.basic(addr)?.unwrap_or_default(); - let pre_state = AccountState::from_account_info( - db_acc.nonce, - db_acc.balance, - db_acc.code.as_ref().map(|code| code.original_bytes()), - ); + let db_code = db_acc.code.as_ref(); + let db_code_hash = db_acc.code_hash; + + // Geth always includes the contract code in the prestate. However, + // the code hash will be KECCAK_EMPTY if the account is an EOA. Therefore + // we need to filter it out. + let pre_code = db_code + .map(|code| code.original_bytes()) + .or_else(|| { + if db_code_hash == KECCAK_EMPTY { + None + } else { + db.code_by_hash(db_code_hash).ok().map(|code| code.original_bytes()) + } + }) + .map(Into::into); + + let pre_state = + AccountState::from_account_info(db_acc.nonce, db_acc.balance, pre_code); + let post_state = AccountState::from_account_info( - changed_acc.nonce, - changed_acc.balance, - changed_acc.code.as_ref().map(|code| code.original_bytes()), + changed_acc.info.nonce, + changed_acc.info.balance, + changed_acc.info.code.as_ref().map(|code| code.original_bytes()), ); state_diff.pre.insert(addr, pre_state); state_diff.post.insert(addr, post_state); + + // determine the change type + let pre_change = if changed_acc.is_created() { + AccountChangeKind::Create + } else { + AccountChangeKind::Modify + }; + let post_change = if changed_acc.is_selfdestructed() { + AccountChangeKind::SelfDestruct + } else { + AccountChangeKind::Modify + }; + + change_types.insert(addr, (pre_change, post_change)); } - self.update_storage_from_trace(&mut state_diff.pre, false); - self.update_storage_from_trace(&mut state_diff.post, true); + + self.update_storage_from_trace_diff_mode(&mut state_diff.pre, DiffStateKind::Pre); + self.update_storage_from_trace_diff_mode(&mut state_diff.post, DiffStateKind::Post); // ensure we're only keeping changed entries - state_diff.retain_changed(); + state_diff.retain_changed().remove_zero_storage_values(); + self.diff_traces(&mut state_diff.pre, &mut state_diff.post, change_types); Ok(PreStateFrame::Diff(state_diff)) } } - fn update_storage_from_trace( + /// Updates the account storage for all nodes in the trace for pre-state mode. + #[inline] + fn update_storage_from_trace_prestate_mode( &self, account_states: &mut BTreeMap, - post_value: bool, + kind: DiffStateKind, ) { for node in self.nodes.iter() { - node.geth_update_account_storage(account_states, post_value); + node.geth_update_account_storage(account_states, kind); } } + + /// Updates the account storage for all nodes in the trace for diff mode. + #[inline] + fn update_storage_from_trace_diff_mode( + &self, + account_states: &mut BTreeMap, + kind: DiffStateKind, + ) { + for node in self.nodes.iter() { + node.geth_update_account_storage_diff_mode(account_states, kind); + } + } + + /// Returns the difference between the pre and post state of the transaction depending on the + /// kind of changes of that account (pre,post) + fn diff_traces( + &self, + pre: &mut BTreeMap, + post: &mut BTreeMap, + change_type: HashMap, + ) { + // Don't keep destroyed accounts in the post state + post.retain(|addr, post_state| { + // only keep accounts that are not created + if change_type.get(addr).map(|ty| !ty.1.is_selfdestruct()).unwrap_or(false) { + return false + } + if let Some(pre_state) = pre.get(addr) { + // remove any unchanged account info + post_state.remove_matching_account_info(pre_state); + } + + true + }); + + // Don't keep created accounts the pre state + pre.retain(|addr, _pre_state| { + // only keep accounts that are not created + change_type.get(addr).map(|ty| !ty.0.is_created()).unwrap_or(true) + }); + } } diff --git a/crates/revm/revm-inspectors/src/tracing/mod.rs b/crates/revm/revm-inspectors/src/tracing/mod.rs index f5b1794c93..51c5a1e2d3 100644 --- a/crates/revm/revm-inspectors/src/tracing/mod.rs +++ b/crates/revm/revm-inspectors/src/tracing/mod.rs @@ -24,7 +24,7 @@ mod types; mod utils; use crate::tracing::{ arena::PushTraceKind, - types::{CallTraceNode, StorageChange}, + types::{CallTraceNode, StorageChange, StorageChangeReason}, utils::gas_used, }; pub use builder::{ @@ -335,7 +335,6 @@ impl TracingInspector { step.memory.resize(interp.memory.len()); } } - if self.config.record_state_diff { let op = interp.current_opcode(); @@ -355,7 +354,12 @@ impl TracingInspector { ) => { // SAFETY: (Address,key) exists if part if StorageChange let value = data.journaled_state.state[address].storage[key].present_value(); - let change = StorageChange { key: *key, value, had_value: *had_value }; + let reason = match op { + opcode::SLOAD => StorageChangeReason::SLOAD, + opcode::SSTORE => StorageChangeReason::SSTORE, + _ => unreachable!(), + }; + let change = StorageChange { key: *key, value, had_value: *had_value, reason }; Some(change) } _ => None, diff --git a/crates/revm/revm-inspectors/src/tracing/types.rs b/crates/revm/revm-inspectors/src/tracing/types.rs index b5a607f434..8d2529029a 100644 --- a/crates/revm/revm-inspectors/src/tracing/types.rs +++ b/crates/revm/revm-inspectors/src/tracing/types.rs @@ -4,7 +4,9 @@ use crate::tracing::{config::TraceStyle, utils::convert_memory}; use alloy_sol_types::decode_revert_reason; use reth_primitives::{Address, Bytes, B256, U256, U64}; use reth_rpc_types::trace::{ - geth::{AccountState, CallFrame, CallLogFrame, GethDefaultTracingOptions, StructLog}, + geth::{ + AccountState, CallFrame, CallLogFrame, DiffStateKind, GethDefaultTracingOptions, StructLog, + }, parity::{ Action, ActionType, CallAction, CallOutput, CallType, ChangedType, CreateAction, CreateOutput, Delta, SelfdestructAction, StateDiff, TraceOutput, TransactionTrace, @@ -251,6 +253,16 @@ impl CallTraceNode { stack.extend(self.call_step_stack().into_iter().rev()); } + /// Returns all changed slots and the recorded changes + fn changed_storage_slots(&self) -> BTreeMap> { + let mut changed_slots: BTreeMap> = BTreeMap::new(); + for change in self.trace.steps.iter().filter_map(|s| s.storage_change) { + changed_slots.entry(change.key).or_default().push(change); + } + + changed_slots + } + /// Returns a list of all steps in this trace in the order they were executed /// /// If the step is a call, the id of the child trace is set. @@ -310,7 +322,7 @@ impl CallTraceNode { // iterate over all storage diffs for change in self.trace.steps.iter().filter_map(|s| s.storage_change) { - let StorageChange { key, value, had_value } = change; + let StorageChange { key, value, had_value, .. } = change; let b256_value = B256::from(value); match acc.storage.entry(key.into()) { Entry::Vacant(entry) => { @@ -507,18 +519,18 @@ impl CallTraceNode { /// [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. + /// * `kind` - if [DiffStateKind::Post], it adds storage values after trace transaction + /// execution, if [DiffStateKind::Pre], returns the storage values before trace execution. pub(crate) fn geth_update_account_storage( &self, account_states: &mut BTreeMap, - post_value: bool, + kind: DiffStateKind, ) { let addr = self.trace.address; let acc_state = account_states.entry(addr).or_default(); for change in self.trace.steps.iter().filter_map(|s| s.storage_change) { - let StorageChange { key, value, had_value } = change; - let value_to_insert = if post_value { + let StorageChange { key, value, had_value, .. } = change; + let value_to_insert = if kind.is_post() { B256::from(value) } else { match had_value { @@ -529,6 +541,57 @@ impl CallTraceNode { acc_state.storage.insert(key.into(), value_to_insert); } } + + /// Updates the account storage for all accounts that were touched in the trace. + /// + /// Depending on the [DiffStateKind] this will either insert the initial value + /// [DiffStateKind::Pre] or the final value [DiffStateKind::Post] of the storage slot. + pub(crate) fn geth_update_account_storage_diff_mode( + &self, + account_states: &mut BTreeMap, + kind: DiffStateKind, + ) { + let addr = self.execution_address(); + let changed_slots = self.changed_storage_slots(); + + // loop over all changed slots and track the storage changes of that slot + for (slot, changes) in changed_slots { + let account = account_states.entry(addr).or_default(); + + let mut initial_value = account.storage.get(&B256::from(slot)).copied().map(Into::into); + let mut final_value = None; + + for change in changes { + if initial_value.is_none() { + // set the initial value for the first storage change depending on the change + // reason + initial_value = match change.reason { + StorageChangeReason::SSTORE => Some(change.had_value.unwrap_or_default()), + StorageChangeReason::SLOAD => Some(change.value), + }; + } + + if change.reason == StorageChangeReason::SSTORE { + // keep track of the actual state value that's updated on sstore + final_value = Some(change.value); + } + } + + if final_value.is_none() || initial_value.is_none() { + continue + } + + if initial_value == final_value { + // unchanged + continue + } + + let value_to_write = + if kind.is_post() { final_value } else { initial_value }.expect("exists; qed"); + + account.storage.insert(B256::from(slot), B256::from(value_to_write)); + } + } } pub(crate) struct CallTraceStepStackItem<'a> { @@ -666,10 +729,20 @@ impl CallTraceStep { } } +/// Represents the source of a storage change - e.g., whether it came +/// from an SSTORE or SLOAD instruction. +#[allow(clippy::upper_case_acronyms)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) enum StorageChangeReason { + SLOAD, + SSTORE, +} + /// Represents a storage change during execution #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub(crate) struct StorageChange { pub(crate) key: U256, pub(crate) value: U256, pub(crate) had_value: Option, + pub(crate) reason: StorageChangeReason, } diff --git a/crates/rpc/rpc-types/src/eth/trace/geth/mod.rs b/crates/rpc/rpc-types/src/eth/trace/geth/mod.rs index 7e2c576f49..fe4019e9e8 100644 --- a/crates/rpc/rpc-types/src/eth/trace/geth/mod.rs +++ b/crates/rpc/rpc-types/src/eth/trace/geth/mod.rs @@ -11,7 +11,10 @@ pub use self::{ call::{CallConfig, CallFrame, CallLogFrame}, four_byte::FourByteFrame, noop::NoopFrame, - pre_state::{AccountState, DiffMode, PreStateConfig, PreStateFrame, PreStateMode}, + pre_state::{ + AccountChangeKind, AccountState, DiffMode, DiffStateKind, PreStateConfig, PreStateFrame, + PreStateMode, + }, }; mod call; diff --git a/crates/rpc/rpc-types/src/eth/trace/geth/pre_state.rs b/crates/rpc/rpc-types/src/eth/trace/geth/pre_state.rs index d867ebd0b8..d4d187fd9d 100644 --- a/crates/rpc/rpc-types/src/eth/trace/geth/pre_state.rs +++ b/crates/rpc/rpc-types/src/eth/trace/geth/pre_state.rs @@ -54,6 +54,9 @@ pub struct PreStateMode(pub BTreeMap); /// Represents the account states before and after the transaction is executed. /// /// This corresponds to the [DiffMode] of the [PreStateConfig]. +/// +/// This will only contain changed [AccountState]s, created accounts will not be included in the pre +/// state and selfdestructed accounts will not be included in the post state. #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct DiffMode { @@ -71,7 +74,7 @@ impl DiffMode { /// This will remove all unchanged [AccountState]s from the sets. /// /// In other words it removes entries that are equal (unchanged) in both the pre and post sets. - pub fn retain_changed(&mut self) { + pub fn retain_changed(&mut self) -> &mut Self { self.pre.retain(|address, pre| { if let btree_map::Entry::Occupied(entry) = self.post.entry(*address) { if entry.get() == pre { @@ -83,6 +86,38 @@ impl DiffMode { true }); + self + } + + /// Removes all zero values from the storage of the [AccountState]s. + pub fn remove_zero_storage_values(&mut self) { + self.pre.values_mut().for_each(|state| { + state.storage.retain(|_, value| *value != B256::ZERO); + }); + self.post.values_mut().for_each(|state| { + state.storage.retain(|_, value| *value != B256::ZERO); + }); + } +} + +/// Helper type for [DiffMode] to represent a specific set +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum DiffStateKind { + /// Corresponds to the pre state of the [DiffMode] + Pre, + /// Corresponds to the post state of the [DiffMode] + Post, +} + +impl DiffStateKind { + /// Returns true if this is the pre state of the [DiffMode] + pub fn is_pre(&self) -> bool { + matches!(self, DiffStateKind::Pre) + } + + /// Returns true if this is the post state of the [DiffMode] + pub fn is_post(&self) -> bool { + matches!(self, DiffStateKind::Post) } } @@ -117,6 +152,47 @@ impl AccountState { storage: Default::default(), } } + + /// Removes balance,nonce or code if they match the given account info. + /// + /// This is useful for comparing pre vs post state and only keep changed values in post state. + pub fn remove_matching_account_info(&mut self, other: &AccountState) { + if self.balance == other.balance { + self.balance = None; + } + if self.nonce == other.nonce { + self.nonce = None; + } + if self.code == other.code { + self.code = None; + } + } +} + +/// Helper type to track the kind of change of an [AccountState]. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, Serialize, Deserialize)] +pub enum AccountChangeKind { + #[default] + Modify, + Create, + SelfDestruct, +} + +impl AccountChangeKind { + /// Returns true if the account was created + pub fn is_created(&self) -> bool { + matches!(self, AccountChangeKind::Create) + } + + /// Returns true the account was modified + pub fn is_modified(&self) -> bool { + matches!(self, AccountChangeKind::Modify) + } + + /// Returns true the account was modified + pub fn is_selfdestruct(&self) -> bool { + matches!(self, AccountChangeKind::SelfDestruct) + } } #[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]