feat(blockifier): replace the aliases in the state diff

crates/blockifier/src/state/stateful_compression.rs

@@ -138,8 +138,37 @@ impl<'a, S: StateReader> AliasUpdater<'a, S> {
     }
 }
 
+/// Compresses the state diff by replacing the addresses and storage keys with aliases.
+pub fn compress<S: StateReader>(
+    state_diff: &StateMaps,
+    state: &S,
+    alias_contract_address: ContractAddress,
+) -> CompressionResult<StateMaps> {
+    let alias_compressor = AliasCompressor { state, alias_contract_address };
+
+    let mut nonces = HashMap::new();
+    for (contract_address, nonce) in state_diff.nonces.iter() {
+        nonces.insert(alias_compressor.compress_address(contract_address)?, *nonce);
+    }
+    let mut class_hashes = HashMap::new();
+    for (contract_address, class_hash) in state_diff.class_hashes.iter() {
+        class_hashes.insert(alias_compressor.compress_address(contract_address)?, *class_hash);
+    }
+    let mut storage = HashMap::new();
+    for ((contract_address, key), value) in state_diff.storage.iter() {
+        storage.insert(
+            (
+                alias_compressor.compress_address(contract_address)?,
+                alias_compressor.compress_storage_key(key, contract_address)?,
+            ),
+            *value,
+        );
+    }
+
+    Ok(StateMaps { nonces, class_hashes, storage, ..state_diff.clone() })
+}
+
 /// Replaces contact addresses and storage keys with aliases.
-#[allow(dead_code)]
 struct AliasCompressor<'a, S: StateReader> {
     state: &'a S,
     alias_contract_address: ContractAddress,

crates/blockifier/src/state/stateful_compression_test.rs

@@ -1,27 +1,116 @@
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::LazyLock;
 
 use assert_matches::assert_matches;
 use rstest::rstest;
-use starknet_api::core::{ClassHash, ContractAddress, PatriciaKey};
+use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce, PatriciaKey};
+use starknet_api::felt;
 use starknet_api::state::StorageKey;
 use starknet_types_core::felt::Felt;
 
 use super::{
+    compress,
     state_diff_with_alias_allocation,
+    Alias,
+    AliasKey,
     AliasUpdater,
     ALIAS_COUNTER_STORAGE_KEY,
     INITIAL_AVAILABLE_ALIAS,
     MAX_NON_COMPRESSED_CONTRACT_ADDRESS,
+    MIN_VALUE_FOR_ALIAS_ALLOC,
 };
-use crate::state::cached_state::{CachedState, StorageEntry};
+use crate::state::cached_state::{CachedState, StateMaps, StorageEntry};
 use crate::state::state_api::{State, StateReader};
 use crate::state::stateful_compression::{AliasCompressor, CompressionError};
 use crate::test_utils::dict_state_reader::DictStateReader;
 
 static ALIAS_CONTRACT_ADDRESS: LazyLock<ContractAddress> =
     LazyLock::new(|| ContractAddress(PatriciaKey::try_from(Felt::TWO).unwrap()));
 
+/// Decompresses the state diff by replacing the aliases with addresses and storage keys.
+fn decompress<S: StateReader>(
+    state_diff: &StateMaps,
+    state: &S,
+    alias_contract_address: ContractAddress,
+    alias_keys: HashSet<AliasKey>,
+) -> StateMaps {
+    let alias_decompressor = AliasDecompressorUtil::new(state, alias_contract_address, alias_keys);
+
+    let mut nonces = HashMap::new();
+    for (alias_contract_address, nonce) in state_diff.nonces.iter() {
+        nonces.insert(alias_decompressor.decompress_address(alias_contract_address), *nonce);
+    }
+    let mut class_hashes = HashMap::new();
+    for (alias_contract_address, class_hash) in state_diff.class_hashes.iter() {
+        class_hashes
+            .insert(alias_decompressor.decompress_address(alias_contract_address), *class_hash);
+    }
+    let mut storage = HashMap::new();
+    for ((alias_contract_address, alias_storage_key), value) in state_diff.storage.iter() {
+        let contract_address = alias_decompressor.decompress_address(alias_contract_address);
+        storage.insert(
+            (
+                contract_address,
+                alias_decompressor.decompress_storage_key(alias_storage_key, &contract_address),
+            ),
+            *value,
+        );
+    }
+
+    StateMaps { nonces, class_hashes, storage, ..state_diff.clone() }
+}
+
+/// Replaces aliases with the original contact addresses and storage keys.
+struct AliasDecompressorUtil {
+    reversed_alias_mapping: HashMap<Alias, AliasKey>,
+}
+
+impl AliasDecompressorUtil {
+    fn new<S: StateReader>(
+        state: &S,
+        alias_contract_address: ContractAddress,
+        alias_keys: HashSet<AliasKey>,
+    ) -> Self {
+        let mut reversed_alias_mapping = HashMap::new();
+        for alias_key in alias_keys.into_iter() {
+            reversed_alias_mapping.insert(
+                state.get_storage_at(alias_contract_address, alias_key).unwrap(),
+                alias_key,
+            );
+        }
+        Self { reversed_alias_mapping }
+    }
+
+    fn decompress_address(&self, contract_address_alias: &ContractAddress) -> ContractAddress {
+        if contract_address_alias.0 >= MIN_VALUE_FOR_ALIAS_ALLOC {
+            ContractAddress::try_from(
+                *self.restore_alias_key(Felt::from(*contract_address_alias)).key(),
+            )
+            .unwrap()
+        } else {
+            *contract_address_alias
+        }
+    }
+
+    fn decompress_storage_key(
+        &self,
+        storage_key_alias: &StorageKey,
+        contact_address: &ContractAddress,
+    ) -> StorageKey {
+        if storage_key_alias.0 >= MIN_VALUE_FOR_ALIAS_ALLOC
+            && contact_address > &MAX_NON_COMPRESSED_CONTRACT_ADDRESS
+        {
+            self.restore_alias_key(*storage_key_alias.0)
+        } else {
+            *storage_key_alias
+        }
+    }
+
+    fn restore_alias_key(&self, alias: Alias) -> AliasKey {
+        *self.reversed_alias_mapping.get(&alias).unwrap()
+    }
+}
+
 fn insert_to_alias_contract(
     storage: &mut HashMap<StorageEntry, Felt>,
     key: StorageKey,
@@ -240,3 +329,72 @@ fn test_alias_compressor() {
         Err(CompressionError::MissedAlias(key)) if key == missed_key.into()
     );
 }
