refactor(execution): make block state of TransactionExecutor an Optio…

…n to allow moving it to chunks

crates/blockifier/src/blockifier/stateful_validator.rs

@@ -110,7 +110,7 @@ impl<S: StateReader> StatefulValidator<S> {
         // Run pre-validation in charge fee mode to perform fee and balance related checks.
         let charge_fee = true;
         tx.perform_pre_validation_stage(
-            &mut self.tx_executor.state,
+            self.tx_executor.block_state.as_mut().expect("The block state should be `Some`."),
             tx_context,
             charge_fee,
             strict_nonce_check,
@@ -127,7 +127,12 @@ impl<S: StateReader> StatefulValidator<S> {
         tx_info: &TransactionInfo,
         deploy_account_tx_hash: Option<TransactionHash>,
     ) -> StatefulValidatorResult<bool> {
-        let nonce = self.tx_executor.state.get_nonce_at(tx_info.sender_address())?;
+        let nonce = self
+            .tx_executor
+            .block_state
+            .as_ref()
+            .expect("The block state should be `Some`.")
+            .get_nonce_at(tx_info.sender_address())?;
         let tx_nonce = tx_info.nonce();
 
         let deploy_account_not_processed =

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -45,21 +45,29 @@ pub struct TransactionExecutor<S: StateReader> {
     pub config: TransactionExecutorConfig,
 
     // State-related fields.
-    pub state: CachedState<S>,
+    // The transaction executor operates at the block level. In concurrency mode, it moves the
+    // block state to the worker executor - operating at the chunk level - and gets it back after
+    // committing the chunk. The block state is wrapped with an Option<_> to allow setting it to
+    // `None` while it is moved to the worker executor.
+    pub block_state: Option<CachedState<S>>,
 }
 
 impl<S: StateReader> TransactionExecutor<S> {
     pub fn new(
-        state: CachedState<S>,
+        block_state: CachedState<S>,
         block_context: BlockContext,
         bouncer_config: BouncerConfig,
         config: TransactionExecutorConfig,
     ) -> Self {
         log::debug!("Initializing Transaction Executor...");
         // Note: the state might not be empty even at this point; it is the creator's
         // responsibility to tune the bouncer according to pre and post block process.
-        let tx_executor =
-            Self { block_context, bouncer: Bouncer::new(bouncer_config), config, state };
+        let tx_executor = Self {
+            block_context,
+            bouncer: Bouncer::new(bouncer_config),
+            config,
+            block_state: Some(block_state),
+        };
         log::debug!("Initialized Transaction Executor.");
 
         tx_executor
@@ -73,7 +81,9 @@ impl<S: StateReader> TransactionExecutor<S> {
         tx: &Transaction,
         charge_fee: bool,
     ) -> TransactionExecutorResult<TransactionExecutionInfo> {
-        let mut transactional_state = TransactionalState::create_transactional(&mut self.state);
+        let mut transactional_state = TransactionalState::create_transactional(
+            self.block_state.as_mut().expect("The block state should be `Some`."),
+        );
         let validate = true;
 
         let tx_execution_result =
@@ -163,11 +173,14 @@ impl<S: StateReader> TransactionExecutor<S> {
         // For fee charging purposes, the nonce-increment cost is taken into consideration when
         // calculating the fees for validation.
         // Note: This assumes that the state is reset between calls to validate.
-        self.state.increment_nonce(tx_info.sender_address())?;
+        self.block_state
+            .as_mut()
+            .expect("The block state should be `Some`.")
+            .increment_nonce(tx_info.sender_address())?;
 
         let limit_steps_by_resources = true;
         let validate_call_info = account_tx.validate_tx(
-            &mut self.state,
+            self.block_state.as_mut().expect("The block state should be `Some`."),
             &mut execution_resources,
             tx_context.clone(),
             &mut remaining_gas,
@@ -177,7 +190,11 @@ impl<S: StateReader> TransactionExecutor<S> {
         let tx_receipt = TransactionReceipt::from_account_tx(
             account_tx,
             &tx_context,
-            &self.state.get_actual_state_changes()?,
+            &self
+                .block_state
+                .as_mut()
+                .expect("The block state should be `Some`.")
+                .get_actual_state_changes()?,
             &execution_resources,
             validate_call_info.iter(),
             0,
@@ -196,18 +213,28 @@ impl<S: StateReader> TransactionExecutor<S> {
         // This is done by taking all the visited PCs of each contract, and compress them to one
         // representative for each visited segment.
         let visited_segments = self
-            .state
+            .block_state
+            .as_ref()
+            .expect("The block state should be `Some`.")
             .visited_pcs
             .iter()
             .map(|(class_hash, class_visited_pcs)| -> TransactionExecutorResult<_> {
-                let contract_class = self.state.get_compiled_contract_class(*class_hash)?;
+                let contract_class = self
+                    .block_state
+                    .as_ref()
+                    .expect("The block state should be `Some`.")
+                    .get_compiled_contract_class(*class_hash)?;
                 Ok((*class_hash, contract_class.get_visited_segments(class_visited_pcs)?))
             })
             .collect::<TransactionExecutorResult<_>>()?;
 
         log::debug!("Final block weights: {:?}.", self.bouncer.get_accumulated_weights());
         Ok((
-            self.state.to_state_diff()?.into(),
+            self.block_state
+                .as_mut()
+                .expect("The block state should be `Some`.")
+                .to_state_diff()?
+                .into(),
             visited_segments,
             *self.bouncer.get_accumulated_weights(),
         ))

crates/blockifier/src/blockifier/transaction_executor_test.rs

@@ -338,7 +338,15 @@ fn test_execute_txs_bouncing(block_context: BlockContext) {
     assert!(results[2].is_ok());
 
     // Check state.
-    assert_eq!(tx_executor.state.get_nonce_at(account_address).unwrap(), nonce!(2_u32));
+    assert_eq!(
+        tx_executor
+            .block_state
+            .as_ref()
+            .expect("The block state should be `Some`.")
+            .get_nonce_at(account_address)
+            .unwrap(),
+        nonce!(2_u32)
+    );
 
     // Check idempotency: excess transactions should not be added.
     let remaining_txs = &txs[expected_offset..];
@@ -352,5 +360,13 @@ fn test_execute_txs_bouncing(block_context: BlockContext) {
     assert_eq!(remaining_tx_results.len(), 2);
     assert!(remaining_tx_results[0].is_ok());
     assert!(remaining_tx_results[1].is_ok());
-    assert_eq!(tx_executor.state.get_nonce_at(account_address).unwrap(), nonce!(4_u32));
+    assert_eq!(
+        tx_executor
+            .block_state
+            .as_ref()
+            .expect("The block state should be `Some`.")
+            .get_nonce_at(account_address)
+            .unwrap(),
+        nonce!(4_u32)
+    );
 }

crates/native_blockifier/src/py_block_executor_test.rs

@@ -53,8 +53,13 @@ fn global_contract_cache_update() {
 
     assert_eq!(block_executor.global_contract_cache.lock().cache_size(), 0);
 
-    let queried_contract_class =
-        block_executor.tx_executor().state.get_compiled_contract_class(class_hash).unwrap();
+    let queried_contract_class = block_executor
+        .tx_executor()
+        .block_state
+        .as_ref()
+        .expect("The block state should be `Some`.")
+        .get_compiled_contract_class(class_hash)
+        .unwrap();
 
     assert_eq!(queried_contract_class, contract_class);
     assert_eq!(block_executor.global_contract_cache.lock().cache_size(), 1);