Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

bundle builder: add Option for state_original and state_present #1641

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 52 additions & 21 deletions crates/revm/src/db/states/bundle_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,8 +18,8 @@ use std::{
#[derive(Debug)]
pub struct BundleBuilder {
states: HashSet<Address>,
state_original: HashMap<Address, AccountInfo>,
state_present: HashMap<Address, AccountInfo>,
state_original: HashMap<Address, Option<AccountInfo>>,
state_present: HashMap<Address, Option<AccountInfo>>,
state_storage: HashMap<Address, HashMap<U256, (U256, U256)>>,

reverts: BTreeSet<(u64, Address)>,
Expand Down Expand Up @@ -103,13 +103,21 @@ impl BundleBuilder {
}

/// Collect account info of BundleState state
pub fn state_original_account_info(mut self, address: Address, original: AccountInfo) -> Self {
pub fn state_original_account_info(
mut self,
address: Address,
original: Option<AccountInfo>,
) -> Self {
self.set_state_original_account_info(address, original);
self
}

/// Collect account info of BundleState state
pub fn state_present_account_info(mut self, address: Address, present: AccountInfo) -> Self {
pub fn state_present_account_info(
mut self,
address: Address,
present: Option<AccountInfo>,
) -> Self {
self.set_state_present_account_info(address, present);
self
}
Expand Down Expand Up @@ -173,7 +181,7 @@ impl BundleBuilder {
pub fn set_state_original_account_info(
&mut self,
address: Address,
original: AccountInfo,
original: Option<AccountInfo>,
) -> &mut Self {
self.states.insert(address);
self.state_original.insert(address, original);
Expand All @@ -184,7 +192,7 @@ impl BundleBuilder {
pub fn set_state_present_account_info(
&mut self,
address: Address,
present: AccountInfo,
present: Option<AccountInfo>,
) -> &mut Self {
self.states.insert(address);
self.state_present.insert(address, present);
Expand Down Expand Up @@ -255,8 +263,8 @@ impl BundleBuilder {
})
.unwrap_or_default();
let bundle_account = BundleAccount::new(
self.state_original.remove(&address),
self.state_present.remove(&address),
self.state_original.remove(&address).unwrap_or_default(),
self.state_present.remove(&address).unwrap_or_default(),
storage,
AccountStatus::Changed,
);
Expand Down Expand Up @@ -327,12 +335,12 @@ impl BundleBuilder {
}

/// Mutable getter for `state_original` field
pub fn get_state_original_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
pub fn get_state_original_mut(&mut self) -> &mut HashMap<Address, Option<AccountInfo>> {
&mut self.state_original
}

/// Mutable getter for `state_present` field
pub fn get_state_present_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
pub fn get_state_present_mut(&mut self) -> &mut HashMap<Address, Option<AccountInfo>> {
&mut self.state_present
}

Expand Down Expand Up @@ -958,12 +966,12 @@ mod tests {
BundleState::builder(0..=0)
.state_present_account_info(
account1(),
AccountInfo {
Some(AccountInfo {
nonce: 1,
balance: U256::from(10),
code_hash: KECCAK_EMPTY,
code: None,
},
}),
)
.state_storage(
account1(),
Expand All @@ -972,12 +980,12 @@ mod tests {
.state_address(account2())
.state_present_account_info(
account2(),
AccountInfo {
Some(AccountInfo {
nonce: 1,
balance: U256::from(10),
code_hash: KECCAK_EMPTY,
code: None,
},
}),
)
.revert_address(0, account1())
.revert_account_info(0, account1(), Some(None))
Expand All @@ -991,12 +999,12 @@ mod tests {
BundleState::builder(0..=0)
.state_present_account_info(
account1(),
AccountInfo {
Some(AccountInfo {
nonce: 3,
balance: U256::from(20),
code_hash: KECCAK_EMPTY,
code: None,
},
}),
)
.state_storage(
account1(),
Expand Down Expand Up @@ -1213,12 +1221,12 @@ mod tests {
};

let present_state = BundleState::builder(2..=2)
.state_present_account_info(address1, account1_changed.clone())
.state_present_account_info(address1, Some(account1_changed.clone()))
.build();
assert_eq!(present_state.reverts.len(), 1);
let previous_state = BundleState::builder(1..=1)
.state_present_account_info(address1, account1)
.state_present_account_info(address2, account2.clone())
.state_present_account_info(address1, Some(account1))
.state_present_account_info(address2, Some(account2.clone()))
.build();
assert_eq!(previous_state.reverts.len(), 1);

Expand Down Expand Up @@ -1251,14 +1259,14 @@ mod tests {
assert!(builder.get_state_original_mut().is_empty());
builder
.get_state_original_mut()
.insert(account1(), AccountInfo::default());
.insert(account1(), Some(AccountInfo::default()));
assert!(builder.get_state_original_mut().contains_key(&account1()));

// Test get_state_present_mut
assert!(builder.get_state_present_mut().is_empty());
builder
.get_state_present_mut()
.insert(account1(), AccountInfo::default());
.insert(account1(), Some(AccountInfo::default()));
assert!(builder.get_state_present_mut().contains_key(&account1()));

// Test get_state_storage_mut
Expand Down Expand Up @@ -1301,4 +1309,27 @@ mod tests {
.insert(B256::default(), Bytecode::default());
assert!(builder.get_contracts_mut().contains_key(&B256::default()));
}

#[test]
fn destroyed_account() {
let address1 = account1();
let address2 = account2();

// Test to insert None as present state account info
let present_state = BundleState::builder(2..=2)
.state_present_account_info(address1, None)
.build();
assert!(present_state.state.get(&address1).unwrap().info.is_none());

// Test to insert None as original state account info
let original_state = BundleState::builder(2..=2)
.state_original_account_info(address2, None)
.build();
assert!(original_state
.state
.get(&address2)
.unwrap()
.original_info
.is_none());
}
}
Loading