From ff60390026e21934b912b584dc58fa97c27757ec Mon Sep 17 00:00:00 2001 From: Yoav Gross Date: Wed, 18 Dec 2024 14:34:25 +0200 Subject: [PATCH 1/4] feat(blockifier): define alias compressor struct --- .../src/state/stateful_compression.rs | 53 ++++++++++++++++++ .../src/state/stateful_compression_test.rs | 56 +++++++++++++++++++ 2 files changed, 109 insertions(+) diff --git a/crates/blockifier/src/state/stateful_compression.rs b/crates/blockifier/src/state/stateful_compression.rs index 073546278b..da1d1099df 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,42 @@ 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, +} + +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, + contact_address: &ContractAddress, + ) -> CompressionResult { + 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 { + 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 9cdd118cbe..bcd6461009 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() + ); +} From 73a571cc473b9c1f38b3bc8c1a4df4b1524e4979 Mon Sep 17 00:00:00 2001 From: Yoav Gross Date: Wed, 18 Dec 2024 14:34:25 +0200 Subject: [PATCH 2/4] 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 da1d1099df..3022d272dc 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 3a1c42978356a17d0505f36ae73a95945be5d5c1 Mon Sep 17 00:00:00 2001 From: Yoav Gross Date: Thu, 19 Dec 2024 16:36:05 +0200 Subject: [PATCH 3/4] 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..e9a4a06bce 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); } From a33587dc747b7feba86786f661e4211fc7c658a2 Mon Sep 17 00:00:00 2001 From: Yoav Gross Date: Wed, 18 Dec 2024 17:06:52 +0200 Subject: [PATCH 4/4] feat(blockifier): return compressed state diff at finalizing --- .../src/blockifier/transaction_executor.rs | 49 +++++++++++++------ .../src/state_reader/utils.rs | 2 +- .../src/py_block_executor.rs | 2 +- 3 files changed, 36 insertions(+), 17 deletions(-) diff --git a/crates/blockifier/src/blockifier/transaction_executor.rs b/crates/blockifier/src/blockifier/transaction_executor.rs index 5ef913aa4b..a547fe217e 100644 --- a/crates/blockifier/src/blockifier/transaction_executor.rs +++ b/crates/blockifier/src/blockifier/transaction_executor.rs @@ -16,7 +16,11 @@ 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::state_diff_with_alias_allocation; +use crate::state::stateful_compression::{ + compress, + state_diff_with_alias_allocation, + CompressionError, +}; use crate::transaction::errors::TransactionExecutionError; use crate::transaction::objects::TransactionExecutionInfo; use crate::transaction::transaction_execution::Transaction; @@ -36,6 +40,8 @@ pub enum TransactionExecutorError { StateError(#[from] StateError), #[error(transparent)] TransactionExecutionError(#[from] TransactionExecutionError), + #[error(transparent)] + CompressionError(#[from] CompressionError), } pub type TransactionExecutorResult = Result; @@ -141,12 +147,16 @@ impl TransactionExecutor { results } - /// Returns the state diff, a list of contract class hash with the corresponding list of - /// visited segment values and the block weights. + /// Returns the state diff, the compressed state diff, a list of contract class hash with the + /// corresponding list of visited segment values and the block weights. pub fn finalize( &mut self, - ) -> TransactionExecutorResult<(CommitmentStateDiff, VisitedSegmentsMapping, BouncerWeights)> - { + ) -> TransactionExecutorResult<( + CommitmentStateDiff, + Option, + VisitedSegmentsMapping, + BouncerWeights, + )> { // Get the visited segments of each contract class. // This is done by taking all the visited PCs of each contract, and compress them to one // representative for each visited segment. @@ -168,19 +178,28 @@ impl TransactionExecutor { log::debug!("Final block weights: {:?}.", self.bouncer.get_accumulated_weights()); let mut block_state = self.block_state.take().expect(BLOCK_STATE_ACCESS_ERR); - let state_diff = if self.block_context.versioned_constants.enable_stateful_compression { - state_diff_with_alias_allocation( - &mut block_state, - self.block_context + let (state_diff, compressed_state_diff) = + if self.block_context.versioned_constants.enable_stateful_compression { + let alias_contract_address = self + .block_context .versioned_constants .os_constants .os_contract_addresses - .alias_contract_address(), - )? - } else { - block_state.to_state_diff()?.state_maps - }; - Ok((state_diff.into(), visited_segments, *self.bouncer.get_accumulated_weights())) + .alias_contract_address(); + let allocated_state_diff = + state_diff_with_alias_allocation(&mut block_state, alias_contract_address)?; + let compressed_allocated_state_diff = + compress(&allocated_state_diff, &block_state, alias_contract_address)?; + (allocated_state_diff, Some(compressed_allocated_state_diff.into())) + } else { + (block_state.to_state_diff()?.state_maps, None) + }; + Ok(( + state_diff.into(), + compressed_state_diff, + visited_segments, + *self.bouncer.get_accumulated_weights(), + )) } } diff --git a/crates/blockifier_reexecution/src/state_reader/utils.rs b/crates/blockifier_reexecution/src/state_reader/utils.rs index 4e9bde49bc..2cd835171e 100644 --- a/crates/blockifier_reexecution/src/state_reader/utils.rs +++ b/crates/blockifier_reexecution/src/state_reader/utils.rs @@ -233,7 +233,7 @@ pub fn reexecute_and_verify_correctness< } // Finalize block and read actual statediff. - let (actual_state_diff, _, _) = + let (actual_state_diff, _, _, _) = transaction_executor.finalize().expect("Couldn't finalize block"); assert_eq_state_diff!(expected_state_diff, actual_state_diff); diff --git a/crates/native_blockifier/src/py_block_executor.rs b/crates/native_blockifier/src/py_block_executor.rs index 58bcc6f9e6..017ade0524 100644 --- a/crates/native_blockifier/src/py_block_executor.rs +++ b/crates/native_blockifier/src/py_block_executor.rs @@ -260,7 +260,7 @@ impl PyBlockExecutor { &mut self, ) -> NativeBlockifierResult<(PyStateDiff, PyVisitedSegmentsMapping, Py)> { log::debug!("Finalizing execution..."); - let (commitment_state_diff, visited_pcs, block_weights) = self.tx_executor().finalize()?; + let (commitment_state_diff, _, visited_pcs, block_weights) = self.tx_executor().finalize()?; let visited_pcs = visited_pcs .into_iter() .map(|(class_hash, class_visited_pcs_vec)| {