diff --git a/crates/blockifier/src/state.rs b/crates/blockifier/src/state.rs index 8aa857c963..179737b9f3 100644 --- a/crates/blockifier/src/state.rs +++ b/crates/blockifier/src/state.rs @@ -6,3 +6,4 @@ pub mod error_format_test; pub mod errors; pub mod global_cache; pub mod state_api; +pub mod stateful_compression; diff --git a/crates/blockifier/src/state/stateful_compression.rs b/crates/blockifier/src/state/stateful_compression.rs new file mode 100644 index 0000000000..d998aa14b9 --- /dev/null +++ b/crates/blockifier/src/state/stateful_compression.rs @@ -0,0 +1,71 @@ +use std::collections::HashMap; + +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::state_api::{StateReader, StateResult}; + +#[cfg(test)] +#[path = "stateful_compression_test.rs"] +pub mod stateful_compression_test; + +type Alias = Felt; +type AliasKey = StorageKey; + +// The address of the alias contract. +const ALIAS_CONTRACT_ADDRESS: ContractAddress = ContractAddress::new(Felt::TWO); +// The storage key of the alias counter in the alias contract. +const ALIAS_COUNTER_STORAGE_KEY: StorageKey = StorageKey(PatriciaKey::new_unchecked(Felt::ZERO)); +// The minimal value for a key to be allocated an alias. Smaller keys are serialized as is (their +// alias is identical to the key). +const MIN_VALUE_FOR_ALIAS_ALLOC: Felt = Felt::from_hex_unchecked("0x80"); + +/// Generate updates for the alias contract with the new keys. +struct AliasUpdater<'a, S: StateReader> { + state: &'a CachedState, + new_aliases: HashMap, + next_free_alias: Alias, +} + +impl<'a, S: StateReader> AliasUpdater<'a, S> { + fn new(state: &'a CachedState) -> StateResult { + let next_free_alias = + state.get_storage_at(ALIAS_CONTRACT_ADDRESS, ALIAS_COUNTER_STORAGE_KEY)?; + Ok(Self { + state, + new_aliases: HashMap::new(), + next_free_alias: if next_free_alias == Felt::ZERO { + // Aliasing first time. + MIN_VALUE_FOR_ALIAS_ALLOC + } else { + next_free_alias + }, + }) + } + + /// 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 >= PatriciaKey::try_from(MIN_VALUE_FOR_ALIAS_ALLOC)? + && self.state.get_storage_at(ALIAS_CONTRACT_ADDRESS, *alias_key)? == Felt::ZERO + && !self.new_aliases.contains_key(alias_key) + { + self.new_aliases.insert(*alias_key, self.next_free_alias); + self.next_free_alias += Felt::ONE; + } + Ok(()) + } + + /// Inserts the counter of the alias contract. Returns the storage updates for the alias + /// contract. + fn finalize_updates(mut self) -> HashMap { + if !self.new_aliases.is_empty() || self.next_free_alias == MIN_VALUE_FOR_ALIAS_ALLOC { + self.new_aliases.insert(ALIAS_COUNTER_STORAGE_KEY, self.next_free_alias); + } + self.new_aliases + .into_iter() + .map(|(key, alias)| ((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 new file mode 100644 index 0000000000..78ac556c53 --- /dev/null +++ b/crates/blockifier/src/state/stateful_compression_test.rs @@ -0,0 +1,98 @@ +use std::collections::HashMap; + +use rstest::rstest; +use starknet_api::state::StorageKey; +use starknet_types_core::felt::Felt; + +use super::{ + AliasUpdater, + ALIAS_CONTRACT_ADDRESS, + ALIAS_COUNTER_STORAGE_KEY, + MIN_VALUE_FOR_ALIAS_ALLOC, +}; +use crate::state::cached_state::{CachedState, StorageEntry}; +use crate::test_utils::dict_state_reader::DictStateReader; + +fn insert_to_alias_contract( + storage: &mut HashMap, + key: StorageKey, + value: Felt, +) { + storage.insert((ALIAS_CONTRACT_ADDRESS, key), value); +} + +fn initial_state(n_existing_aliases: u8) -> CachedState { + let mut state_reader = DictStateReader::default(); + if n_existing_aliases > 0 { + let high_alias_key = MIN_VALUE_FOR_ALIAS_ALLOC * Felt::TWO; + insert_to_alias_contract( + &mut state_reader.storage_view, + ALIAS_COUNTER_STORAGE_KEY, + MIN_VALUE_FOR_ALIAS_ALLOC + Felt::from(n_existing_aliases), + ); + for i in 0..n_existing_aliases { + insert_to_alias_contract( + &mut state_reader.storage_view, + (high_alias_key + Felt::from(i)).try_into().unwrap(), + MIN_VALUE_FOR_ALIAS_ALLOC + Felt::from(i), + ); + } + } + + CachedState::new(state_reader) +} + +/// Tests the alias contract updater with an empty state. +#[rstest] +#[case::no_update(vec![], vec![])] +#[case::low_update(vec![MIN_VALUE_FOR_ALIAS_ALLOC - 1], vec![])] +#[case::single_update(vec![MIN_VALUE_FOR_ALIAS_ALLOC], vec![MIN_VALUE_FOR_ALIAS_ALLOC])] +#[case::some_update( + vec![ + MIN_VALUE_FOR_ALIAS_ALLOC + 1, + MIN_VALUE_FOR_ALIAS_ALLOC - 1, + MIN_VALUE_FOR_ALIAS_ALLOC, + MIN_VALUE_FOR_ALIAS_ALLOC + 2, + MIN_VALUE_FOR_ALIAS_ALLOC, + ], + vec![ + MIN_VALUE_FOR_ALIAS_ALLOC + 1, + MIN_VALUE_FOR_ALIAS_ALLOC, + MIN_VALUE_FOR_ALIAS_ALLOC + 2, + ] +)] +fn test_alias_updater( + #[case] keys: Vec, + #[case] expected_alias_keys: Vec, + #[values(0, 2)] n_existing_aliases: u8, +) { + 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(); + for key in keys { + alias_contract_updater.insert_alias(&StorageKey::try_from(key).unwrap()).unwrap(); + } + let storage_diff = alias_contract_updater.finalize_updates(); + + // Test the new aliases. + let mut expected_storage_diff = HashMap::new(); + let mut expected_next_alias = MIN_VALUE_FOR_ALIAS_ALLOC + Felt::from(n_existing_aliases); + for key in &expected_alias_keys { + insert_to_alias_contract( + &mut expected_storage_diff, + StorageKey::try_from(*key).unwrap(), + expected_next_alias, + ); + expected_next_alias += Felt::ONE; + } + if !expected_alias_keys.is_empty() || n_existing_aliases == 0 { + insert_to_alias_contract( + &mut expected_storage_diff, + ALIAS_COUNTER_STORAGE_KEY, + expected_next_alias, + ); + } + + assert_eq!(storage_diff, expected_storage_diff); +} diff --git a/crates/starknet_api/src/core.rs b/crates/starknet_api/src/core.rs index 8b1a055462..deff31a066 100644 --- a/crates/starknet_api/src/core.rs +++ b/crates/starknet_api/src/core.rs @@ -121,6 +121,10 @@ impl ContractAddress { Err(StarknetApiError::OutOfRange { string: format!("[0x2, {})", l2_address_upper_bound) }) } + + pub const fn new(val: Felt) -> Self { + Self(PatriciaKey(val)) + } } impl From for Felt { @@ -357,6 +361,10 @@ impl PatriciaKey { pub fn key(&self) -> &StarkHash { &self.0 } + + pub const fn new_unchecked(val: StarkHash) -> Self { + Self(val) + } } impl From for PatriciaKey {