diff --git a/crates/blockifier/src/blockifier/transaction_executor.rs b/crates/blockifier/src/blockifier/transaction_executor.rs index bec2bd5b8c..91d550850b 100644 --- a/crates/blockifier/src/blockifier/transaction_executor.rs +++ b/crates/blockifier/src/blockifier/transaction_executor.rs @@ -16,6 +16,7 @@ use crate::context::BlockContext; use crate::state::cached_state::{CachedState, CommitmentStateDiff, TransactionalState}; use crate::state::errors::StateError; use crate::state::state_api::{StateReader, StateResult}; +use crate::state::stateful_compression::state_diff_with_alias_allocation; use crate::transaction::errors::TransactionExecutionError; use crate::transaction::objects::TransactionExecutionInfo; use crate::transaction::transaction_execution::Transaction; @@ -166,16 +167,13 @@ impl TransactionExecutor { .collect::>()?; log::debug!("Final block weights: {:?}.", self.bouncer.get_accumulated_weights()); - Ok(( - self.block_state - .as_mut() - .expect(BLOCK_STATE_ACCESS_ERR) - .to_state_diff()? - .state_maps - .into(), - visited_segments, - *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)? + } else { + block_state.to_state_diff()?.state_maps + }; + Ok((state_diff.into(), visited_segments, *self.bouncer.get_accumulated_weights())) } } diff --git a/crates/blockifier/src/state.rs b/crates/blockifier/src/state.rs index b9f5b19289..179737b9f3 100644 --- a/crates/blockifier/src/state.rs +++ b/crates/blockifier/src/state.rs @@ -6,5 +6,4 @@ pub mod error_format_test; pub mod errors; pub mod global_cache; pub mod state_api; -#[allow(dead_code)] pub mod stateful_compression; diff --git a/crates/blockifier/src/state/stateful_compression.rs b/crates/blockifier/src/state/stateful_compression.rs index 9617817bf0..15f2bca8a6 100644 --- a/crates/blockifier/src/state/stateful_compression.rs +++ b/crates/blockifier/src/state/stateful_compression.rs @@ -1,11 +1,11 @@ -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::sync::LazyLock; use starknet_api::core::{ContractAddress, PatriciaKey}; use starknet_api::state::StorageKey; use starknet_types_core::felt::Felt; -use super::cached_state::{CachedState, StorageEntry}; +use super::cached_state::{CachedState, StateMaps, StorageEntry}; use super::state_api::{StateReader, StateResult}; #[cfg(test)] @@ -24,11 +24,49 @@ static ALIAS_CONTRACT_ADDRESS: LazyLock = // 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())); +// The maximal contract address for which aliases are not used and all keys are serialized as is, +// without compression. +static MAX_NON_COMPRESSED_CONTRACT_ADDRESS: LazyLock = + LazyLock::new(|| ContractAddress(PatriciaKey::try_from(Felt::from_hex_unchecked("0xf")).unwrap())); // The minimal value for a key to be allocated an alias. Smaller keys are serialized as is (their // alias is identical to the key). static MIN_VALUE_FOR_ALIAS_ALLOC: LazyLock = LazyLock::new(|| PatriciaKey::try_from(INITIAL_AVAILABLE_ALIAS).unwrap()); +/// Allocates aliases for the new addresses and storage keys in the alias contract. +/// Iterates over the addresses in ascending order. For each address, sets an alias for the new +/// storage keys (in ascending order) and for the address itself. +pub fn state_diff_with_alias_allocation( + state: &mut CachedState, +) -> StateResult { + let mut state_diff = state.to_state_diff()?.state_maps; + + // Collect the addresses and the storage keys that need aliases. + let addresses: BTreeSet = + state_diff.get_modified_contracts().into_iter().collect(); + let mut sorted_storage_keys = HashMap::new(); + for (address, storage_key) in state_diff.storage.keys() { + if address > &*MAX_NON_COMPRESSED_CONTRACT_ADDRESS { + sorted_storage_keys.entry(address).or_insert_with(BTreeSet::new).insert(storage_key); + } + } + + // Iterate over the addresses and the storage keys and update the aliases. + let mut alias_updater = AliasUpdater::new(state)?; + for address in addresses { + if let Some(storage_keys) = sorted_storage_keys.get(&address) { + for key in storage_keys { + alias_updater.insert_alias(key)?; + } + } + alias_updater.insert_alias(&StorageKey(address.0))?; + } + + let alias_storage_updates = alias_updater.finalize_updates(); + state_diff.storage.extend(alias_storage_updates); + Ok(state_diff) +} + /// Generate updates for the alias contract with the new keys. struct AliasUpdater<'a, S: StateReader> { state: &'a CachedState, diff --git a/crates/blockifier/src/state/stateful_compression_test.rs b/crates/blockifier/src/state/stateful_compression_test.rs index d0717221cd..2e0e502a20 100644 --- a/crates/blockifier/src/state/stateful_compression_test.rs +++ b/crates/blockifier/src/state/stateful_compression_test.rs @@ -1,16 +1,20 @@ use std::collections::HashMap; use rstest::rstest; +use starknet_api::core::{ClassHash, ContractAddress}; 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, }; use crate::state::cached_state::{CachedState, StorageEntry}; +use crate::state::state_api::{State, StateReader}; use crate::test_utils::dict_state_reader::DictStateReader; fn insert_to_alias_contract( @@ -96,3 +100,86 @@ fn test_alias_updater( assert_eq!(storage_diff, expected_storage_diff); } + +#[test] +fn test_iterate_aliases() { + let mut state = initial_state(0); + state + .set_storage_at(ContractAddress::from(0x201_u16), StorageKey::from(0x307_u16), Felt::ONE) + .unwrap(); + state + .set_storage_at(ContractAddress::from(0x201_u16), StorageKey::from(0x309_u16), Felt::TWO) + .unwrap(); + state + .set_storage_at(ContractAddress::from(0x201_u16), StorageKey::from(0x304_u16), Felt::THREE) + .unwrap(); + state + .set_storage_at( + *MAX_NON_COMPRESSED_CONTRACT_ADDRESS, + StorageKey::from(0x301_u16), + Felt::ONE, + ) + .unwrap(); + state.get_class_hash_at(ContractAddress::from(0x202_u16)).unwrap(); + 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; + assert_eq!( + storage_diff, + vec![ + ( + (*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(0x304_u16)), + INITIAL_AVAILABLE_ALIAS + Felt::ONE + ), + ( + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x307_u16)), + INITIAL_AVAILABLE_ALIAS + Felt::TWO + ), + ( + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x309_u16)), + INITIAL_AVAILABLE_ALIAS + Felt::THREE + ), + ( + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x201_u16)), + INITIAL_AVAILABLE_ALIAS + Felt::from(4_u8) + ), + ( + (*ALIAS_CONTRACT_ADDRESS, StorageKey::from(0x202_u16)), + INITIAL_AVAILABLE_ALIAS + Felt::from(5_u8) + ), + ((ContractAddress::from(0x201_u16), StorageKey::from(0x304_u16)), Felt::THREE), + ((ContractAddress::from(0x201_u16), StorageKey::from(0x307_u16)), Felt::ONE), + ((ContractAddress::from(0x201_u16), StorageKey::from(0x309_u16)), Felt::TWO), + ((*MAX_NON_COMPRESSED_CONTRACT_ADDRESS, StorageKey::from(0x301_u16)), Felt::ONE), + ] + .into_iter() + .collect() + ); +} + +#[rstest] +fn test_read_only_state(#[values(0, 2)] n_existing_aliases: u8) { + let mut state = initial_state(n_existing_aliases); + state + .set_storage_at(ContractAddress::from(0x200_u16), StorageKey::from(0x300_u16), Felt::ZERO) + .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 expected_storage_diff = if n_existing_aliases == 0 { + HashMap::from([( + (*ALIAS_CONTRACT_ADDRESS, *ALIAS_COUNTER_STORAGE_KEY), + INITIAL_AVAILABLE_ALIAS, + )]) + } else { + HashMap::new() + }; + assert_eq!(storage_diff, expected_storage_diff); +}