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,6 +40,8 @@ pub enum TransactionExecutorError {
     StateError(#[from] StateError),
     #[error(transparent)]
     TransactionExecutionError(#[from] TransactionExecutionError),
+    #[error(transparent)]
+    CompressionError(#[from] CompressionError),
 }
 
 pub type TransactionExecutorResult<T> = Result<T, TransactionExecutorError>;
@@ -141,12 +147,16 @@ impl<S: StateReader> TransactionExecutor<S> {
         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<CommitmentStateDiff>,
+        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<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((
+            state_diff.into(),
+            compressed_state_diff,
+            visited_segments,
+            *self.bouncer.get_accumulated_weights(),
+        ))
     }
 }
 

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

crates/native_blockifier/src/py_block_executor.rs

@@ -260,7 +260,7 @@ impl PyBlockExecutor {
         &mut self,
     ) -> NativeBlockifierResult<(PyStateDiff, PyVisitedSegmentsMapping, Py<PyBytes>)> {
         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)| {