From cc30c5909fb774f2bbb60a937113b9ccf01d7234 Mon Sep 17 00:00:00 2001 From: yoavGrs <97383386+yoavGrs@users.noreply.github.com> Date: Tue, 17 Dec 2024 17:38:19 +0200 Subject: [PATCH] feat(blockifier): aliases updater struct (#2534) fix(blockifier): fix test_write_at_validate_and_execute (#2529) feat(blockifier): aliases updater struct --- crates/blockifier/src/state.rs | 2 + .../src/state/stateful_compression.rs | 85 ++++++++++++++++ .../src/state/stateful_compression_test.rs | 98 +++++++++++++++++++ 3 files changed, 185 insertions(+) create mode 100644 crates/blockifier/src/state/stateful_compression.rs create mode 100644 crates/blockifier/src/state/stateful_compression_test.rs diff --git a/crates/blockifier/src/state.rs b/crates/blockifier/src/state.rs index 8aa857c963..b9f5b19289 100644 --- a/crates/blockifier/src/state.rs +++ b/crates/blockifier/src/state.rs @@ -6,3 +6,5 @@ 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 new file mode 100644 index 0000000000..9617817bf0 --- /dev/null +++ b/crates/blockifier/src/state/stateful_compression.rs @@ -0,0 +1,85 @@ +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: Option, +} + +impl<'a, S: StateReader> AliasUpdater<'a, S> { + fn new(state: &'a CachedState) -> StateResult { + let stored_counter = + 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) }, + }) + } + + /// 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) + { + let alias_to_allocate = match self.next_free_alias { + Some(alias) => alias, + None => INITIAL_AVAILABLE_ALIAS, + }; + self.new_aliases.insert(*alias_key, alias_to_allocate); + self.next_free_alias = Some(alias_to_allocate + Felt::ONE); + } + Ok(()) + } + + /// Inserts the counter of the alias contract. Returns the storage updates for the alias + /// contract. + fn finalize_updates(mut self) -> HashMap { + match self.next_free_alias { + None => { + self.new_aliases.insert(*ALIAS_COUNTER_STORAGE_KEY, INITIAL_AVAILABLE_ALIAS); + } + Some(alias) => { + if !self.new_aliases.is_empty() { + self.new_aliases.insert(*ALIAS_COUNTER_STORAGE_KEY, 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); +}