refactor(blockifier): move functions to state maps

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -109,7 +109,7 @@ impl<S: StateReader> TransactionExecutor<S> {
         match tx_execution_result {
             Ok(tx_execution_info) => {
                 let tx_state_changes_keys =
-                    transactional_state.get_actual_state_changes()?.into_keys();
+                    transactional_state.get_actual_state_changes()?.state_maps.into_keys();
                 self.bouncer.try_update(
                     &transactional_state,
                     &tx_state_changes_keys,

crates/blockifier/src/bouncer_test.rs

@@ -248,7 +248,8 @@ fn test_bouncer_try_update(#[case] added_ecdsa: usize, #[case] scenario: &'stati
         },
         ..Default::default()
     };
-    let tx_state_changes_keys = transactional_state.get_actual_state_changes().unwrap().into_keys();
+    let tx_state_changes_keys =
+        transactional_state.get_actual_state_changes().unwrap().state_maps.into_keys();
 
     // TODO(Yoni, 1/10/2024): simplify this test and move tx-too-large cases out.
 

crates/blockifier/src/concurrency/worker_logic.rs

@@ -15,12 +15,7 @@ use crate::concurrency::utils::lock_mutex_in_array;
 use crate::concurrency::versioned_state::ThreadSafeVersionedState;
 use crate::concurrency::TxIndex;
 use crate::context::BlockContext;
-use crate::state::cached_state::{
-    ContractClassMapping,
-    StateChanges,
-    StateMaps,
-    TransactionalState,
-};
+use crate::state::cached_state::{ContractClassMapping, StateMaps, TransactionalState};
 use crate::state::state_api::{StateReader, UpdatableState};
 use crate::transaction::objects::{
     TransactionExecutionInfo,
@@ -233,7 +228,7 @@ impl<'a, S: StateReader> WorkerExecutor<'a, S> {
         let mut execution_output = lock_mutex_in_array(&self.execution_outputs, tx_index);
         let writes = &execution_output.as_ref().expect(EXECUTION_OUTPUTS_UNWRAP_ERROR).writes;
         // TODO(Yoni): get rid of this clone.
-        let mut tx_state_changes_keys = StateChanges { state_maps: writes.clone() }.into_keys();
+        let mut tx_state_changes_keys = writes.clone().into_keys();
         let tx_result =
             &mut execution_output.as_mut().expect(EXECUTION_OUTPUTS_UNWRAP_ERROR).result;
 

crates/blockifier/src/state/cached_state.rs

@@ -360,6 +360,28 @@ impl StateMaps {
             ),
         }
     }
+
+    pub fn get_modified_contracts(&self) -> HashSet<ContractAddress> {
+        // Storage updates.
+        let mut modified_contracts: HashSet<ContractAddress> =
+            self.storage.keys().map(|address_key_pair| address_key_pair.0).collect();
+        // Nonce updates.
+        modified_contracts.extend(self.nonces.keys());
+        // Class hash updates (deployed contracts + replace_class syscall).
+        modified_contracts.extend(self.class_hashes.keys());
+
+        modified_contracts
+    }
+
+    pub fn into_keys(self) -> StateChangesKeys {
+        StateChangesKeys {
+            modified_contracts: self.get_modified_contracts(),
+            nonce_keys: self.nonces.into_keys().collect(),
+            class_hash_keys: self.class_hashes.into_keys().collect(),
+            storage_keys: self.storage.into_keys().collect(),
+            compiled_class_hash_keys: self.compiled_class_hashes.into_keys().collect(),
+        }
+    }
 }
 /// Caches read and write requests.
 /// The tracked changes are needed for block state commitment.
@@ -673,24 +695,12 @@ impl StateChanges {
         merged_state_changes
     }
 
-    pub fn get_modified_contracts(&self) -> HashSet<ContractAddress> {
-        // Storage updates.
-        let mut modified_contracts: HashSet<ContractAddress> =
-            self.state_maps.storage.keys().map(|address_key_pair| address_key_pair.0).collect();
-        // Nonce updates.
-        modified_contracts.extend(self.state_maps.nonces.keys());
-        // Class hash updates (deployed contracts + replace_class syscall).
-        modified_contracts.extend(self.state_maps.class_hashes.keys());
-
-        modified_contracts
-    }
-
     pub fn count_for_fee_charge(
         &self,
         sender_address: Option<ContractAddress>,
         fee_token_address: ContractAddress,
     ) -> StateChangesCount {
-        let mut modified_contracts = self.get_modified_contracts();
+        let mut modified_contracts = self.state_maps.get_modified_contracts();
 
         // For account transactions, we need to compute the transaction fee before we can execute
         // the fee transfer, and the fee should cover the state changes that happen in the
@@ -716,16 +726,6 @@ impl StateChanges {
             n_modified_contracts: modified_contracts.len(),
         }
     }
-
-    pub fn into_keys(self) -> StateChangesKeys {
-        StateChangesKeys {
-            modified_contracts: self.get_modified_contracts(),
-            nonce_keys: self.state_maps.nonces.into_keys().collect(),
-            class_hash_keys: self.state_maps.class_hashes.into_keys().collect(),
-            storage_keys: self.state_maps.storage.into_keys().collect(),
-            compiled_class_hash_keys: self.state_maps.compiled_class_hashes.into_keys().collect(),
-        }
-    }
 }
 
 /// Holds the number of state changes.

crates/blockifier/src/state/cached_state_test.rs

@@ -435,7 +435,7 @@ fn test_contract_cache_is_used() {
 #[test]
 fn test_cache_get_write_keys() {
     // Trivial case.
-    assert_eq!(StateChanges::default().into_keys(), StateChangesKeys::default());
+    assert_eq!(StateMaps::default().into_keys(), StateChangesKeys::default());
 
     // Interesting case.
     let some_felt = felt!("0x1");
@@ -448,21 +448,19 @@ fn test_cache_get_write_keys() {
 
     let class_hash0 = class_hash!("0x300");
 
-    let state_changes = StateChanges {
-        state_maps: StateMaps {
-            nonces: HashMap::from([(contract_address0, Nonce(some_felt))]),
-            class_hashes: HashMap::from([
-                (contract_address1, some_class_hash),
-                (contract_address2, some_class_hash),
-            ]),
-            storage: HashMap::from([
-                ((contract_address1, storage_key!(0x300_u16)), some_felt),
-                ((contract_address1, storage_key!(0x600_u16)), some_felt),
-                ((contract_address3, storage_key!(0x600_u16)), some_felt),
-            ]),
-            compiled_class_hashes: HashMap::from([(class_hash0, compiled_class_hash!(0x3_u16))]),
-            declared_contracts: HashMap::default(),
-        },
+    let state_maps = StateMaps {
+        nonces: HashMap::from([(contract_address0, Nonce(some_felt))]),
+        class_hashes: HashMap::from([
+            (contract_address1, some_class_hash),
+            (contract_address2, some_class_hash),
+        ]),
+        storage: HashMap::from([
+            ((contract_address1, storage_key!(0x300_u16)), some_felt),
+            ((contract_address1, storage_key!(0x600_u16)), some_felt),
+            ((contract_address3, storage_key!(0x600_u16)), some_felt),
+        ]),
+        compiled_class_hashes: HashMap::from([(class_hash0, compiled_class_hash!(0x3_u16))]),
+        declared_contracts: HashMap::default(),
     };
 
     let expected_keys = StateChangesKeys {
@@ -482,7 +480,7 @@ fn test_cache_get_write_keys() {
         ]),
     };
 
-    assert_eq!(state_changes.into_keys(), expected_keys);
+    assert_eq!(state_maps.into_keys(), expected_keys);
 }
 
 #[test]

crates/blockifier/src/transaction/account_transactions_test.rs

@@ -1395,7 +1395,7 @@ fn test_count_actual_storage_changes(
         ..Default::default()
     };
 
-    assert_eq!(expected_modified_contracts, state_changes_1.get_modified_contracts());
+    assert_eq!(expected_modified_contracts, state_changes_1.state_maps.get_modified_contracts());
     assert_eq!(expected_storage_updates_1, state_changes_1.state_maps.storage);
     assert_eq!(state_changes_count_1, expected_state_changes_count_1);
 
@@ -1432,7 +1432,7 @@ fn test_count_actual_storage_changes(
         ..Default::default()
     };
 
-    assert_eq!(expected_modified_contracts_2, state_changes_2.get_modified_contracts());
+    assert_eq!(expected_modified_contracts_2, state_changes_2.state_maps.get_modified_contracts());
     assert_eq!(expected_storage_updates_2, state_changes_2.state_maps.storage);
     assert_eq!(state_changes_count_2, expected_state_changes_count_2);
 
@@ -1480,7 +1480,7 @@ fn test_count_actual_storage_changes(
 
     assert_eq!(
         expected_modified_contracts_transfer,
-        state_changes_transfer.get_modified_contracts()
+        state_changes_transfer.state_maps.get_modified_contracts()
     );
     assert_eq!(expected_storage_update_transfer, state_changes_transfer.state_maps.storage);
     assert_eq!(state_changes_count_3, expected_state_changes_count_3);

crates/blockifier/src/transaction/transaction_execution.rs

@@ -210,7 +210,7 @@ impl<U: UpdatableState> ExecutableTransaction<U> for Transaction {
         // Check if the transaction is too large to fit any block.
         // TODO(Yoni, 1/8/2024): consider caching these two.
         let tx_execution_summary = tx_execution_info.summarize(&block_context.versioned_constants);
-        let mut tx_state_changes_keys = state.get_actual_state_changes()?.into_keys();
+        let mut tx_state_changes_keys = state.get_actual_state_changes()?.state_maps.into_keys();
         tx_state_changes_keys.update_sequencer_key_in_storage(
             &block_context.to_tx_context(self),
             &tx_execution_info,