From b90ba8d743711df2120e2939284a9c574a27ebd7 Mon Sep 17 00:00:00 2001 From: Yoav Gross Date: Sun, 15 Dec 2024 15:49:35 +0200 Subject: [PATCH] feat(blockifier): pass alias contract address --- .../src/blockifier/transaction_executor.rs | 9 ++++- .../src/state/stateful_compression.rs | 18 ++++++---- .../src/state/stateful_compression_test.rs | 34 +++++++++++-------- 3 files changed, 39 insertions(+), 22 deletions(-) diff --git a/crates/blockifier/src/blockifier/transaction_executor.rs b/crates/blockifier/src/blockifier/transaction_executor.rs index 91d550850b..5ef913aa4b 100644 --- a/crates/blockifier/src/blockifier/transaction_executor.rs +++ b/crates/blockifier/src/blockifier/transaction_executor.rs @@ -169,7 +169,14 @@ impl TransactionExecutor { log::debug!("Final block weights: {:?}.", self.bouncer.get_accumulated_weights()); let mut block_state = self.block_state.take().expect(BLOCK_STATE_ACCESS_ERR); let state_diff = if self.block_context.versioned_constants.enable_stateful_compression { - state_diff_with_alias_allocation(&mut block_state)? + state_diff_with_alias_allocation( + &mut block_state, + self.block_context + .versioned_constants + .os_constants + .os_contract_addresses + .alias_contract_address(), + )? } else { block_state.to_state_diff()?.state_maps }; diff --git a/crates/blockifier/src/state/stateful_compression.rs b/crates/blockifier/src/state/stateful_compression.rs index ed0017ebd1..fcfa10cbe9 100644 --- a/crates/blockifier/src/state/stateful_compression.rs +++ b/crates/blockifier/src/state/stateful_compression.rs @@ -18,8 +18,6 @@ type AliasKey = StorageKey; // The initial alias available for allocation. const INITIAL_AVAILABLE_ALIAS: Felt = Felt::from_hex_unchecked("0x80"); -// The address of the alias contract. -const ALIAS_CONTRACT_ADDRESS: ContractAddress = ContractAddress(PatriciaKey::TWO); // The storage key of the alias counter in the alias contract. const ALIAS_COUNTER_STORAGE_KEY: StorageKey = StorageKey(PatriciaKey::ZERO); // The maximal contract address for which aliases are not used and all keys are serialized as is, @@ -37,6 +35,7 @@ static MIN_VALUE_FOR_ALIAS_ALLOC: LazyLock = /// storage keys (in ascending order) and for the address itself. pub fn state_diff_with_alias_allocation( state: &mut CachedState, + alias_contract_address: ContractAddress, ) -> StateResult { let mut state_diff = state.to_state_diff()?.state_maps; @@ -54,7 +53,7 @@ pub fn state_diff_with_alias_allocation( } // Iterate over the addresses and the storage keys and update the aliases. - let mut alias_updater = AliasUpdater::new(state)?; + let mut alias_updater = AliasUpdater::new(state, alias_contract_address)?; for contract_address in contract_addresses { if let Some(storage_keys) = contract_address_to_sorted_storage_keys.get(&contract_address) { for key in storage_keys { @@ -74,23 +73,28 @@ struct AliasUpdater<'a, S: StateReader> { state: &'a CachedState, new_aliases: HashMap, next_free_alias: Option, + alias_contract_address: ContractAddress, } impl<'a, S: StateReader> AliasUpdater<'a, S> { - fn new(state: &'a CachedState) -> StateResult { + fn new( + state: &'a CachedState, + alias_contract_address: ContractAddress, + ) -> StateResult { let stored_counter = - state.get_storage_at(ALIAS_CONTRACT_ADDRESS, ALIAS_COUNTER_STORAGE_KEY)?; + state.get_storage_at(alias_contract_address, ALIAS_COUNTER_STORAGE_KEY)?; Ok(Self { state, new_aliases: HashMap::new(), next_free_alias: if stored_counter == Felt::ZERO { None } else { Some(stored_counter) }, + alias_contract_address, }) } /// Inserts the alias key to the updates if it's not already aliased. fn insert_alias(&mut self, alias_key: &AliasKey) -> StateResult<()> { if alias_key.0 >= *MIN_VALUE_FOR_ALIAS_ALLOC - && self.state.get_storage_at(ALIAS_CONTRACT_ADDRESS, *alias_key)? == Felt::ZERO + && self.state.get_storage_at(self.alias_contract_address, *alias_key)? == Felt::ZERO && !self.new_aliases.contains_key(alias_key) { let alias_to_allocate = match self.next_free_alias { @@ -119,7 +123,7 @@ impl<'a, S: StateReader> AliasUpdater<'a, S> { self.new_aliases .into_iter() - .map(|(key, alias)| ((ALIAS_CONTRACT_ADDRESS, key), alias)) + .map(|(key, alias)| ((self.alias_contract_address, key), alias)) .collect() } } diff --git a/crates/blockifier/src/state/stateful_compression_test.rs b/crates/blockifier/src/state/stateful_compression_test.rs index fd389a6a2c..bb9f334282 100644 --- a/crates/blockifier/src/state/stateful_compression_test.rs +++ b/crates/blockifier/src/state/stateful_compression_test.rs @@ -1,14 +1,14 @@ use std::collections::HashMap; +use std::sync::LazyLock; use rstest::rstest; -use starknet_api::core::{ClassHash, ContractAddress}; +use starknet_api::core::{ClassHash, ContractAddress, PatriciaKey}; use starknet_api::state::StorageKey; use starknet_types_core::felt::Felt; use super::{ state_diff_with_alias_allocation, AliasUpdater, - ALIAS_CONTRACT_ADDRESS, ALIAS_COUNTER_STORAGE_KEY, INITIAL_AVAILABLE_ALIAS, MAX_NON_COMPRESSED_CONTRACT_ADDRESS, @@ -17,12 +17,15 @@ use crate::state::cached_state::{CachedState, StorageEntry}; use crate::state::state_api::{State, StateReader}; use crate::test_utils::dict_state_reader::DictStateReader; +static ALIAS_CONTRACT_ADDRESS: LazyLock = + LazyLock::new(|| ContractAddress(PatriciaKey::try_from(Felt::TWO).unwrap())); + fn insert_to_alias_contract( storage: &mut HashMap, key: StorageKey, value: Felt, ) { - storage.insert((ALIAS_CONTRACT_ADDRESS, key), value); + storage.insert((*ALIAS_CONTRACT_ADDRESS, key), value); } fn initial_state(n_existing_aliases: u8) -> CachedState { @@ -73,7 +76,8 @@ fn test_alias_updater( let mut state = initial_state(n_existing_aliases); // Insert the keys into the alias contract updater and finalize the updates. - let mut alias_contract_updater = AliasUpdater::new(&mut state).unwrap(); + let mut alias_contract_updater = + AliasUpdater::new(&mut state, *ALIAS_CONTRACT_ADDRESS).unwrap(); for key in keys { alias_contract_updater.insert_alias(&StorageKey::try_from(key).unwrap()).unwrap(); } @@ -124,33 +128,34 @@ fn test_iterate_aliases() { state.set_class_hash_at(ContractAddress::from(0x202_u16), ClassHash(Felt::ONE)).unwrap(); state.increment_nonce(ContractAddress::from(0x200_u16)).unwrap(); - let storage_diff = state_diff_with_alias_allocation(&mut state).unwrap().storage; + let storage_diff = + state_diff_with_alias_allocation(&mut state, *ALIAS_CONTRACT_ADDRESS).unwrap().storage; assert_eq!( storage_diff, vec![ ( - (ALIAS_CONTRACT_ADDRESS, ALIAS_COUNTER_STORAGE_KEY), + (*ALIAS_CONTRACT_ADDRESS, ALIAS_COUNTER_STORAGE_KEY), INITIAL_AVAILABLE_ALIAS + Felt::from(6_u8) ), - ((ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x200_u16)), INITIAL_AVAILABLE_ALIAS), + ((*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x200_u16)), INITIAL_AVAILABLE_ALIAS), ( - (ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x304_u16)), + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x304_u16)), INITIAL_AVAILABLE_ALIAS + Felt::ONE ), ( - (ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x307_u16)), + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x307_u16)), INITIAL_AVAILABLE_ALIAS + Felt::TWO ), ( - (ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x309_u16)), + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x309_u16)), INITIAL_AVAILABLE_ALIAS + Felt::THREE ), ( - (ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x201_u16)), + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x201_u16)), INITIAL_AVAILABLE_ALIAS + Felt::from(4_u8) ), ( - (ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x202_u16)), + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x202_u16)), INITIAL_AVAILABLE_ALIAS + Felt::from(5_u8) ), ((ContractAddress::from(0x201_u16), StorageKey::from(0x304_u16)), Felt::THREE), @@ -171,11 +176,12 @@ fn test_read_only_state(#[values(0, 2)] n_existing_aliases: u8) { .unwrap(); state.get_nonce_at(ContractAddress::from(0x201_u16)).unwrap(); state.get_class_hash_at(ContractAddress::from(0x202_u16)).unwrap(); - let storage_diff = state_diff_with_alias_allocation(&mut state).unwrap().storage; + let storage_diff = + state_diff_with_alias_allocation(&mut state, *ALIAS_CONTRACT_ADDRESS).unwrap().storage; let expected_storage_diff = if n_existing_aliases == 0 { HashMap::from([( - (ALIAS_CONTRACT_ADDRESS, ALIAS_COUNTER_STORAGE_KEY), + (*ALIAS_CONTRACT_ADDRESS, ALIAS_COUNTER_STORAGE_KEY), INITIAL_AVAILABLE_ALIAS, )]) } else {