diff --git a/crates/blockifier/src/blockifier/transaction_executor.rs b/crates/blockifier/src/blockifier/transaction_executor.rs index f3d493fd29..9aa1ad102c 100644 --- a/crates/blockifier/src/blockifier/transaction_executor.rs +++ b/crates/blockifier/src/blockifier/transaction_executor.rs @@ -6,6 +6,8 @@ use std::sync::Mutex; use itertools::FoldWhile::{Continue, Done}; use itertools::Itertools; use starknet_api::core::ClassHash; +#[cfg(feature = "concurrency")] +use std::collections::{HashMap, HashSet}; use thiserror::Error; use crate::blockifier::config::TransactionExecutorConfig; @@ -240,25 +242,23 @@ 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> = 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)); + for (class_hash, class_visited_pcs) in locked_execution_output.visited_pcs { + visited_pcs.entry(class_hash).or_default().extend(class_visited_pcs); + } + } let block_state_after_commit = Arc::try_unwrap(worker_executor) .unwrap_or_else(|_| { @@ -268,7 +268,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 4b4be73ed0..fd36781f74 100644 --- a/crates/blockifier/src/concurrency/versioned_state.rs +++ b/crates/blockifier/src/concurrency/versioned_state.rs @@ -199,7 +199,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(); } @@ -208,8 +212,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 3525657244..467ccb9ba3 100644 --- a/crates/blockifier/src/concurrency/versioned_state_test.rs +++ b/crates/blockifier/src/concurrency/versioned_state_test.rs @@ -642,8 +642,9 @@ fn test_versioned_proxy_state_flow( for proxy in versioned_proxy_states { drop(proxy); } - let modified_block_state = - safe_versioned_state.into_inner_state().commit_chunk_and_recover_block_state(4); + let modified_block_state = 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 1ed7dda60c..8122521087 100644 --- a/crates/blockifier/src/concurrency/worker_logic.rs +++ b/crates/blockifier/src/concurrency/worker_logic.rs @@ -291,8 +291,14 @@ 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) } }