diff --git a/crates/blockifier/src/blockifier/transaction_executor.rs b/crates/blockifier/src/blockifier/transaction_executor.rs index 224fb95387..f48f8de11a 100644 --- a/crates/blockifier/src/blockifier/transaction_executor.rs +++ b/crates/blockifier/src/blockifier/transaction_executor.rs @@ -109,7 +109,7 @@ impl TransactionExecutor { 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, diff --git a/crates/blockifier/src/bouncer_test.rs b/crates/blockifier/src/bouncer_test.rs index 24ee237b9e..588be7b658 100644 --- a/crates/blockifier/src/bouncer_test.rs +++ b/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. diff --git a/crates/blockifier/src/concurrency/worker_logic.rs b/crates/blockifier/src/concurrency/worker_logic.rs index b50a011c4d..1699e5a4ab 100644 --- a/crates/blockifier/src/concurrency/worker_logic.rs +++ b/crates/blockifier/src/concurrency/worker_logic.rs @@ -233,7 +233,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; diff --git a/crates/blockifier/src/state/cached_state.rs b/crates/blockifier/src/state/cached_state.rs index 2272268112..1824c3d8a2 100644 --- a/crates/blockifier/src/state/cached_state.rs +++ b/crates/blockifier/src/state/cached_state.rs @@ -360,6 +360,28 @@ impl StateMaps { ), } } + + pub fn get_modified_contracts(&self) -> HashSet { + // Storage updates. + let mut modified_contracts: HashSet = + 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 { - // Storage updates. - let mut modified_contracts: HashSet = - 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, 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. diff --git a/crates/blockifier/src/state/cached_state_test.rs b/crates/blockifier/src/state/cached_state_test.rs index 620b969c60..9d32af747b 100644 --- a/crates/blockifier/src/state/cached_state_test.rs +++ b/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] diff --git a/crates/blockifier/src/transaction/account_transactions_test.rs b/crates/blockifier/src/transaction/account_transactions_test.rs index a5a1d8da17..a2b26b634b 100644 --- a/crates/blockifier/src/transaction/account_transactions_test.rs +++ b/crates/blockifier/src/transaction/account_transactions_test.rs @@ -1386,7 +1386,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); @@ -1423,7 +1423,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); @@ -1471,7 +1471,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); diff --git a/crates/blockifier/src/transaction/transaction_execution.rs b/crates/blockifier/src/transaction/transaction_execution.rs index e077370b86..abe3deec77 100644 --- a/crates/blockifier/src/transaction/transaction_execution.rs +++ b/crates/blockifier/src/transaction/transaction_execution.rs @@ -210,7 +210,7 @@ impl ExecutableTransaction 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(); - 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,