refactor(concurrency): use visited_pcs in execute_chunk

crates/blockifier/src/blockifier/transaction_executor.rs

@@ -1,4 +1,6 @@
 #[cfg(feature = "concurrency")]
+use std::collections::{HashMap, HashSet};
+#[cfg(feature = "concurrency")]
 use std::sync::Arc;
 #[cfg(feature = "concurrency")]
 use std::sync::Mutex;
@@ -240,25 +242,23 @@ impl<S: StateReader + Send + Sync> TransactionExecutor<S> {
         });
 
         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<ClassHash, HashSet<usize>> = 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<S: StateReader + Send + Sync> TransactionExecutor<S> {
                      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

crates/blockifier/src/concurrency/versioned_state.rs

@@ -199,7 +199,11 @@ impl<S: StateReader> VersionedState<S> {
 }
 
 impl<U: UpdatableState> VersionedState<U> {
-    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<ClassHash, HashSet<usize>>,
+    ) -> U {
         if n_committed_txs == 0 {
             return self.into_initial_state();
         }
@@ -208,8 +212,7 @@ impl<U: UpdatableState> VersionedState<U> {
         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
     }
 }

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!(

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<ClassHash, HashSet<usize>>,
+    ) -> U {
+        self.state
+            .into_inner_state()
+            .commit_chunk_and_recover_block_state(n_committed_txs, visited_pcs)
     }
 }