+
+#[test]
+fn test_compression() {
+    let state_reader = DictStateReader {
+        storage_view: (200_u16..206)
+            .map(|x| ((*ALIAS_CONTRACT_ADDRESS, StorageKey::from(x)), Felt::from(x + 100)))
+            .collect(),
+        ..Default::default()
+    };
+
+    // State diff with values that should not be compressed.
+    let base_state_diff = StateMaps {
+        nonces: vec![(ContractAddress::from(30_u16), Nonce(Felt::ONE))].into_iter().collect(),
+        class_hashes: vec![(ContractAddress::from(31_u16), ClassHash(Felt::ONE))]
+            .into_iter()
+            .collect(),
+        storage: vec![((ContractAddress::from(10_u16), StorageKey::from(205_u16)), Felt::TWO)]
+            .into_iter()
+            .collect(),
+        compiled_class_hashes: vec![(ClassHash(felt!("0x400")), CompiledClassHash(felt!("0x401")))]
+            .into_iter()
+            .collect(),
+        declared_contracts: vec![(ClassHash(felt!("0x402")), true)].into_iter().collect(),
+    };
+
+    let compressed_base_state_diff =
+        compress(&base_state_diff, &state_reader, *ALIAS_CONTRACT_ADDRESS).unwrap();
+    assert_eq!(compressed_base_state_diff, base_state_diff);
+
+    // Add to the state diff values that should be compressed.
+    let mut state_diff = base_state_diff.clone();
+    state_diff.extend(&StateMaps {
+        nonces: vec![(ContractAddress::from(200_u16), Nonce(Felt::ZERO))].into_iter().collect(),
+        class_hashes: vec![(ContractAddress::from(201_u16), ClassHash(Felt::ZERO))]
+            .into_iter()
+            .collect(),
+        storage: vec![
+            ((ContractAddress::from(202_u16), StorageKey::from(203_u16)), Felt::ZERO),
+            ((ContractAddress::from(32_u16), StorageKey::from(204_u16)), Felt::ONE),
+        ]
+        .into_iter()
+        .collect(),
+        ..Default::default()
+    });
+
+    let mut expected_compressed_state_diff = base_state_diff.clone();
+    expected_compressed_state_diff.extend(&StateMaps {
+        nonces: vec![(ContractAddress::from(300_u16), Nonce(Felt::ZERO))].into_iter().collect(),
+        class_hashes: vec![(ContractAddress::from(301_u16), ClassHash(Felt::ZERO))]
+            .into_iter()
+            .collect(),
+        storage: vec![
+            ((ContractAddress::from(302_u16), StorageKey::from(303_u16)), Felt::ZERO),
+            ((ContractAddress::from(32_u16), StorageKey::from(304_u16)), Felt::ONE),
+        ]
+        .into_iter()
+        .collect(),
+        ..Default::default()
+    });
+
+    let compressed_state_diff =
+        compress(&state_diff, &state_reader, *ALIAS_CONTRACT_ADDRESS).unwrap();
+    assert_eq!(compressed_state_diff, expected_compressed_state_diff);
+
+    let alias_keys = state_reader.storage_view.keys().map(|(_, key)| *key).collect();
+    let decompressed_state_diff =
+        decompress(&compressed_state_diff, &state_reader, *ALIAS_CONTRACT_ADDRESS, alias_keys);
+    assert_eq!(decompressed_state_diff, state_diff);
+}