From 9ef99d427e6c97fc2240f82266c82d0e756993c8 Mon Sep 17 00:00:00 2001 From: Yoav Gross Date: Wed, 18 Dec 2024 14:34:25 +0200 Subject: [PATCH 1/2] feat(blockifier): replace the aliases in the state diff --- .../src/state/stateful_compression.rs | 31 ++++++++- .../src/state/stateful_compression_test.rs | 69 ++++++++++++++++++- 2 files changed, 97 insertions(+), 3 deletions(-) diff --git a/crates/blockifier/src/state/stateful_compression.rs b/crates/blockifier/src/state/stateful_compression.rs index 7c06cacdd2..5595d4ec85 100644 --- a/crates/blockifier/src/state/stateful_compression.rs +++ b/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( + state_diff: &StateMaps, + state: &S, + alias_contract_address: ContractAddress, +) -> CompressionResult { + 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, diff --git a/crates/blockifier/src/state/stateful_compression_test.rs b/crates/blockifier/src/state/stateful_compression_test.rs index bcd6461009..916a0e393c 100644 --- a/crates/blockifier/src/state/stateful_compression_test.rs +++ b/crates/blockifier/src/state/stateful_compression_test.rs @@ -3,18 +3,19 @@ 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::state::StorageKey; use starknet_types_core::felt::Felt; use super::{ + compress, state_diff_with_alias_allocation, AliasUpdater, ALIAS_COUNTER_STORAGE_KEY, INITIAL_AVAILABLE_ALIAS, MAX_NON_COMPRESSED_CONTRACT_ADDRESS, }; -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; @@ -240,3 +241,67 @@ 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); +} From 1644fd52ff98f9c4cce7b8b6c8499a9450c1ad34 Mon Sep 17 00:00:00 2001 From: Yoav Gross Date: Thu, 19 Dec 2024 16:36:05 +0200 Subject: [PATCH 2/2] test(blockifier): stateful compress and decompress --- .../src/state/stateful_compression_test.rs | 95 ++++++++++++++++++- 1 file changed, 94 insertions(+), 1 deletion(-) diff --git a/crates/blockifier/src/state/stateful_compression_test.rs b/crates/blockifier/src/state/stateful_compression_test.rs index 916a0e393c..72898f112b 100644 --- a/crates/blockifier/src/state/stateful_compression_test.rs +++ b/crates/blockifier/src/state/stateful_compression_test.rs @@ -1,19 +1,23 @@ -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, 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, StateMaps, StorageEntry}; use crate::state::state_api::{State, StateReader}; @@ -23,6 +27,90 @@ use crate::test_utils::dict_state_reader::DictStateReader; static ALIAS_CONTRACT_ADDRESS: LazyLock = LazyLock::new(|| ContractAddress(PatriciaKey::try_from(Felt::TWO).unwrap())); +/// Decompresses the state diff by replacing the aliases with addresses and storage keys. +fn decompress( + state_diff: &StateMaps, + state: &S, + alias_contract_address: ContractAddress, + alias_keys: HashSet, +) -> 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, +} + +impl AliasDecompressorUtil { + fn new( + state: &S, + alias_contract_address: ContractAddress, + alias_keys: HashSet, + ) -> 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, key: StorageKey, @@ -304,4 +392,9 @@ fn test_compression() { 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); }