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.
@@ -205,6 +193,10 @@ impl<S: StateReader> VersionedState<S> {
     }
 }
 
+// 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>>;
 
@@ -236,6 +228,17 @@ impl<S: StateReader> Clone for ThreadSafeVersionedState<S> {
     }
 }
 
+impl<U: UpdatableState> ThreadSafeVersionedState<U> {
+    pub fn commit_chunk_and_recover_block_state(self, index: TxIndex) -> U{
+        let mut versioned_state = self.into_inner_state();
+        let writes = versioned_state.get_writes_up_to_index(index);
+        let class_hash_to_class = versioned_state.compiled_contract_classes.get_writes_up_to_index(index);
+        // TODO(barak, 01/08/2024): Add visited_pcs argument to `apply_writes`.
+        versioned_state.initial_state.apply_writes(&writes, &class_hash_to_class, &HashMap::default());
+        versioned_state.into_initial_state()
+    }
+}
+
 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
+    );
 }

crates/blockifier/src/concurrency/worker_logic.rs

@@ -20,7 +20,7 @@ use crate::context::BlockContext;
 use crate::execution::execution_utils::stark_felt_to_felt;
 use crate::fee::fee_utils::get_sequencer_balance_keys;
 use crate::state::cached_state::{
-    ContractClassMapping, StateChanges, StateMaps, TransactionalState,
+    ContractClassMapping, StateChanges, StateMaps, TransactionalState
 };
 use crate::state::state_api::{StateReader, UpdatableState};
 use crate::transaction::errors::TransactionExecutionError;