refactor(concurrency): use visited_pcs in execute_chunk

starkware-libs · Jun 17, 2024 · a0b44fc · a0b44fc
1 parent 4505d16
commit a0b44fc
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 26 deletions.
diff --git 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<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::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<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

diff --git a/crates/blockifier/src/concurrency/versioned_state.rs b/crates/blockifier/src/concurrency/versioned_state.rs
@@ -192,7 +192,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();
         }
@@ -201,8 +205,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
     }
 }

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