From 24a88ab1879005a6a6c58f4f0aaaa471a02c1705 Mon Sep 17 00:00:00 2001
From: Noa Oved <104720318+noaov1@users.noreply.github.com>
Date: Sun, 5 May 2024 16:19:34 +0300
Subject: [PATCH] refactor(state,concurrency): add validate read set and apply
 writes to versioned state proxy (#1860)

---
 .../src/concurrency/versioned_state_proxy.rs   | 18 +++++++++++-------
 .../concurrency/versioned_state_proxy_test.rs  | 12 +++++-------
 2 files changed, 16 insertions(+), 14 deletions(-)

diff --git a/crates/blockifier/src/concurrency/versioned_state_proxy.rs b/crates/blockifier/src/concurrency/versioned_state_proxy.rs
index 94aa273acb..5de7d36d8f 100644
--- a/crates/blockifier/src/concurrency/versioned_state_proxy.rs
+++ b/crates/blockifier/src/concurrency/versioned_state_proxy.rs
@@ -69,7 +69,7 @@ impl<S: StateReader> VersionedState<S> {
     // accessing this function should be protected by a mutex to ensure thread safety.
     // TODO: Consider coupling the tx index with the read set to ensure any mismatch between them
     // will cause the validation to fail.
-    pub fn validate_read_set(&mut self, tx_index: TxIndex, reads: &StateMaps) -> bool {
+    fn validate_read_set(&mut self, tx_index: TxIndex, reads: &StateMaps) -> bool {
         // If is the first transaction in the chunk, then the read set is valid. Since it has no
         // predecessors, there's nothing to compare it to.
         if tx_index == 0 {
@@ -117,7 +117,7 @@ impl<S: StateReader> VersionedState<S> {
         true
     }
 
-    pub fn apply_writes(
+    fn apply_writes(
         &mut self,
         tx_index: TxIndex,
         writes: &StateMaps,
@@ -152,10 +152,6 @@ impl<S: StateReader> ThreadSafeVersionedState<S> {
     pub fn pin_version(&self, tx_index: TxIndex) -> VersionedStateProxy<S> {
         VersionedStateProxy { tx_index, state: self.0.clone() }
     }
-
-    pub fn state(&self) -> LockedVersionedState<'_, S> {
-        self.0.lock().expect("Failed to acquire state lock.")
-    }
 }
 
 pub struct VersionedStateProxy<S: StateReader> {
@@ -164,9 +160,17 @@ pub struct VersionedStateProxy<S: StateReader> {
 }
 
 impl<S: StateReader> VersionedStateProxy<S> {
-    pub fn state(&self) -> LockedVersionedState<'_, S> {
+    fn state(&self) -> LockedVersionedState<'_, S> {
         self.state.lock().expect("Failed to acquire state lock.")
     }
+
+    pub fn validate_read_set(&self, reads: &StateMaps) -> bool {
+        self.state().validate_read_set(self.tx_index, reads)
+    }
+
+    pub fn apply_writes(&self, writes: &StateMaps, class_hash_to_class: &ContractClassMapping) {
+        self.state().apply_writes(self.tx_index, writes, class_hash_to_class)
+    }
 }
 
 impl<S: StateReader> StateReader for VersionedStateProxy<S> {
diff --git a/crates/blockifier/src/concurrency/versioned_state_proxy_test.rs b/crates/blockifier/src/concurrency/versioned_state_proxy_test.rs
index a4e30cf728..3870d2d9f0 100644
--- a/crates/blockifier/src/concurrency/versioned_state_proxy_test.rs
+++ b/crates/blockifier/src/concurrency/versioned_state_proxy_test.rs
@@ -280,7 +280,7 @@ fn test_validate_read_set(
     let transactional_state = CachedState::from(safe_versioned_state.pin_version(1));
 
     // Validating tx index 0 always succeeds.
-    assert!(safe_versioned_state.state().validate_read_set(0, &StateMaps::default()));
+    assert!(safe_versioned_state.pin_version(0).validate_read_set(&StateMaps::default()));
 
     assert!(transactional_state.cache.borrow().initial_reads.storage.is_empty());
     transactional_state.get_storage_at(contract_address, storage_key).unwrap();
@@ -303,8 +303,8 @@ fn test_validate_read_set(
 
     assert!(
         safe_versioned_state
-            .state()
-            .validate_read_set(1, &transactional_state.cache.borrow().initial_reads)
+            .pin_version(1)
+            .validate_read_set(&transactional_state.cache.borrow().initial_reads)
     );
 }
 
@@ -329,8 +329,7 @@ fn test_apply_writes(
     transactional_states[0].set_contract_class(class_hash, contract_class_0.clone()).unwrap();
     assert_eq!(transactional_states[0].class_hash_to_class.borrow().len(), 1);
 
-    safe_versioned_state.state().apply_writes(
-        0,
+    safe_versioned_state.pin_version(0).apply_writes(
         &transactional_states[0].cache.borrow().writes,
         &transactional_states[0].class_hash_to_class.borrow().clone(),
     );
@@ -358,8 +357,7 @@ fn test_apply_writes_reexecute_scenario(
     // updated.
     assert!(transactional_states[1].get_class_hash_at(contract_address).unwrap() == class_hash);
 
-    safe_versioned_state.state().apply_writes(
-        0,
+    safe_versioned_state.pin_version(0).apply_writes(
         &transactional_states[0].cache.borrow().writes,
         &transactional_states[0].class_hash_to_class.borrow().clone(),
     );