From 824e33fb29fdfa4d5e8f5eb616f6b3a53c6a0699 Mon Sep 17 00:00:00 2001 From: Joaquin Gonzalez Date: Fri, 14 Jun 2024 11:21:41 -0300 Subject: [PATCH 1/5] feat: mutable access to state_storage and revert_storage in BundleBuilder --- crates/revm/src/db/states/bundle_state.rs | 30 ++++++++++++++++------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/crates/revm/src/db/states/bundle_state.rs b/crates/revm/src/db/states/bundle_state.rs index 8dd4edfe26..70479bf40e 100644 --- a/crates/revm/src/db/states/bundle_state.rs +++ b/crates/revm/src/db/states/bundle_state.rs @@ -100,9 +100,17 @@ impl BundleBuilder { } /// Collect storage info of BundleState state - pub fn state_storage(mut self, address: Address, storage: HashMap) -> Self { + pub fn state_storage( + mut self, + address: Address, + storage: &mut HashMap, + ) -> Self { self.states.insert(address); - self.state_storage.insert(address, storage); + let entry = self + .state_storage + .entry(address) + .or_insert_with(HashMap::new); + entry.extend(storage.drain()); self } @@ -138,10 +146,14 @@ impl BundleBuilder { mut self, block_number: u64, address: Address, - storage: Vec<(U256, U256)>, + storage: &mut Vec<(U256, U256)>, ) -> Self { self.reverts.insert((block_number, address)); - self.revert_storage.insert((block_number, address), storage); + let entry = self + .revert_storage + .entry((block_number, address)) + .or_insert_with(Vec::new); + entry.extend(storage.drain(..)); self } @@ -832,7 +844,7 @@ mod tests { ) .state_storage( account1(), - HashMap::from([(slot1(), (U256::from(0), U256::from(10)))]), + &mut HashMap::from([(slot1(), (U256::from(0), U256::from(10)))]), ) .state_address(account2()) .state_present_account_info( @@ -846,7 +858,7 @@ mod tests { ) .revert_address(0, account1()) .revert_account_info(0, account1(), Some(None)) - .revert_storage(0, account1(), vec![(slot1(), U256::from(0))]) + .revert_storage(0, account1(), &mut vec![(slot1(), U256::from(0))]) .revert_account_info(0, account2(), Some(None)) .build() } @@ -865,7 +877,7 @@ mod tests { ) .state_storage( account1(), - HashMap::from([(slot1(), (U256::from(0), U256::from(15)))]), + &mut HashMap::from([(slot1(), (U256::from(0), U256::from(15)))]), ) .revert_address(0, account1()) .revert_account_info( @@ -878,7 +890,7 @@ mod tests { code: None, })), ) - .revert_storage(0, account1(), vec![(slot1(), U256::from(10))]) + .revert_storage(0, account1(), &mut vec![(slot1(), U256::from(10))]) .build() } @@ -1008,7 +1020,7 @@ mod tests { .revert_address(2, account2()) .revert_account_info(0, account1(), Some(None)) .revert_account_info(2, account2(), None) - .revert_storage(0, account1(), vec![(slot1(), U256::from(10))]) + .revert_storage(0, account1(), &mut vec![(slot1(), U256::from(10))]) .build(); assert_eq!(state.reverts.len(), 4); From 20e46eea13285d17096d54f120d27db041ffa90f Mon Sep 17 00:00:00 2001 From: Joaquin Gonzalez Date: Fri, 14 Jun 2024 11:37:00 -0300 Subject: [PATCH 2/5] self review --- crates/revm/src/db/states/bundle_state.rs | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/crates/revm/src/db/states/bundle_state.rs b/crates/revm/src/db/states/bundle_state.rs index 70479bf40e..6e66507133 100644 --- a/crates/revm/src/db/states/bundle_state.rs +++ b/crates/revm/src/db/states/bundle_state.rs @@ -106,10 +106,7 @@ impl BundleBuilder { storage: &mut HashMap, ) -> Self { self.states.insert(address); - let entry = self - .state_storage - .entry(address) - .or_insert_with(HashMap::new); + let entry = self.state_storage.entry(address).or_default(); entry.extend(storage.drain()); self } @@ -152,8 +149,8 @@ impl BundleBuilder { let entry = self .revert_storage .entry((block_number, address)) - .or_insert_with(Vec::new); - entry.extend(storage.drain(..)); + .or_default(); + entry.append(storage); self } From 5f1442b9e7fc78aa87b885c1b36aed0a9f5a2cf5 Mon Sep 17 00:00:00 2001 From: Joaquin Gonzalez Date: Sat, 15 Jun 2024 20:35:25 -0300 Subject: [PATCH 3/5] Revert "self review" This reverts commit 20e46eea13285d17096d54f120d27db041ffa90f. --- crates/revm/src/db/states/bundle_state.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/crates/revm/src/db/states/bundle_state.rs b/crates/revm/src/db/states/bundle_state.rs index 6e66507133..70479bf40e 100644 --- a/crates/revm/src/db/states/bundle_state.rs +++ b/crates/revm/src/db/states/bundle_state.rs @@ -106,7 +106,10 @@ impl BundleBuilder { storage: &mut HashMap, ) -> Self { self.states.insert(address); - let entry = self.state_storage.entry(address).or_default(); + let entry = self + .state_storage + .entry(address) + .or_insert_with(HashMap::new); entry.extend(storage.drain()); self } @@ -149,8 +152,8 @@ impl BundleBuilder { let entry = self .revert_storage .entry((block_number, address)) - .or_default(); - entry.append(storage); + .or_insert_with(Vec::new); + entry.extend(storage.drain(..)); self } From d49151434a7d77467400957ce482e255fb8bc3b0 Mon Sep 17 00:00:00 2001 From: Joaquin Gonzalez Date: Sat, 15 Jun 2024 20:35:37 -0300 Subject: [PATCH 4/5] Revert "feat: mutable access to state_storage and revert_storage in BundleBuilder" This reverts commit 824e33fb29fdfa4d5e8f5eb616f6b3a53c6a0699. --- crates/revm/src/db/states/bundle_state.rs | 30 +++++++---------------- 1 file changed, 9 insertions(+), 21 deletions(-) diff --git a/crates/revm/src/db/states/bundle_state.rs b/crates/revm/src/db/states/bundle_state.rs index 70479bf40e..8dd4edfe26 100644 --- a/crates/revm/src/db/states/bundle_state.rs +++ b/crates/revm/src/db/states/bundle_state.rs @@ -100,17 +100,9 @@ impl BundleBuilder { } /// Collect storage info of BundleState state - pub fn state_storage( - mut self, - address: Address, - storage: &mut HashMap, - ) -> Self { + pub fn state_storage(mut self, address: Address, storage: HashMap) -> Self { self.states.insert(address); - let entry = self - .state_storage - .entry(address) - .or_insert_with(HashMap::new); - entry.extend(storage.drain()); + self.state_storage.insert(address, storage); self } @@ -146,14 +138,10 @@ impl BundleBuilder { mut self, block_number: u64, address: Address, - storage: &mut Vec<(U256, U256)>, + storage: Vec<(U256, U256)>, ) -> Self { self.reverts.insert((block_number, address)); - let entry = self - .revert_storage - .entry((block_number, address)) - .or_insert_with(Vec::new); - entry.extend(storage.drain(..)); + self.revert_storage.insert((block_number, address), storage); self } @@ -844,7 +832,7 @@ mod tests { ) .state_storage( account1(), - &mut HashMap::from([(slot1(), (U256::from(0), U256::from(10)))]), + HashMap::from([(slot1(), (U256::from(0), U256::from(10)))]), ) .state_address(account2()) .state_present_account_info( @@ -858,7 +846,7 @@ mod tests { ) .revert_address(0, account1()) .revert_account_info(0, account1(), Some(None)) - .revert_storage(0, account1(), &mut vec![(slot1(), U256::from(0))]) + .revert_storage(0, account1(), vec![(slot1(), U256::from(0))]) .revert_account_info(0, account2(), Some(None)) .build() } @@ -877,7 +865,7 @@ mod tests { ) .state_storage( account1(), - &mut HashMap::from([(slot1(), (U256::from(0), U256::from(15)))]), + HashMap::from([(slot1(), (U256::from(0), U256::from(15)))]), ) .revert_address(0, account1()) .revert_account_info( @@ -890,7 +878,7 @@ mod tests { code: None, })), ) - .revert_storage(0, account1(), &mut vec![(slot1(), U256::from(10))]) + .revert_storage(0, account1(), vec![(slot1(), U256::from(10))]) .build() } @@ -1020,7 +1008,7 @@ mod tests { .revert_address(2, account2()) .revert_account_info(0, account1(), Some(None)) .revert_account_info(2, account2(), None) - .revert_storage(0, account1(), &mut vec![(slot1(), U256::from(10))]) + .revert_storage(0, account1(), vec![(slot1(), U256::from(10))]) .build(); assert_eq!(state.reverts.len(), 4); From f14f173421a92e374673c462ba948ec1d4128f69 Mon Sep 17 00:00:00 2001 From: Joaquin Gonzalez Date: Sat, 15 Jun 2024 20:48:01 -0300 Subject: [PATCH 5/5] feat: Add mutable getters for bundle state fields --- crates/revm/src/db/states/bundle_state.rs | 111 ++++++++++++++++++++++ 1 file changed, 111 insertions(+) diff --git a/crates/revm/src/db/states/bundle_state.rs b/crates/revm/src/db/states/bundle_state.rs index 8dd4edfe26..bd916b7ad6 100644 --- a/crates/revm/src/db/states/bundle_state.rs +++ b/crates/revm/src/db/states/bundle_state.rs @@ -233,6 +233,53 @@ impl BundleBuilder { pub fn get_states(&self) -> &HashSet
{ &self.states } + + /// Mutable getter for `states` field + pub fn get_states_mut(&mut self) -> &mut HashSet
{ + &mut self.states + } + + /// Mutable getter for `state_original` field + pub fn get_state_original_mut(&mut self) -> &mut HashMap { + &mut self.state_original + } + + /// Mutable getter for `state_present` field + pub fn get_state_present_mut(&mut self) -> &mut HashMap { + &mut self.state_present + } + + /// Mutable getter for `state_storage` field + pub fn get_state_storage_mut(&mut self) -> &mut HashMap> { + &mut self.state_storage + } + + /// Mutable getter for `reverts` field + pub fn get_reverts_mut(&mut self) -> &mut BTreeSet<(u64, Address)> { + &mut self.reverts + } + + /// Mutable getter for `revert_range` field + pub fn get_revert_range_mut(&mut self) -> &mut RangeInclusive { + &mut self.revert_range + } + + /// Mutable getter for `revert_account` field + pub fn get_revert_account_mut( + &mut self, + ) -> &mut HashMap<(u64, Address), Option>> { + &mut self.revert_account + } + + /// Mutable getter for `revert_storage` field + pub fn get_revert_storage_mut(&mut self) -> &mut HashMap<(u64, Address), Vec<(U256, U256)>> { + &mut self.revert_storage + } + + /// Mutable getter for `contracts` field + pub fn get_contracts_mut(&mut self) -> &mut HashMap { + &mut self.contracts + } } /// Bundle retention policy for applying substate to the bundle. @@ -1102,4 +1149,68 @@ mod tests { // account2 got inserted assert_eq!(test.state.get(&address2).unwrap().info, Some(account2)); } + + #[test] + fn test_getters() { + let mut builder = BundleBuilder::new(0..=3); + + // Test get_states and get_states_mut + assert!(builder.get_states().is_empty()); + builder.get_states_mut().insert(account1()); + assert!(builder.get_states().contains(&account1())); + + // Test get_state_original_mut + assert!(builder.get_state_original_mut().is_empty()); + builder + .get_state_original_mut() + .insert(account1(), 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()); + assert!(builder.get_state_present_mut().contains_key(&account1())); + + // Test get_state_storage_mut + assert!(builder.get_state_storage_mut().is_empty()); + builder + .get_state_storage_mut() + .insert(account1(), HashMap::new()); + assert!(builder.get_state_storage_mut().contains_key(&account1())); + + // Test get_reverts_mut + assert!(builder.get_reverts_mut().is_empty()); + builder.get_reverts_mut().insert((0, account1())); + assert!(builder.get_reverts_mut().contains(&(0, account1()))); + + // Test get_revert_range_mut + assert_eq!(builder.get_revert_range_mut().clone(), 0..=3); + + // Test get_revert_account_mut + assert!(builder.get_revert_account_mut().is_empty()); + builder + .get_revert_account_mut() + .insert((0, account1()), Some(None)); + assert!(builder + .get_revert_account_mut() + .contains_key(&(0, account1()))); + + // Test get_revert_storage_mut + assert!(builder.get_revert_storage_mut().is_empty()); + builder + .get_revert_storage_mut() + .insert((0, account1()), vec![(slot1(), U256::from(0))]); + assert!(builder + .get_revert_storage_mut() + .contains_key(&(0, account1()))); + + // Test get_contracts_mut + assert!(builder.get_contracts_mut().is_empty()); + builder + .get_contracts_mut() + .insert(B256::default(), Bytecode::default()); + assert!(builder.get_contracts_mut().contains_key(&B256::default())); + } }