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<TransactionExecutionOutput> {
         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();
@@ -135,11 +136,13 @@ impl<S: StateReader> TransactionExecutor<S> {
     pub fn execute_txs_sequentially_inner(
         &mut self,
         txs: &[Transaction],
-    ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>> {
+    ) -> Vec<TransactionExecutorResult<(TransactionExecutionInfo, StateMaps)>> {
         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, state_diff)))
+                }
                 Err(TransactionExecutorError::BlockFull) => break,
                 Err(error) => results.push(Err(error)),
             }
@@ -233,14 +236,18 @@ impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
         txs: &[Transaction],
     ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>> {
         #[cfg(not(feature = "cairo_native"))]
-        return self.execute_txs_sequentially_inner(txs);
+        return self
+            .execute_txs_sequentially_inner(txs)
+            .into_iter()
+            .map(|result| result.map(|(info, _)| info))
+            .collect();
         #[cfg(feature = "cairo_native")]
         {
             // TODO meshi: find a way to access the contract class manager config from transaction
             // executor.
             let txs = txs.to_vec();
             std::thread::scope(|s| {
-                std::thread::Builder::new()
+                let execution_outputs = std::thread::Builder::new()
                     // when running Cairo natively, the real stack is used and could get overflowed
                     // (unlike the VM where the stack is simulated in the heap as a memory segment).
                     //
@@ -256,7 +263,9 @@ impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
                     .spawn_scoped(s, || self.execute_txs_sequentially_inner(&txs))
                     .expect("Failed to spawn thread")
                     .join()
-                    .expect("Failed to join thread.")
+                    .expect("Failed to join thread.");
+
+                execution_outputs.into_iter().map(|result| result.map(|(info, _)| info)).collect()
             })
         }
     }

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