feat(blockifier): define alias compressor struct

crates/blockifier/src/state/stateful_compression.rs

@@ -2,9 +2,12 @@ use std::collections::{BTreeSet, HashMap};
 
 use starknet_api::core::{ContractAddress, PatriciaKey};
 use starknet_api::state::StorageKey;
+use starknet_api::StarknetApiError;
 use starknet_types_core::felt::Felt;
+use thiserror::Error;
 
 use super::cached_state::{CachedState, StateMaps, StorageEntry};
+use super::errors::StateError;
 use super::state_api::{StateReader, StateResult};
 
 #[cfg(test)]
@@ -14,6 +17,17 @@ pub mod stateful_compression_test;
 type Alias = Felt;
 type AliasKey = StorageKey;
 
+#[derive(Debug, Error)]
+pub enum CompressionError {
+    #[error("Missing key in alias contract: {:#064x}", ***.0)]
+    MissedAlias(AliasKey),
+    #[error(transparent)]
+    StateError(#[from] StateError),
+    #[error(transparent)]
+    StarknetApiError(#[from] StarknetApiError),
+}
+pub type CompressionResult<T> = Result<T, CompressionError>;
+
 // The initial alias available for allocation.
 const INITIAL_AVAILABLE_ALIAS_HEX: &str = "0x80";
 const INITIAL_AVAILABLE_ALIAS: Felt = Felt::from_hex_unchecked(INITIAL_AVAILABLE_ALIAS_HEX);
@@ -123,3 +137,43 @@ impl<'a, S: StateReader> AliasUpdater<'a, S> {
             .collect()
     }
 }
+
+/// Replaces contact addresses and storage keys with aliases.
+#[allow(dead_code)]
+struct AliasCompressor<'a, S: StateReader> {
+    state: &'a S,
+    alias_contract_address: ContractAddress,
+}
+
+#[allow(dead_code)]
+impl<S: StateReader> AliasCompressor<'_, S> {
+    fn compress_address(
+        &self,
+        contract_address: &ContractAddress,
+    ) -> CompressionResult<ContractAddress> {
+        if contract_address.0 >= MIN_VALUE_FOR_ALIAS_ALLOC {
+            Ok(self.get_alias(StorageKey(contract_address.0))?.try_into()?)
+        } else {
+            Ok(*contract_address)
+        }
+    }
+
+    fn compress_storage_key(
+        &self,
+        storage_key: &StorageKey,
+        contact_address: &ContractAddress,
+    ) -> CompressionResult<StorageKey> {
+        if storage_key.0 >= MIN_VALUE_FOR_ALIAS_ALLOC
+            && contact_address > &MAX_NON_COMPRESSED_CONTRACT_ADDRESS
+        {
+            Ok(self.get_alias(*storage_key)?.try_into()?)
+        } else {
+            Ok(*storage_key)
+        }
+    }
+
+    fn get_alias(&self, alias_key: AliasKey) -> CompressionResult<Alias> {
+        let alias = self.state.get_storage_at(self.alias_contract_address, alias_key)?;
+        if alias == Felt::ZERO { Err(CompressionError::MissedAlias(alias_key)) } else { Ok(alias) }
+    }
+}

crates/blockifier/src/state/stateful_compression_test.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::sync::LazyLock;
 
+use assert_matches::assert_matches;
 use rstest::rstest;
 use starknet_api::core::{ClassHash, ContractAddress, PatriciaKey};
 use starknet_api::state::StorageKey;
@@ -15,6 +16,7 @@ use super::{
 };
 use crate::state::cached_state::{CachedState, 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> =
@@ -184,3 +186,57 @@ fn test_read_only_state(#[values(0, 2)] n_existing_aliases: u8) {
     };
     assert_eq!(storage_diff, expected_storage_diff);
 }
+
+/// Tests the range of alias keys that should be compressed.
+#[test]
+fn test_alias_compressor() {
+    let alias = Felt::from(500_u16);
+
+    let high_key = 200_u16;
+    let high_storage_key = StorageKey::from(high_key);
+    let high_contract_address = ContractAddress::from(high_key);
+
+    let no_aliasing_key = 50_u16;
+    let no_aliasing_storage_key = StorageKey::from(no_aliasing_key);
+    let no_aliasing_contract_address = ContractAddress::from(no_aliasing_key);
+
+    let no_compression_contract_address = ContractAddress::from(10_u16);
+
+    let mut state_reader = DictStateReader::default();
+    insert_to_alias_contract(&mut state_reader.storage_view, high_storage_key, alias);
+    let alias_compressor =
+        AliasCompressor { state: &state_reader, alias_contract_address: *ALIAS_CONTRACT_ADDRESS };
+
+    assert_eq!(
+        alias_compressor.compress_address(&high_contract_address).unwrap(),
+        ContractAddress::try_from(alias).unwrap(),
+    );
+    assert_eq!(
+        alias_compressor.compress_address(&no_aliasing_contract_address).unwrap(),
+        no_aliasing_contract_address,
+    );
+
+    assert_eq!(
+        alias_compressor.compress_storage_key(&high_storage_key, &high_contract_address).unwrap(),
+        StorageKey::try_from(alias).unwrap(),
+    );
+    assert_eq!(
+        alias_compressor
+            .compress_storage_key(&no_aliasing_storage_key, &high_contract_address)
+            .unwrap(),
+        no_aliasing_storage_key,
+    );
+    assert_eq!(
+        alias_compressor
+            .compress_storage_key(&high_storage_key, &no_compression_contract_address)
+            .unwrap(),
+        high_storage_key,
+    );
+
+    let missed_key = 300_u16;
+    let err = alias_compressor.compress_address(&ContractAddress::from(missed_key));
+    assert_matches!(
+        err,
+        Err(CompressionError::MissedAlias(key)) if key == missed_key.into()
+    );
+}