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..345c8f43fe --- /dev/null +++ b/crates/blockifier/src/state/stateful_compression.rs @@ -0,0 +1,78 @@ +use std::collections::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::state_api::{StateReader, StateResult}; + +#[cfg(test)] +#[path = "stateful_compression_test.rs"] +pub mod stateful_compression_test; + +type Alias = Felt; +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())); +// 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()); + +/// 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. + INITIAL_AVAILABLE_ALIAS + } 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 >= *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 == INITIAL_AVAILABLE_ALIAS { + 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..d0717221cd --- /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, + INITIAL_AVAILABLE_ALIAS, +}; +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 = INITIAL_AVAILABLE_ALIAS * Felt::TWO; + insert_to_alias_contract( + &mut state_reader.storage_view, + *ALIAS_COUNTER_STORAGE_KEY, + INITIAL_AVAILABLE_ALIAS + 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(), + INITIAL_AVAILABLE_ALIAS + 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![INITIAL_AVAILABLE_ALIAS - 1], vec![])] +#[case::single_update(vec![INITIAL_AVAILABLE_ALIAS], vec![INITIAL_AVAILABLE_ALIAS])] +#[case::some_update( + vec![ + INITIAL_AVAILABLE_ALIAS + 1, + INITIAL_AVAILABLE_ALIAS - 1, + INITIAL_AVAILABLE_ALIAS, + INITIAL_AVAILABLE_ALIAS + 2, + INITIAL_AVAILABLE_ALIAS, + ], + vec![ + INITIAL_AVAILABLE_ALIAS + 1, + INITIAL_AVAILABLE_ALIAS, + INITIAL_AVAILABLE_ALIAS + 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 = INITIAL_AVAILABLE_ALIAS + 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); +}