feat(blockifier): iterate over aliases on the state diff

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -16,6 +16,7 @@ use crate::context::BlockContext;
 use crate::state::cached_state::{CachedState, CommitmentStateDiff, TransactionalState};
 use crate::state::errors::StateError;
 use crate::state::state_api::{StateReader, StateResult};
+use crate::state::stateful_compression::allocate_aliases;
 use crate::transaction::errors::TransactionExecutionError;
 use crate::transaction::objects::TransactionExecutionInfo;
 use crate::transaction::transaction_execution::Transaction;
@@ -166,13 +167,12 @@ impl<S: StateReader> TransactionExecutor<S> {
             .collect::<TransactionExecutorResult<_>>()?;
 
         log::debug!("Final block weights: {:?}.", self.bouncer.get_accumulated_weights());
+        let mut block_state = self.block_state.take().expect(BLOCK_STATE_ACCESS_ERR);
+        if self.block_context.versioned_constants.enable_stateful_compression {
+            allocate_aliases(&mut block_state)?;
+        }
         Ok((
-            self.block_state
-                .as_mut()
-                .expect(BLOCK_STATE_ACCESS_ERR)
-                .to_state_diff()?
-                .state_maps
-                .into(),
+            block_state.to_state_diff()?.state_maps.into(),
             visited_segments,
             *self.bouncer.get_accumulated_weights(),
         ))

crates/blockifier/src/state/stateful_compression.rs

@@ -1,3 +1,5 @@
+use std::collections::{BTreeSet, HashMap};
+
 use starknet_api::core::{ContractAddress, PatriciaKey};
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
@@ -16,6 +18,9 @@ type AliasKey = StorageKey;
 const ALIAS_CONTRACT_ADDRESS: u8 = 2;
 // The storage key of the alias counter in the alias contract.
 const ALIAS_COUNTER_STORAGE_KEY: u8 = 0;
+// The maximal contract address for which aliases are not used and all keys are serialized as is,
+// without compression.
+const MAX_NON_COMPRESSED_CONTRACT_ADDRESS: u8 = 15;
 // 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");
@@ -26,6 +31,39 @@ pub fn get_alias_contract_address() -> ContractAddress {
 pub fn get_alias_counter_storage_key() -> StorageKey {
     StorageKey::from(ALIAS_COUNTER_STORAGE_KEY)
 }
+pub fn get_max_non_compressed_contract_address() -> ContractAddress {
+    ContractAddress::from(MAX_NON_COMPRESSED_CONTRACT_ADDRESS)
+}
+
+/// Allocates aliases for the new addresses and storage keys in the alias contract.
+/// Iterates over the addresses in ascending order. For each address, sets an alias for the new
+/// storage keys (in ascending order) and for the address itself.
+pub fn allocate_aliases<S: StateReader>(state: &mut CachedState<S>) -> StateResult<()> {
+    let writes = state.borrow_updated_state_cache()?.clone().writes;
+
+    // Collect the addresses and the storage keys that need aliases.
+    let mut addresses = BTreeSet::new();
+    let mut sorted_storage_keys = HashMap::new();
+    addresses.extend(writes.class_hashes.keys().chain(writes.nonces.keys()));
+    for (address, storage_key) in writes.storage.keys() {
+        addresses.insert(address);
+        if address > &get_max_non_compressed_contract_address() {
+            sorted_storage_keys.entry(address).or_insert_with(BTreeSet::new).insert(storage_key);
+        }
+    }
+
+    // Iterate over the addresses and the storage keys and update the aliases.
+    let mut alias_updater = AliasUpdater::new(state)?;
+    for address in addresses {
+        if let Some(storage_keys) = sorted_storage_keys.get(address) {
+            for key in storage_keys {
+                alias_updater.set_alias(key)?;
+            }
+        }
+        alias_updater.set_alias(&StorageKey(address.0))?;
+    }
+    alias_updater.finalize_updates()
+}
 
 /// Updates the alias contract with the new keys.
 struct AliasUpdater<'a, S: StateReader> {

crates/blockifier/src/state/stateful_compression_test.rs

@@ -1,17 +1,20 @@
 use std::collections::HashMap;
 
 use rstest::rstest;
-use starknet_api::core::ContractAddress;
+use starknet_api::core::{ClassHash, ContractAddress};
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
 
 use super::{
+    allocate_aliases,
     get_alias_contract_address,
     get_alias_counter_storage_key,
+    get_max_non_compressed_contract_address,
     AliasUpdater,
     MIN_VALUE_FOR_ALIAS_ALLOC,
 };
 use crate::state::cached_state::CachedState;
+use crate::state::state_api::{State, StateReader};
 use crate::test_utils::dict_state_reader::DictStateReader;
 
 fn insert_to_alias_contract(
@@ -98,3 +101,53 @@ fn test_alias_updater(
 
     assert_eq!(storage_diff, expected_storage_diff);
 }
+
+#[test]
+fn test_iterate_aliases() {
+    let mut state = initial_state(0);
+    state
+        .set_storage_at(ContractAddress::from(0x201_u16), StorageKey::from(0x300_u16), Felt::ONE)
+        .unwrap();
+    state
+        .set_storage_at(
+            get_max_non_compressed_contract_address(),
+            StorageKey::from(0x301_u16),
+            Felt::ONE,
+        )
+        .unwrap();
+    state.get_class_hash_at(ContractAddress::from(0x202_u16)).unwrap();
+    state.set_class_hash_at(ContractAddress::from(0x202_u16), ClassHash::default()).unwrap();
+    state.increment_nonce(ContractAddress::from(0x200_u16)).unwrap();
+
+    allocate_aliases(&mut state).unwrap();
+    let storage_diff = state.to_state_diff().unwrap().state_maps.storage;
+    assert_eq!(
+        storage_diff,
+        vec![
+            (
+                (get_alias_contract_address(), get_alias_counter_storage_key()),
+                MIN_VALUE_FOR_ALIAS_ALLOC + Felt::from(4_u8)
+            ),
+            (
+                (get_alias_contract_address(), StorageKey::from(0x200_u16)),
+                MIN_VALUE_FOR_ALIAS_ALLOC
+            ),
+            (
+                (get_alias_contract_address(), StorageKey::from(0x300_u16)),
+                MIN_VALUE_FOR_ALIAS_ALLOC + Felt::ONE
+            ),
+            (
+                (get_alias_contract_address(), StorageKey::from(0x201_u16)),
+                MIN_VALUE_FOR_ALIAS_ALLOC + Felt::TWO
+            ),
+            (
+                (get_alias_contract_address(), StorageKey::from(0x202_u16)),
+                MIN_VALUE_FOR_ALIAS_ALLOC + Felt::THREE
+            ),
+            ((ContractAddress::from(0x201_u16), StorageKey::from(0x300_u16)), Felt::ONE),
+            ((get_max_non_compressed_contract_address(), StorageKey::from(0x301_u16)), Felt::ONE),
+        ]
+        .into_iter()
+        .collect()
+    );
+}