From 8f4c153a02d94d9e1e9275ba3cf73f599e23c6b6 Mon Sep 17 00:00:00 2001
From: jpgonzalezra <jpgonzalezra@gmail.com>
Date: Sun, 16 Jun 2024 09:01:55 -0300
Subject: [PATCH] feat: mutable access for all fields in BundleBuilder (#1524)

* feat: mutable access to state_storage and revert_storage in BundleBuilder

* self review

* Revert "self review"

This reverts commit 20e46eea13285d17096d54f120d27db041ffa90f.

* Revert "feat: mutable access to state_storage and revert_storage in BundleBuilder"

This reverts commit 824e33fb29fdfa4d5e8f5eb616f6b3a53c6a0699.

* 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<Address> {
         &self.states
     }
+
+    /// Mutable getter for `states` field
+    pub fn get_states_mut(&mut self) -> &mut HashSet<Address> {
+        &mut self.states
+    }
+
+    /// Mutable getter for `state_original` field
+    pub fn get_state_original_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
+        &mut self.state_original
+    }
+
+    /// Mutable getter for `state_present` field
+    pub fn get_state_present_mut(&mut self) -> &mut HashMap<Address, AccountInfo> {
+        &mut self.state_present
+    }
+
+    /// Mutable getter for `state_storage` field
+    pub fn get_state_storage_mut(&mut self) -> &mut HashMap<Address, HashMap<U256, (U256, U256)>> {
+        &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<u64> {
+        &mut self.revert_range
+    }
+
+    /// Mutable getter for `revert_account` field
+    pub fn get_revert_account_mut(
+        &mut self,
+    ) -> &mut HashMap<(u64, Address), Option<Option<AccountInfo>>> {
+        &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<B256, Bytecode> {
+        &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()));
+    }
 }