From cd76836af8211cea2a79b7e965d29b1f7a2d1f36 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 | 19 +++++++++++-------- .../src/state/stateful_compression_test.rs | 16 +++++++++++----- 3 files changed, 30 insertions(+), 14 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 e86649a07f..adc1e3f687 100644 --- a/crates/blockifier/src/state/stateful_compression.rs +++ b/crates/blockifier/src/state/stateful_compression.rs @@ -18,9 +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. -static ALIAS_CONTRACT_ADDRESS: LazyLock = - LazyLock::new(|| ContractAddress(PatriciaKey::try_from(Felt::TWO).unwrap())); // The storage key of the alias counter in the alias contract. static ALIAS_COUNTER_STORAGE_KEY: LazyLock = LazyLock::new(|| StorageKey(PatriciaKey::try_from(Felt::ZERO).unwrap())); @@ -39,6 +36,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; @@ -53,7 +51,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 address in addresses { if let Some(storage_keys) = sorted_storage_keys.get(&address) { for key in storage_keys { @@ -73,23 +71,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 { @@ -118,7 +121,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 2e0e502a20..f54f1a7fc8 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,6 +17,9 @@ 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, @@ -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,7 +128,8 @@ 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![ @@ -171,7 +176,8 @@ 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([(