diff --git a/crates/blockifier/src/state/stateful_compression.rs b/crates/blockifier/src/state/stateful_compression.rs index fbbea66398f..68b4f5b2905 100644 --- a/crates/blockifier/src/state/stateful_compression.rs +++ b/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 = Result; + // 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 AliasCompressor<'_, S> { + fn compress_address( + &self, + contract_address: &ContractAddress, + ) -> CompressionResult { + 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, + contract_address: &ContractAddress, + ) -> CompressionResult { + if storage_key.0 >= MIN_VALUE_FOR_ALIAS_ALLOC + && contract_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 { + 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) } + } +} diff --git a/crates/blockifier/src/state/stateful_compression_test.rs b/crates/blockifier/src/state/stateful_compression_test.rs index 9cdd118cbe7..bcd64610096 100644 --- a/crates/blockifier/src/state/stateful_compression_test.rs +++ b/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 = @@ -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() + ); +}