feat(blockifier): aliases updater struct

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;

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

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<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 = 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<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 = 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);
+}