feat(blockifier): return compressed state diff at finalizing

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -14,7 +14,7 @@ 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::allocate_aliases_in_storage;
+use crate::state::stateful_compression::{allocate_aliases_in_storage, compress, CompressionError};
 use crate::transaction::errors::TransactionExecutionError;
 use crate::transaction::objects::TransactionExecutionInfo;
 use crate::transaction::transaction_execution::Transaction;
@@ -34,10 +34,18 @@ pub enum TransactionExecutorError {
     StateError(#[from] StateError),
     #[error(transparent)]
     TransactionExecutionError(#[from] TransactionExecutionError),
+    #[error(transparent)]
+    CompressionError(#[from] CompressionError),
 }
 
 pub type TransactionExecutorResult<T> = Result<T, TransactionExecutorError>;
 
+pub struct BlockExecutionSummary {
+    pub state_diff: CommitmentStateDiff,
+    pub compressed_state_diff: Option<CommitmentStateDiff>,
+    pub bouncer_weights: BouncerWeights,
+}
+
 /// A transaction executor, used for building a single block.
 pub struct TransactionExecutor<S: StateReader> {
     pub block_context: BlockContext,
@@ -140,21 +148,30 @@ impl<S: StateReader> TransactionExecutor<S> {
 
     /// Returns the state diff and the block weights.
     // TODO(Yoav): Consume "self".
-    pub fn finalize(&mut self) -> TransactionExecutorResult<(CommitmentStateDiff, BouncerWeights)> {
+    pub fn finalize(&mut self) -> TransactionExecutorResult<BlockExecutionSummary> {
         log::debug!("Final block weights: {:?}.", self.bouncer.get_accumulated_weights());
         let block_state = self.block_state.as_mut().expect(BLOCK_STATE_ACCESS_ERR);
+        let alias_contract_address = self
+            .block_context
+            .versioned_constants
+            .os_constants
+            .os_contract_addresses
+            .alias_contract_address();
         if self.block_context.versioned_constants.enable_stateful_compression {
-            allocate_aliases_in_storage(
-                block_state,
-                self.block_context
-                    .versioned_constants
-                    .os_constants
-                    .os_contract_addresses
-                    .alias_contract_address(),
-            )?;
+            allocate_aliases_in_storage(block_state, alias_contract_address)?;
         }
         let state_diff = block_state.to_state_diff()?.state_maps;
-        Ok((state_diff.into(), *self.bouncer.get_accumulated_weights()))
+        let compressed_state_diff =
+            if self.block_context.versioned_constants.enable_stateful_compression {
+                Some(compress(&state_diff, block_state, alias_contract_address)?.into())
+            } else {
+                None
+            };
+        Ok(BlockExecutionSummary {
+            state_diff: state_diff.into(),
+            compressed_state_diff,
+            bouncer_weights: *self.bouncer.get_accumulated_weights(),
+        })
     }
 }
 

crates/blockifier_reexecution/src/state_reader/utils.rs

@@ -237,7 +237,8 @@ pub fn reexecute_and_verify_correctness<
     let block_state = transaction_executor.block_state.clone();
 
     // 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;
@@ -253,20 +257,24 @@ impl PyBlockExecutor {
         })
     }
 
-    /// Returns the state diff and the block weights.
-    pub fn finalize(&mut self) -> NativeBlockifierResult<(PyStateDiff, Py<PyBytes>)> {
+    /// Returns the state diff, the stateful-compressed state diff and the block weights.
+    pub fn finalize(
+        &mut self,
+    ) -> NativeBlockifierResult<(PyStateDiff, Option<PyStateDiff>, Py<PyBytes>)> {
         log::debug!("Finalizing execution...");
-        let (commitment_state_diff, block_weights) = self.tx_executor().finalize()?;
-        let py_state_diff = PyStateDiff::from(commitment_state_diff);
+        let BlockExecutionSummary { state_diff, compressed_state_diff, bouncer_weights } =
+            self.tx_executor().finalize()?;
+        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, raw_block_weights))
+        Ok((py_state_diff, py_compressed_state_diff, 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,
@@ -165,10 +166,11 @@ impl BlockBuilderTrait for BlockBuilder {
             )
             .await?;
         }
-        let (commitment_state_diff, bouncer_weights) = self.executor.close_block()?;
+        let BlockExecutionSummary { state_diff, bouncer_weights, .. } =
+            self.executor.close_block()?;
         Ok(BlockExecutionArtifacts {
             execution_infos,
-            commitment_state_diff,
+            commitment_state_diff: state_diff,
             bouncer_weights,
             l2_gas_used,
         })

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::fee::receipt::TransactionReceipt;
@@ -261,10 +264,11 @@ 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.bouncer_weights,
-        ))
+        Ok(BlockExecutionSummary {
+            state_diff: expected_block_artifacts_copy.commitment_state_diff,
+            compressed_state_diff: None,
+            bouncer_weights: expected_block_artifacts_copy.bouncer_weights,
+        })
     });
 
     let mock_tx_provider = mock_tx_provider_limitless_calls(1, vec![input_txs]);
@@ -293,7 +297,11 @@ 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.bouncer_weights))
+        Ok(BlockExecutionSummary {
+            state_diff: output_block_artifacts.commitment_state_diff,
+            compressed_state_diff: None,
+            bouncer_weights: output_block_artifacts.bouncer_weights,
+        })
     });
     output_block_artifacts_copy
 }

crates/starknet_batcher/src/transaction_executor.rs

@@ -1,9 +1,8 @@
 use blockifier::blockifier::transaction_executor::{
+    BlockExecutionSummary,
     TransactionExecutor,
     TransactionExecutorResult,
 };
-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;
@@ -16,7 +15,7 @@ pub trait TransactionExecutorTrait: Send {
         &mut self,
         txs: &[BlockifierTransaction],
     ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>>;
-    fn close_block(&mut self) -> TransactionExecutorResult<(CommitmentStateDiff, BouncerWeights)>;
+    fn close_block(&mut self) -> TransactionExecutorResult<BlockExecutionSummary>;
 }
 
 impl<S: StateReader + Send + Sync> TransactionExecutorTrait for TransactionExecutor<S> {
@@ -29,7 +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, BouncerWeights)> {
+    fn close_block(&mut self) -> TransactionExecutorResult<BlockExecutionSummary> {
         self.finalize()
     }
 }