From 952373db454baa8d038609b5fbf00273635f1ca1 Mon Sep 17 00:00:00 2001
From: Noa Oved <104720318+noaov1@users.noreply.github.com>
Date: Thu, 6 Jun 2024 17:35:47 +0300
Subject: [PATCH] refactor(concurrency): use state reader generic, instead of
 updatable, when possible (#1965)

---
 crates/blockifier/src/state/cached_state.rs   | 22 ++++++++++---------
 .../src/transaction/account_transaction.rs    | 20 ++++++++---------
 2 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/crates/blockifier/src/state/cached_state.rs b/crates/blockifier/src/state/cached_state.rs
index b96c689096..5aea15be98 100644
--- a/crates/blockifier/src/state/cached_state.rs
+++ b/crates/blockifier/src/state/cached_state.rs
@@ -476,16 +476,16 @@ impl StateCache {
 
 /// Wraps a mutable reference to a `State` object, exposing its API.
 /// Used to pass ownership to a `CachedState`.
-pub struct MutRefState<'a, U: UpdatableState + ?Sized>(&'a mut U);
+pub struct MutRefState<'a, S: StateReader + ?Sized>(&'a mut S);
 
-impl<'a, U: UpdatableState + ?Sized> MutRefState<'a, U> {
-    pub fn new(state: &'a mut U) -> Self {
+impl<'a, S: StateReader + ?Sized> MutRefState<'a, S> {
+    pub fn new(state: &'a mut S) -> Self {
         Self(state)
     }
 }
 
 /// Proxies inner object to expose `State` functionality.
-impl<'a, U: UpdatableState + ?Sized> StateReader for MutRefState<'a, U> {
+impl<'a, S: StateReader + ?Sized> StateReader for MutRefState<'a, S> {
     fn get_storage_at(
         &self,
         contract_address: ContractAddress,
@@ -513,16 +513,21 @@ impl<'a, U: UpdatableState + ?Sized> StateReader for MutRefState<'a, U> {
 
 pub type TransactionalState<'a, U> = CachedState<MutRefState<'a, U>>;
 
-/// Adds the ability to perform a transactional execution.
-impl<'a, U: UpdatableState> TransactionalState<'a, U> {
+impl<'a, S: StateReader> TransactionalState<'a, S> {
     /// Creates a transactional instance from the given updatable state.
     /// It allows performing buffered modifying actions on the given state, which
     /// will either all happen (will be updated in the state and committed)
     /// or none of them (will be discarded).
-    pub fn create_transactional(state: &mut U) -> TransactionalState<'_, U> {
+    pub fn create_transactional(state: &mut S) -> TransactionalState<'_, S> {
         CachedState::new(MutRefState::new(state))
     }
 
+    /// Drops `self`.
+    pub fn abort(self) {}
+}
+
+/// Adds the ability to perform a transactional execution.
+impl<'a, U: UpdatableState> TransactionalState<'a, U> {
     /// Commits changes in the child (wrapping) state to its parent.
     pub fn commit(self) {
         let state = self.state.0;
@@ -533,9 +538,6 @@ impl<'a, U: UpdatableState> TransactionalState<'a, U> {
             &self.visited_pcs,
         )
     }
-
-    /// Drops `self`.
-    pub fn abort(self) {}
 }
 
 type StorageDiff = IndexMap<ContractAddress, IndexMap<StorageKey, StarkFelt>>;
diff --git a/crates/blockifier/src/transaction/account_transaction.rs b/crates/blockifier/src/transaction/account_transaction.rs
index 648648320d..f4f006d01b 100644
--- a/crates/blockifier/src/transaction/account_transaction.rs
+++ b/crates/blockifier/src/transaction/account_transaction.rs
@@ -303,9 +303,9 @@ impl AccountTransaction {
         Ok(())
     }
 
-    fn handle_fee<U: UpdatableState>(
+    fn handle_fee<S: StateReader>(
         &self,
-        state: &mut TransactionalState<'_, U>,
+        state: &mut TransactionalState<'_, S>,
         tx_context: Arc<TransactionContext>,
         actual_fee: Fee,
         charge_fee: bool,
@@ -372,8 +372,8 @@ impl AccountTransaction {
     /// manipulates the state to avoid that part.
     /// Note: the returned transfer call info is partial, and should be completed at the commit
     /// stage, as well as the actual sequencer balance.
-    fn concurrency_execute_fee_transfer<U: UpdatableState>(
-        state: &mut TransactionalState<'_, U>,
+    fn concurrency_execute_fee_transfer<S: StateReader>(
+        state: &mut TransactionalState<'_, S>,
         tx_context: Arc<TransactionContext>,
         actual_fee: Fee,
     ) -> TransactionExecutionResult<CallInfo> {
@@ -413,9 +413,9 @@ impl AccountTransaction {
         }
     }
 
-    fn run_non_revertible<U: UpdatableState>(
+    fn run_non_revertible<S: StateReader>(
         &self,
-        state: &mut TransactionalState<'_, U>,
+        state: &mut TransactionalState<'_, S>,
         tx_context: Arc<TransactionContext>,
         remaining_gas: &mut u64,
         validate: bool,
@@ -476,9 +476,9 @@ impl AccountTransaction {
         }
     }
 
-    fn run_revertible<U: UpdatableState>(
+    fn run_revertible<S: StateReader>(
         &self,
-        state: &mut TransactionalState<'_, U>,
+        state: &mut TransactionalState<'_, S>,
         tx_context: Arc<TransactionContext>,
         remaining_gas: &mut u64,
         validate: bool,
@@ -617,9 +617,9 @@ impl AccountTransaction {
     }
 
     /// Runs validation and execution.
-    fn run_or_revert<U: UpdatableState>(
+    fn run_or_revert<S: StateReader>(
         &self,
-        state: &mut TransactionalState<'_, U>,
+        state: &mut TransactionalState<'_, S>,
         remaining_gas: &mut u64,
         tx_context: Arc<TransactionContext>,
         validate: bool,