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};
@@ -27,6 +27,8 @@ pub mod transaction_executor_test;
 pub const BLOCK_STATE_ACCESS_ERR: &str = "Error: The block state should be `Some`.";
 pub const DEFAULT_STACK_SIZE: usize = 60 * 1024 * 1024;
 
+pub type TransactionExecutionOutput = (TransactionExecutionInfo, StateMaps);
+
 #[derive(Debug, Error)]
 pub enum TransactionExecutorError {
     #[error("Transaction cannot be added to the current block, block capacity reached.")]
@@ -103,7 +105,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 +118,16 @@ 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();
@@ -132,14 +136,16 @@ impl<S: StateReader> TransactionExecutor<S> {
         }
     }
 
-    pub fn execute_txs_sequentially_inner(
+    fn execute_txs_sequentially_inner(
         &mut self,
         txs: &[Transaction],
-    ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>> {
+    ) -> Vec<TransactionExecutorResult<TransactionExecutionOutput>> {
         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)),
             }
@@ -196,6 +202,9 @@ impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
         if !self.config.concurrency_config.enabled {
             log::debug!("Executing transactions sequentially.");
             self.execute_txs_sequentially(txs)
+                .into_iter()
+                .map(|result| result.map(|(info, _)| info))
+                .collect()
         } else {
             log::debug!("Executing transactions concurrently.");
             let chunk_size = self.config.concurrency_config.chunk_size;
@@ -228,40 +237,41 @@ impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
         }
     }
 
-    pub fn execute_txs_sequentially(
+    fn execute_txs_sequentially(
         &mut self,
         txs: &[Transaction],
-    ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>> {
+    ) -> Vec<TransactionExecutorResult<TransactionExecutionOutput>> {
         #[cfg(not(feature = "cairo_native"))]
         return self.execute_txs_sequentially_inner(txs);
+
         #[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()
-                    // 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).
-                    //
-                    // We pre-allocate the stack here, and not during Native execution (not trivial), so it
-                    // needs to be big enough ahead.
-                    // However, making it very big is wasteful (especially with multi-threading).
-                    // So, the stack size should support calls with a reasonable gas limit, for extremely deep
-                    // recursions to reach out-of-gas before hitting the bottom of the recursion.
-                    //
-                    // The gas upper bound is MAX_POSSIBLE_SIERRA_GAS, and sequencers must not raise it without
-                    // adjusting the stack size.
-                    .stack_size(self.config.stack_size)
-                    .spawn_scoped(s, || self.execute_txs_sequentially_inner(&txs))
-                    .expect("Failed to spawn thread")
-                    .join()
-                    .expect("Failed to join thread.")
+                // 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).
+                //
+                // We pre-allocate the stack here, and not during Native execution (not trivial), so it
+                // needs to be big enough ahead.
+                // However, making it very big is wasteful (especially with multi-threading).
+                // So, the stack size should support calls with a reasonable gas limit, for extremely deep
+                // recursions to reach out-of-gas before hitting the bottom of the recursion.
+                //
+                // The gas upper bound is MAX_POSSIBLE_SIERRA_GAS, and sequencers must not raise it without
+                // adjusting the stack size.
+                .stack_size(self.config.stack_size)
+                .spawn_scoped(s, || self.execute_txs_sequentially_inner(&txs))
+                .expect("Failed to spawn thread")
+                .join()
+                .expect("Failed to join thread.")
             })
         }
     }
 
-    pub fn execute_chunk(
+    fn execute_chunk(
         &mut self,
         chunk: &[Transaction],
     ) -> Vec<TransactionExecutorResult<TransactionExecutionInfo>> {

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