refactor(concurrency): use apply_writes in the commit of VersionedState

crates/blockifier/src/concurrency/versioned_state.rs

@@ -8,7 +8,7 @@ use starknet_api::state::StorageKey;
 use crate::concurrency::versioned_storage::VersionedStorage;
 use crate::concurrency::TxIndex;
 use crate::execution::contract_class::ContractClass;
-use crate::state::cached_state::{CachedState, ContractClassMapping, StateMaps};
+use crate::state::cached_state::{ContractClassMapping, StateMaps};
 use crate::state::errors::StateError;
 use crate::state::state_api::{StateReader, StateResult, UpdatableState};
 
@@ -71,18 +71,6 @@ impl<S: StateReader> VersionedState<S> {
         }
     }
 
-    pub fn commit<T>(&mut self, from_index: TxIndex, parent_state: &mut CachedState<T>)
-    where
-        T: StateReader,
-    {
-        let writes = self.get_writes_up_to_index(from_index);
-
-        parent_state.update_cache(
-            &writes,
-            self.compiled_contract_classes.get_writes_up_to_index(from_index),
-        );
-    }
-
     // TODO(Mohammad, 01/04/2024): Store the read set (and write set) within a shared
     // object (probabily `VersionedState`). As RefCell operations are not thread-safe. Therefore,
     // accessing this function should be protected by a mutex to ensure thread safety.
@@ -199,12 +187,27 @@ impl<S: StateReader> VersionedState<S> {
         }
     }
 
-    #[allow(dead_code)]
     fn into_initial_state(self) -> S {
         self.initial_state
     }
 }
 
+impl<U: UpdatableState> VersionedState<U> {
+    pub fn commit_chunk_and_recover_block_state(mut self, commit_index: TxIndex) -> U {
+        let writes = self.get_writes_up_to_index(commit_index);
+        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
+    }
+}
+
+// TODO(barak, 01/07/2024): Re-consider the API (pub functions) of VersionedState,
+// ThreadSafeVersionedState and VersionedStateProxy.
+// TODO(barak, 01/07/2024): Re-consider the necessity ot ThreadSafeVersionedState once the worker
+// logic is completed.
 pub struct ThreadSafeVersionedState<S: StateReader>(Arc<Mutex<VersionedState<S>>>);
 pub type LockedVersionedState<'a, S> = MutexGuard<'a, VersionedState<S>>;
 
@@ -217,7 +220,6 @@ impl<S: StateReader> ThreadSafeVersionedState<S> {
         VersionedStateProxy { tx_index, state: self.0.clone() }
     }
 
-    #[allow(dead_code)]
     fn into_inner_state(self) -> VersionedState<S> {
         Arc::try_unwrap(self.0)
             .unwrap_or_else(|_| {
@@ -237,6 +239,12 @@ impl<S: StateReader> Clone for ThreadSafeVersionedState<S> {
     }
 }
 
+impl<U: UpdatableState> ThreadSafeVersionedState<U> {
+    pub fn commit_chunk_and_recover_block_state(self, commit_index: TxIndex) -> U {
+        self.into_inner_state().commit_chunk_and_recover_block_state(commit_index)
+    }
+}
+
 pub struct VersionedStateProxy<S: StateReader> {
     pub tx_index: TxIndex,
     pub state: Arc<Mutex<VersionedState<S>>>,

crates/blockifier/src/concurrency/versioned_state_test.rs

@@ -532,7 +532,6 @@ fn test_versioned_proxy_state_flow(
     let contract_address = contract_address!("0x1");
     let class_hash = ClassHash(stark_felt!(27_u8));
 
-    let mut block_state = CachedState::from(DictStateReader::default());
     let mut versioned_proxy_states: Vec<VersionedStateProxy<CachedState<DictStateReader>>> =
         (0..4).map(|i| safe_versioned_state.pin_version(i)).collect();
 
@@ -566,8 +565,13 @@ fn test_versioned_proxy_state_flow(
     }
 
     // Check the final state.
-    safe_versioned_state.0.lock().unwrap().commit(4, &mut block_state);
+    for proxy in versioned_proxy_states {
+        drop(proxy);
+    }
+    let modified_block_state = safe_versioned_state.commit_chunk_and_recover_block_state(4);
 
-    assert!(block_state.get_class_hash_at(contract_address).unwrap() == class_hash_3);
-    assert!(block_state.get_compiled_contract_class(class_hash).unwrap() == contract_class_2);
+    assert!(modified_block_state.get_class_hash_at(contract_address).unwrap() == class_hash_3);
+    assert!(
+        modified_block_state.get_compiled_contract_class(class_hash).unwrap() == contract_class_2
+    );
 }