Skip to content

Commit

Permalink
feat: Add call depth to vm.stopAndReturnStateDiff() results (foundr…
Browse files Browse the repository at this point in the history 
…y-rs#7234)

* Added depth member to AccountAccess

* Combine AccountAccess structs

---------

Co-authored-by: ercembu <ercembulut@gmail.com>
  • Loading branch information
@RPate97 @ercembu @zeroXbrock
2 people authored and zeroXbrock committed Feb 28, 2024
1 parent 8b18a13 commit ccb2331
Show file tree
Hide file tree
Showing 6 changed files with 132 additions and 114 deletions.
5 changes: 5 additions & 0 deletions crates/cheatcodes/assets/cheatcodes.json
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions crates/cheatcodes/spec/src/vm.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -203,6 +203,8 @@ interface Vm {
bool reverted;
/// An ordered list of storage accesses made during an account access operation.
StorageAccess[] storageAccesses;
/// Call depth traversed during the recording of state differences
uint64 depth;
}

/// The storage accessed during an `AccountAccess`.
Expand Down
1 change: 0 additions & 1 deletion crates/cheatcodes/src/evm.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -543,7 +543,6 @@ fn get_state_diff(state: &mut Cheatcodes) -> Result {
.unwrap_or_default()
.into_iter()
.flatten()
.map(|record| record.access)
.collect::<Vec<_>>();
Ok(res.abi_encode())
}
147 changes: 65 additions & 82 deletions crates/cheatcodes/src/inspector.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ use crate::{
ExpectedRevert, ExpectedRevertKind,
},
CheatsConfig, CheatsCtxt, Error, Result, Vm,
Vm::AccountAccess,
};
use alloy_primitives::{Address, Bytes, B256, U256, U64};
use alloy_rpc_types::request::{TransactionInput, TransactionRequest};
Expand Down Expand Up @@ -83,14 +84,6 @@ pub struct BroadcastableTransaction {
/// List of transactions that can be broadcasted.
pub type BroadcastableTransactions = VecDeque<BroadcastableTransaction>;

#[derive(Clone, Debug)]
pub struct AccountAccess {
/// The account access.
pub access: crate::Vm::AccountAccess,
/// The call depth the account was accessed.
pub depth: u64,
}

/// An EVM inspector that handles calls to various cheatcodes, each with their own behavior.
///
/// Cheatcodes can be called by contracts during execution to modify the VM environment, such as
Expand Down Expand Up @@ -434,10 +427,11 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
reverted: false,
deployedCode: vec![],
storageAccesses: vec![],
depth: data.journaled_state.depth(),
};
// Ensure that we're not selfdestructing a context recording was initiated on
if let Some(last) = account_accesses.last_mut() {
last.push(AccountAccess { access, depth: data.journaled_state.depth() });
last.push(access);
}
}
}
Expand Down Expand Up @@ -541,18 +535,14 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
reverted: false,
deployedCode: vec![],
storageAccesses: vec![],
};
let access = AccountAccess {
access: account_access,
// use current depth; EXT* opcodes are not creating new contexts
depth: data.journaled_state.depth(),
};
// Record the EXT* call as an account access at the current depth
// (future storage accesses will be recorded in a new "Resume" context)
if let Some(last) = recorded_account_diffs_stack.last_mut() {
last.push(access);
last.push(account_access);
} else {
recorded_account_diffs_stack.push(vec![access]);
recorded_account_diffs_stack.push(vec![account_access]);
}
}
_ => (),
Expand Down Expand Up @@ -883,23 +873,21 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
// updated with the revert status of this call, since the EVM does not mark accounts
// as "warm" if the call from which they were accessed is reverted
recorded_account_diffs_stack.push(vec![AccountAccess {
access: crate::Vm::AccountAccess {
chainInfo: crate::Vm::ChainInfo {
forkId: data.db.active_fork_id().unwrap_or_default(),
chainId: U256::from(data.env.cfg.chain_id),
},
accessor: call.context.caller,
account: call.contract,
kind,
initialized,
oldBalance: old_balance,
newBalance: U256::ZERO, // updated on call_end
value: call.transfer.value,
data: call.input.to_vec(),
reverted: false,
deployedCode: vec![],
storageAccesses: vec![], // updated on step
chainInfo: crate::Vm::ChainInfo {
forkId: data.db.active_fork_id().unwrap_or_default(),
chainId: U256::from(data.env.cfg.chain_id),
},
accessor: call.context.caller,
account: call.contract,
kind,
initialized,
oldBalance: old_balance,
newBalance: U256::ZERO, // updated on call_end
value: call.transfer.value,
data: call.input.to_vec(),
reverted: false,
deployedCode: vec![],
storageAccesses: vec![], // updated on step
depth: data.journaled_state.depth(),
}]);
}
Expand Down Expand Up @@ -1004,9 +992,8 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
// accordance with EVM behavior
if status.is_revert() {
last_recorded_depth.iter_mut().for_each(|element| {
element.access.reverted = true;
element.reverted = true;
element
.access
.storageAccesses
.iter_mut()
.for_each(|storage_access| storage_access.reverted = true);
Expand All @@ -1019,8 +1006,8 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
if call_access.depth == data.journaled_state.depth() {
if let Ok((acc, _)) = data.journaled_state.load_account(call.contract, data.db)
{
debug_assert!(access_is_call(call_access.access.kind));
call_access.access.newBalance = acc.info.balance;
debug_assert!(access_is_call(call_access.kind));
call_access.newBalance = acc.info.balance;
}
}
// Merge the last depth's AccountAccesses into the AccountAccesses at the current
Expand Down Expand Up @@ -1267,23 +1254,21 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
// Record the create context as an account access and create a new vector to record all
// subsequent account accesses
recorded_account_diffs_stack.push(vec![AccountAccess {
access: crate::Vm::AccountAccess {
chainInfo: crate::Vm::ChainInfo {
forkId: data.db.active_fork_id().unwrap_or_default(),
chainId: U256::from(data.env.cfg.chain_id),
},
accessor: call.caller,
account: address,
kind: crate::Vm::AccountAccessKind::Create,
initialized: true,
oldBalance: U256::ZERO, // updated on create_end
newBalance: U256::ZERO, // updated on create_end
value: call.value,
data: call.init_code.to_vec(),
reverted: false,
deployedCode: vec![], // updated on create_end
storageAccesses: vec![], // updated on create_end
chainInfo: crate::Vm::ChainInfo {
forkId: data.db.active_fork_id().unwrap_or_default(),
chainId: U256::from(data.env.cfg.chain_id),
},
accessor: call.caller,
account: address,
kind: crate::Vm::AccountAccessKind::Create,
initialized: true,
oldBalance: U256::ZERO, // updated on create_end
newBalance: U256::ZERO, // updated on create_end
value: call.value,
data: call.init_code.to_vec(),
reverted: false,
deployedCode: vec![], // updated on create_end
storageAccesses: vec![], // updated on create_end
depth,
}]);
}
Expand Down Expand Up @@ -1357,9 +1342,8 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
// accordance with EVM behavior
if status.is_revert() {
last_depth.iter_mut().for_each(|element| {
element.access.reverted = true;
element.reverted = true;
element
.access
.storageAccesses
.iter_mut()
.for_each(|storage_access| storage_access.reverted = true);
Expand All @@ -1372,15 +1356,15 @@ impl<DB: DatabaseExt> Inspector<DB> for Cheatcodes {
// percolated up to a higher depth.
if create_access.depth == data.journaled_state.depth() {
debug_assert_eq!(
create_access.access.kind as u8,
create_access.kind as u8,
crate::Vm::AccountAccessKind::Create as u8
);
if let Some(address) = address {
if let Ok((created_acc, _)) =
data.journaled_state.load_account(address, data.db)
{
create_access.access.newBalance = created_acc.info.balance;
create_access.access.deployedCode = created_acc
create_access.newBalance = created_acc.info.balance;
create_access.deployedCode = created_acc
.info
.code
.clone()
Expand Down Expand Up @@ -1464,23 +1448,21 @@ fn apply_create2_deployer<DB: DatabaseExt>(
if let Some(recorded_account_diffs_stack) = diffs_stack {
let calldata = [&salt.to_be_bytes::<32>()[..], &call.init_code[..]].concat();
recorded_account_diffs_stack.push(vec![AccountAccess {
access: crate::Vm::AccountAccess {
chainInfo: crate::Vm::ChainInfo {
forkId: data.db.active_fork_id().unwrap_or_default(),
chainId: U256::from(data.env.cfg.chain_id),
},
accessor: call.caller,
account: DEFAULT_CREATE2_DEPLOYER,
kind: crate::Vm::AccountAccessKind::Call,
initialized: true,
oldBalance: U256::ZERO, // updated on create_end
newBalance: U256::ZERO, // updated on create_end
value: call.value,
data: calldata,
reverted: false,
deployedCode: vec![], // updated on create_end
storageAccesses: vec![], // updated on create_end
chainInfo: crate::Vm::ChainInfo {
forkId: data.db.active_fork_id().unwrap_or_default(),
chainId: U256::from(data.env.cfg.chain_id),
},
accessor: call.caller,
account: DEFAULT_CREATE2_DEPLOYER,
kind: crate::Vm::AccountAccessKind::Call,
initialized: true,
oldBalance: U256::ZERO, // updated on create_end
newBalance: U256::ZERO, // updated on create_end
value: call.value,
data: calldata,
reverted: false,
deployedCode: vec![], // updated on create_end
storageAccesses: vec![], // updated on create_end
depth: data.journaled_state.depth(),
}])
}
Expand Down Expand Up @@ -1589,32 +1571,33 @@ fn append_storage_access(
// 2. If there's an existing Resume record, then add the storage access to it.
// 3. Otherwise, create a new Resume record based on the current context.
if last.len() == 1 {
last.first_mut().unwrap().access.storageAccesses.push(storage_access);
last.first_mut().unwrap().storageAccesses.push(storage_access);
} else {
let last_record = last.last_mut().unwrap();
if last_record.access.kind as u8 == crate::Vm::AccountAccessKind::Resume as u8 {
last_record.access.storageAccesses.push(storage_access);
if last_record.kind as u8 == crate::Vm::AccountAccessKind::Resume as u8 {
last_record.storageAccesses.push(storage_access);
} else {
let entry = last.first().unwrap();
let resume_record = crate::Vm::AccountAccess {
chainInfo: crate::Vm::ChainInfo {
forkId: entry.access.chainInfo.forkId,
chainId: entry.access.chainInfo.chainId,
forkId: entry.chainInfo.forkId,
chainId: entry.chainInfo.chainId,
},
accessor: entry.access.accessor,
account: entry.access.account,
accessor: entry.accessor,
account: entry.account,
kind: crate::Vm::AccountAccessKind::Resume,
initialized: entry.access.initialized,
initialized: entry.initialized,
storageAccesses: vec![storage_access],
reverted: entry.access.reverted,
reverted: entry.reverted,
// The remaining fields are defaults
oldBalance: U256::ZERO,
newBalance: U256::ZERO,
value: U256::ZERO,
data: vec![],
deployedCode: vec![],
depth: entry.depth,
};
last.push(AccountAccess { access: resume_record, depth: entry.depth });
last.push(resume_record);
}
}
}
Expand Down
Loading

0 comments on commit ccb2331

Please sign in to comment.