diff --git a/crates/blockifier/src/blockifier/transaction_executor.rs b/crates/blockifier/src/blockifier/transaction_executor.rs index f3d493fd29..f85425a264 100644 --- a/crates/blockifier/src/blockifier/transaction_executor.rs +++ b/crates/blockifier/src/blockifier/transaction_executor.rs @@ -20,6 +20,7 @@ use crate::transaction::errors::TransactionExecutionError; use crate::transaction::objects::TransactionExecutionInfo; use crate::transaction::transaction_execution::Transaction; use crate::transaction::transactions::ExecutableTransaction; +use std::collections::HashMap; #[cfg(test)] #[path = "transaction_executor_test.rs"] @@ -240,25 +241,22 @@ impl TransactionExecutor { }); let n_committed_txs = worker_executor.scheduler.get_n_committed_txs(); - let tx_execution_results = worker_executor - .execution_outputs - .iter() - .fold_while(Vec::new(), |mut results, execution_output| { - if results.len() >= n_committed_txs { - Done(results) - } else { - let locked_execution_output = execution_output - .lock() - .expect("Failed to lock execution output.") - .take() - .expect("Output must be ready."); - results.push( - locked_execution_output.result.map_err(TransactionExecutorError::from), - ); - Continue(results) - } - }) - .into_inner(); + let mut tx_execution_results = Vec::new(); + let mut visited_pcs = HashMap::new(); + for execution_output in worker_executor.execution_outputs.iter() { + if tx_execution_results.len() >= n_committed_txs { + break; + } + let locked_execution_output = execution_output + .lock() + .expect("Failed to lock execution output.") + .take() + .expect("Output must be ready."); + tx_execution_results.push( + locked_execution_output.result.map_err(TransactionExecutorError::from), + ); + visited_pcs.extend(locked_execution_output.visited_pcs); + } let block_state_after_commit = Arc::try_unwrap(worker_executor) .unwrap_or_else(|_| { @@ -268,7 +266,7 @@ impl TransactionExecutor { it." ) }) - .commit_chunk_and_recover_block_state(n_committed_txs); + .commit_chunk_and_recover_block_state(n_committed_txs, visited_pcs); self.block_state.replace(block_state_after_commit); tx_execution_results diff --git a/crates/blockifier/src/concurrency/versioned_state.rs b/crates/blockifier/src/concurrency/versioned_state.rs index 1823baf2f3..9908e0a9f0 100644 --- a/crates/blockifier/src/concurrency/versioned_state.rs +++ b/crates/blockifier/src/concurrency/versioned_state.rs @@ -192,7 +192,11 @@ impl VersionedState { } impl VersionedState { - pub fn commit_chunk_and_recover_block_state(mut self, n_committed_txs: usize) -> U { + pub fn commit_chunk_and_recover_block_state( + mut self, + n_committed_txs: usize, + visited_pcs: HashMap>, + ) -> U { if n_committed_txs == 0 { return self.into_initial_state(); } @@ -201,8 +205,7 @@ impl VersionedState { let class_hash_to_class = self.compiled_contract_classes.get_writes_up_to_index(commit_index); let mut state = self.into_initial_state(); - // TODO(barak, 01/08/2024): Add visited_pcs argument to `apply_writes`. - state.apply_writes(&writes, &class_hash_to_class, &HashMap::default()); + state.apply_writes(&writes, &class_hash_to_class, &visited_pcs); state } } diff --git a/crates/blockifier/src/concurrency/versioned_state_test.rs b/crates/blockifier/src/concurrency/versioned_state_test.rs index 7482e45f93..9b1cf2466f 100644 --- a/crates/blockifier/src/concurrency/versioned_state_test.rs +++ b/crates/blockifier/src/concurrency/versioned_state_test.rs @@ -569,7 +569,7 @@ fn test_versioned_proxy_state_flow( drop(proxy); } let modified_block_state = - safe_versioned_state.into_inner_state().commit_chunk_and_recover_block_state(4); + safe_versioned_state.into_inner_state().commit_chunk_and_recover_block_state(4, HashMap::new()); assert!(modified_block_state.get_class_hash_at(contract_address).unwrap() == class_hash_3); assert!( diff --git a/crates/blockifier/src/concurrency/worker_logic.rs b/crates/blockifier/src/concurrency/worker_logic.rs index 207b002466..397b86c71b 100644 --- a/crates/blockifier/src/concurrency/worker_logic.rs +++ b/crates/blockifier/src/concurrency/worker_logic.rs @@ -291,8 +291,12 @@ impl<'a, S: StateReader> WorkerExecutor<'a, S> { } impl<'a, U: UpdatableState> WorkerExecutor<'a, U> { - pub fn commit_chunk_and_recover_block_state(self, n_committed_txs: usize) -> U { - self.state.into_inner_state().commit_chunk_and_recover_block_state(n_committed_txs) + pub fn commit_chunk_and_recover_block_state( + self, + n_committed_txs: usize, + visited_pcs: HashMap>, + ) -> U { + self.state.into_inner_state().commit_chunk_and_recover_block_state(n_committed_txs, visited_pcs) } }