diff --git a/Cargo.lock b/Cargo.lock index d57ab8f4f9..36e43a8734 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1492,6 +1492,7 @@ dependencies = [ "infra_utils", "itertools 0.12.1", "keccak", + "lazy_static", "log", "num-bigint 0.4.6", "num-integer", diff --git a/crates/blockifier/Cargo.toml b/crates/blockifier/Cargo.toml index b752a6531f..7cf47d61a0 100644 --- a/crates/blockifier/Cargo.toml +++ b/crates/blockifier/Cargo.toml @@ -35,6 +35,7 @@ indexmap.workspace = true infra_utils.workspace = true itertools.workspace = true keccak.workspace = true +lazy_static.workspace = true log.workspace = true num-bigint.workspace = true num-integer.workspace = true 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..077897bca3 --- /dev/null +++ b/crates/blockifier/src/state/stateful_compression.rs @@ -0,0 +1,83 @@ +use std::collections::HashMap; + +use lazy_static::lazy_static; +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"); + +lazy_static! { + // The address of the alias contract. + static ref ALIAS_CONTRACT_ADDRESS: ContractAddress = ContractAddress( + PatriciaKey::try_from(Felt::TWO).unwrap() + ); + // const ALIAS_CONTRACT_ADDRESS: ContractAddress = ContractAddress::new(Felt::TWO); + // The storage key of the alias counter in the alias contract. + static ref ALIAS_COUNTER_STORAGE_KEY: StorageKey = 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 ref MIN_VALUE_FOR_ALIAS_ALLOC: PatriciaKey = + 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); +}