chore(blockifier): can statediff be returned from execute?

Author: avivg-starkware

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -11,7 +11,7 @@ use crate::blockifier::config::TransactionExecutorConfig;
 use crate::bouncer::{Bouncer, BouncerWeights};
 use crate::concurrency::worker_logic::WorkerExecutor;
 use crate::context::BlockContext;
-use crate::state::cached_state::{CachedState, CommitmentStateDiff, TransactionalState};
+use crate::state::cached_state::{CachedState, CommitmentStateDiff, StateMaps, TransactionalState};
 use crate::state::errors::StateError;
 use crate::state::state_api::{StateReader, StateResult};
 use crate::state::stateful_compression::{allocate_aliases_in_storage, compress, CompressionError};
@@ -103,7 +103,7 @@ impl<S: StateReader> TransactionExecutor<S> {
     pub fn execute(
         &mut self,
         tx: &Transaction,
-    ) -> TransactionExecutorResult<TransactionExecutionInfo> {
+    ) -> TransactionExecutorResult<(TransactionExecutionInfo, StateMaps)> {
         let mut transactional_state = TransactionalState::create_transactional(
             self.block_state.as_mut().expect(BLOCK_STATE_ACCESS_ERR),
         );
@@ -116,14 +116,15 @@ impl<S: StateReader> TransactionExecutor<S> {
             Ok(tx_execution_info) => {
                 let tx_state_changes_keys =
                     transactional_state.get_actual_state_changes()?.state_maps.into_keys();
+                let state_diff = transactional_state.to_state_diff()?.state_maps;
                 self.bouncer.try_update(
                     &transactional_state,
                     &tx_state_changes_keys,
                     &tx_execution_info.summarize(&self.block_context.versioned_constants),
                     &tx_execution_info.receipt.resources,
                 )?;
                 transactional_state.commit();
-                Ok(tx_execution_info)
+                Ok((tx_execution_info, state_diff))
             }
             Err(error) => {
                 transactional_state.abort();
@@ -139,7 +140,7 @@ impl<S: StateReader> TransactionExecutor<S> {
         let mut results = Vec::new();
         for tx in txs {
             match self.execute(tx) {
-                Ok(tx_execution_info) => results.push(Ok(tx_execution_info)),
+                Ok((tx_execution_info, _state_diff)) => results.push(Ok(tx_execution_info)),
                 Err(TransactionExecutorError::BlockFull) => break,
                 Err(error) => results.push(Err(error)),
             }

crates/blockifier/src/blockifier/transaction_executor_test.rs

@@ -59,7 +59,7 @@ fn tx_executor_test_body<S: StateReader>(
     // TODO(Arni, 30/03/2024): Consider adding a test for the transaction execution info. If A test
     // should not be added, rename the test to `test_bouncer_info`.
     // TODO(Arni, 30/03/2024): Test all bouncer weights.
-    let _tx_execution_info = tx_executor.execute(&tx).unwrap();
+    let (_tx_execution_info, _state_diff) = tx_executor.execute(&tx).unwrap();
     let bouncer_weights = tx_executor.bouncer.get_accumulated_weights();
     assert_eq!(bouncer_weights.state_diff_size, expected_bouncer_weights.state_diff_size);
     assert_eq!(

crates/native_blockifier/src/py_block_executor.rs

@@ -197,7 +197,7 @@ impl PyBlockExecutor {
         optional_py_class_info: Option<PyClassInfo>,
     ) -> NativeBlockifierResult<Py<PyBytes>> {
         let tx: Transaction = py_tx(tx, optional_py_class_info).expect(PY_TX_PARSING_ERR);
-        let tx_execution_info = self.tx_executor().execute(&tx)?;
+        let (tx_execution_info, _state_diff) = self.tx_executor().execute(&tx)?;
         let thin_tx_execution_info =
             ThinTransactionExecutionInfo::from_tx_execution_info(tx_execution_info);