feat(blockifier): aliases updater struct

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;

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<S>,
+    new_aliases: HashMap<AliasKey, Alias>,
+    next_free_alias: Alias,
+}
+
+impl<'a, S: StateReader> AliasUpdater<'a, S> {
+    fn new(state: &'a CachedState<S>) -> StateResult<Self> {
+        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<StorageEntry, Felt> {
+        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()
+    }
+}

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<StorageEntry, Felt>,
+    key: StorageKey,
+    value: Felt,
+) {
+    storage.insert((ALIAS_CONTRACT_ADDRESS, key), value);
+}
+
+fn initial_state(n_existing_aliases: u8) -> CachedState<DictStateReader> {
+    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<Felt>,
+    #[case] expected_alias_keys: Vec<Felt>,
+    #[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);
+}

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<ContractAddress> 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<u128> for PatriciaKey {