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)| {