feat(blockifier): return compressed state diff at finalizing

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,11 +40,20 @@ pub enum TransactionExecutorError {
     StateError(#[from] StateError),
     #[error(transparent)]
     TransactionExecutionError(#[from] TransactionExecutionError),
+    #[error(transparent)]
+    CompressionError(#[from] CompressionError),
 }
 
 pub type TransactionExecutorResult<T> = Result<T, TransactionExecutorError>;
 pub type VisitedSegmentsMapping = Vec<(ClassHash, Vec<usize>)>;
 
+pub struct BlockExecutionSummary {
+    pub state_diff: CommitmentStateDiff,
+    pub compressed_state_diff: Option<CommitmentStateDiff>,
+    pub visited_segments: VisitedSegmentsMapping,
+    pub bouncer_weights: BouncerWeights,
+}
+
 /// A transaction executor, used for building a single block.
 pub struct TransactionExecutor<S: StateReader> {
     pub block_context: BlockContext,
@@ -144,10 +157,7 @@ impl<S: StateReader> TransactionExecutor<S> {
     /// Returns the state diff, a list of contract class hash with the corresponding list of
     /// visited segment values and the block weights.
     // TODO(Yoav): Consume "self".
-    pub fn finalize(
-        &mut self,
-    ) -> TransactionExecutorResult<(CommitmentStateDiff, VisitedSegmentsMapping, BouncerWeights)>
-    {
+    pub fn finalize(&mut self) -> TransactionExecutorResult<BlockExecutionSummary> {
         // 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.
@@ -169,19 +179,28 @@ impl<S: StateReader> TransactionExecutor<S> {
 
         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(BlockExecutionSummary {
+            state_diff: state_diff.into(),
+            compressed_state_diff,
+            visited_segments,
+            bouncer_weights: *self.bouncer.get_accumulated_weights(),
+        })
     }
 }
 

crates/blockifier_reexecution/src/state_reader/utils.rs

@@ -233,8 +233,8 @@ pub fn reexecute_and_verify_correctness<
     }
 
     // Finalize block and read actual statediff.
-    let (actual_state_diff, _, _) =
-        transaction_executor.finalize().expect("Couldn't finalize block");
+    let actual_state_diff =
+        transaction_executor.finalize().expect("Couldn't finalize block").state_diff;
 
     assert_eq_state_diff!(expected_state_diff, actual_state_diff);
 

crates/native_blockifier/src/py_block_executor.rs

@@ -4,7 +4,11 @@ use std::collections::HashMap;
 
 use blockifier::abi::constants as abi_constants;
 use blockifier::blockifier::config::{ContractClassManagerConfig, TransactionExecutorConfig};
-use blockifier::blockifier::transaction_executor::{TransactionExecutor, TransactionExecutorError};
+use blockifier::blockifier::transaction_executor::{
+    BlockExecutionSummary,
+    TransactionExecutor,
+    TransactionExecutorError,
+};
 use blockifier::bouncer::BouncerConfig;
 use blockifier::context::{BlockContext, ChainInfo, FeeTokenAddresses};
 use blockifier::execution::call_info::CallInfo;
@@ -258,25 +262,36 @@ impl PyBlockExecutor {
     /// visited segment values and the block weights.
     pub fn finalize(
         &mut self,
-    ) -> NativeBlockifierResult<(PyStateDiff, PyVisitedSegmentsMapping, Py<PyBytes>)> {
+    ) -> NativeBlockifierResult<(
+        PyStateDiff,
+        Option<PyStateDiff>,
+        PyVisitedSegmentsMapping,
+        Py<PyBytes>,
+    )> {
         log::debug!("Finalizing execution...");
-        let (commitment_state_diff, visited_pcs, block_weights) = self.tx_executor().finalize()?;
-        let visited_pcs = visited_pcs
+        let BlockExecutionSummary {
+            state_diff,
+            compressed_state_diff,
+            visited_segments,
+            bouncer_weights,
+        } = self.tx_executor().finalize()?;
+        let visited_pcs = visited_segments
             .into_iter()
             .map(|(class_hash, class_visited_pcs_vec)| {
                 (PyFelt::from(class_hash), class_visited_pcs_vec)
             })
             .collect();
-        let py_state_diff = PyStateDiff::from(commitment_state_diff);
+        let py_state_diff = PyStateDiff::from(state_diff);
+        let py_compressed_state_diff = compressed_state_diff.map(PyStateDiff::from);
 
         let serialized_block_weights =
-            serde_json::to_vec(&block_weights).expect("Failed serializing bouncer weights.");
+            serde_json::to_vec(&bouncer_weights).expect("Failed serializing bouncer weights.");
         let raw_block_weights =
             Python::with_gil(|py| PyBytes::new(py, &serialized_block_weights).into());
 
         log::debug!("Finalized execution.");
 
-        Ok((py_state_diff, visited_pcs, raw_block_weights))
+        Ok((py_state_diff, py_compressed_state_diff, visited_pcs, raw_block_weights))
     }
 
     // Storage Alignment API.

crates/starknet_batcher/src/block_builder.rs

@@ -3,6 +3,7 @@ use std::collections::BTreeMap;
 use async_trait::async_trait;
 use blockifier::blockifier::config::TransactionExecutorConfig;
 use blockifier::blockifier::transaction_executor::{
+    BlockExecutionSummary,
     TransactionExecutor,
     TransactionExecutorError as BlockifierTransactionExecutorError,
     TransactionExecutorResult,
@@ -163,12 +164,12 @@ impl BlockBuilderTrait for BlockBuilder {
             )
             .await?;
         }
-        let (commitment_state_diff, visited_segments_mapping, bouncer_weights) =
+        let BlockExecutionSummary { state_diff, visited_segments, bouncer_weights, .. } =
             self.executor.close_block()?;
         Ok(BlockExecutionArtifacts {
             execution_infos,
-            commitment_state_diff,
-            visited_segments_mapping,
+            commitment_state_diff: state_diff,
+            visited_segments_mapping: visited_segments,
             bouncer_weights,
         })
     }

crates/starknet_batcher/src/block_builder_test.rs

@@ -1,5 +1,8 @@
 use assert_matches::assert_matches;
-use blockifier::blockifier::transaction_executor::TransactionExecutorError;
+use blockifier::blockifier::transaction_executor::{
+    BlockExecutionSummary,
+    TransactionExecutorError,
+};
 use blockifier::bouncer::BouncerWeights;
 use blockifier::fee::fee_checks::FeeCheckError;
 use blockifier::state::errors::StateError;
@@ -253,11 +256,12 @@ fn transaction_failed_test_expectations() -> TestExpectations {
     let expected_block_artifacts = block_execution_artifacts(execution_infos_mapping);
     let expected_block_artifacts_copy = expected_block_artifacts.clone();
     mock_transaction_executor.expect_close_block().times(1).return_once(move || {
-        Ok((
-            expected_block_artifacts_copy.commitment_state_diff,
-            expected_block_artifacts_copy.visited_segments_mapping,
-            expected_block_artifacts_copy.bouncer_weights,
-        ))
+        Ok(BlockExecutionSummary {
+            state_diff: expected_block_artifacts_copy.commitment_state_diff,
+            compressed_state_diff: None,
+            visited_segments: expected_block_artifacts_copy.visited_segments_mapping,
+            bouncer_weights: expected_block_artifacts_copy.bouncer_weights,
+        })
     });
 
     let mock_tx_provider = mock_tx_provider_limitless_calls(1, vec![input_txs]);
@@ -286,11 +290,12 @@ fn set_close_block_expectations(
     let output_block_artifacts = block_builder_expected_output(block_size);
     let output_block_artifacts_copy = output_block_artifacts.clone();
     mock_transaction_executor.expect_close_block().times(1).return_once(move || {
-        Ok((
-            output_block_artifacts.commitment_state_diff,
-            output_block_artifacts.visited_segments_mapping,
-            output_block_artifacts.bouncer_weights,
-        ))
+        Ok(BlockExecutionSummary {
+            state_diff: output_block_artifacts.commitment_state_diff,
+            compressed_state_diff: None,
+            visited_segments: output_block_artifacts.visited_segments_mapping,
+            bouncer_weights: output_block_artifacts.bouncer_weights,
+        })
     });
     output_block_artifacts_copy
 }

crates/starknet_batcher/src/transaction_executor.rs

@@ -1,10 +1,8 @@
 use blockifier::blockifier::transaction_executor::{
+    BlockExecutionSummary,
     TransactionExecutor,
     TransactionExecutorResult,
-    VisitedSegmentsMapping,
 };
-use blockifier::bouncer::BouncerWeights;
-use blockifier::state::cached_state::CommitmentStateDiff;
 use blockifier::state::state_api::StateReader;
 use blockifier::transaction::objects::TransactionExecutionInfo;
 use blockifier::transaction::transaction_execution::Transaction as BlockifierTransaction;
@@ -17,9 +15,7 @@ pub trait TransactionExecutorTrait: Send {
         &mut self,
         txs: &[BlockifierTransaction],
     ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>>;
-    fn close_block(
-        &mut self,
-    ) -> TransactionExecutorResult<(CommitmentStateDiff, VisitedSegmentsMapping, BouncerWeights)>;
+    fn close_block(&mut self) -> TransactionExecutorResult<BlockExecutionSummary>;
 }
 
 impl<S: StateReader + Send + Sync> TransactionExecutorTrait for TransactionExecutor<S> {
@@ -32,10 +28,7 @@ impl<S: StateReader + Send + Sync> TransactionExecutorTrait for TransactionExecu
     }
     /// Finalizes the block creation and returns the commitment state diff, visited
     /// segments mapping and bouncer.
-    fn close_block(
-        &mut self,
-    ) -> TransactionExecutorResult<(CommitmentStateDiff, VisitedSegmentsMapping, BouncerWeights)>
-    {
+    fn close_block(&mut self) -> TransactionExecutorResult<BlockExecutionSummary> {
         self.finalize()
     }
 }