From 428f370da49048542db8b7765d4fede76db76a17 Mon Sep 17 00:00:00 2001 From: yoavGrs <97383386+yoavGrs@users.noreply.github.com> Date: Wed, 20 Nov 2024 16:56:36 +0200 Subject: [PATCH] refactor(blockifier): move functions to state maps (#2173) --- .../src/blockifier/transaction_executor.rs | 2 +- crates/blockifier/src/bouncer_test.rs | 3 +- .../src/concurrency/worker_logic.rs | 9 +--- crates/blockifier/src/state/cached_state.rs | 46 +++++++++---------- .../blockifier/src/state/cached_state_test.rs | 32 ++++++------- .../transaction/account_transactions_test.rs | 6 +-- .../src/transaction/transaction_execution.rs | 2 +- 7 files changed, 47 insertions(+), 53 deletions(-) diff --git a/crates/blockifier/src/blockifier/transaction_executor.rs b/crates/blockifier/src/blockifier/transaction_executor.rs index 16dcefb6b9..63887ff9d2 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 b672b2c9a1..d24d33223d 100644 --- a/crates/blockifier/src/concurrency/worker_logic.rs +++ b/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; diff --git a/crates/blockifier/src/state/cached_state.rs b/crates/blockifier/src/state/cached_state.rs index 2edb4e7dc3..52b8435eb5 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 fefe41ac0b..52e040fcde 100644 --- a/crates/blockifier/src/transaction/account_transactions_test.rs +++ b/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); diff --git a/crates/blockifier/src/transaction/transaction_execution.rs b/crates/blockifier/src/transaction/transaction_execution.rs index dcf0cf0e41..bea665412e 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(&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,