Add cache to the BlockifierStateAdapter for setters

crates/pallets/starknet/runtime_api/src/lib.rs

@@ -94,8 +94,6 @@ sp_api::decl_runtime_apis! {
         fn is_transaction_fee_disabled() -> bool;
         /// Return messages sent to L1 during tx execution
         fn get_tx_messages_to_l1(tx_hash: TransactionHash) -> Vec<MessageToL1>;
-        /// Check if L1 Message Nonce has not been used
-        fn l1_nonce_unused(nonce: Nonce) -> bool;
     }
 
     pub trait ConvertTransactionRuntimeApi {

crates/pallets/starknet/src/blockifier_state_adapter.rs

@@ -1,5 +1,6 @@
 use alloc::collections::{BTreeMap, BTreeSet};
 use core::marker::PhantomData;
+use std::collections::{HashMap, HashSet};
 
 use blockifier::execution::contract_class::ContractClass;
 use blockifier::state::cached_state::StateChangesCount;
@@ -9,9 +10,11 @@ use indexmap::IndexMap;
 use mc_db::storage::StorageHandler;
 use mp_felt::Felt252Wrapper;
 use mp_state::StateChanges;
+use sp_runtime::Storage;
 use starknet_api::core::{ClassHash, CompiledClassHash, ContractAddress, Nonce};
 use starknet_api::hash::StarkFelt;
 use starknet_api::state::StorageKey;
+use starknet_core::types::contract;
 use starknet_core::types::{BlockId, BlockTag};
 use starknet_crypto::FieldElement;
 
@@ -23,9 +26,11 @@ use crate::{Config, Pallet};
 /// and not an extra layer that would add overhead.
 /// We don't implement those traits directly on the pallet to avoid compilation problems.
 pub struct BlockifierStateAdapter<T: Config> {
-    storage_update: BTreeMap<ContractAddress, Vec<(StorageKey, StarkFelt)>>,
-    class_hash_update: usize,
-    compiled_class_hash_update: usize,
+    storage_update: HashMap<(ContractAddress, StorageKey), StarkFelt>,
+    nonce_update: HashMap<ContractAddress, Nonce>,
+    class_hash_update: HashMap<ContractAddress, ClassHash>,
+    compiled_class_hash_update: HashMap<ClassHash, CompiledClassHash>,
+    contract_class_update: HashMap<ClassHash, ContractClass>,
     _phantom: PhantomData<T>,
 }
 
@@ -35,22 +40,25 @@ where
 {
     fn count_state_changes(&self) -> StateChangesCount {
         let keys = self.storage_update.keys();
-        let n_contract_updated = BTreeSet::from_iter(keys.clone()).len();
+        let n_storage_updates = keys.len();
+        let contract_updated = keys.map(|(contract_address, _)| contract_address).collect::<HashSet<_>>();
         StateChangesCount {
-            n_modified_contracts: n_contract_updated,
-            n_storage_updates: keys.len(),
-            n_class_hash_updates: self.class_hash_update,
-            n_compiled_class_hash_updates: self.compiled_class_hash_update,
+            n_modified_contracts: contract_updated.len(),
+            n_storage_updates,
+            n_class_hash_updates: self.class_hash_update.len(),
+            n_compiled_class_hash_updates: self.compiled_class_hash_update.len(),
         }
     }
 }
 
 impl<T: Config> Default for BlockifierStateAdapter<T> {
     fn default() -> Self {
         Self {
-            storage_update: BTreeMap::default(),
-            class_hash_update: usize::default(),
-            compiled_class_hash_update: usize::default(),
+            storage_update: HashMap::default(),
+            nonce_update: HashMap::default(),
+            class_hash_update: HashMap::default(),
+            compiled_class_hash_update: HashMap::default(),
+            contract_class_update: HashMap::default(),
             _phantom: PhantomData,
         }
     }
@@ -62,27 +70,41 @@ impl<T: Config> StateReader for BlockifierStateAdapter<T> {
 
         match search {
             Ok(Some(value)) => Ok(StarkFelt(value.to_bytes_be())),
-            _ => Err(StateError::StateReadError(format!(
-                "Failed to retrieve storage value for contract {} at key {}",
-                contract_address.0.0, key.0.0
-            ))),
+            None => Ok(StarkFelt::default()),
+            // _ => Err(StateError::StateReadError(format!(
+            //      "Failed to retrieve storage value for contract {} at key {}",
+            //      contract_address.0.0, key.0.0
+            // ))),
         }
     }
 
     fn get_nonce_at(&self, contract_address: ContractAddress) -> StateResult<Nonce> {
-        Ok(Pallet::<T>::nonce(contract_address))
+        Ok(self.nonce_update.get(&contract_address).cloned().unwrap_or_else(|| Pallet::<T>::nonce(contract_address)))
     }
 
     fn get_class_hash_at(&self, contract_address: ContractAddress) -> StateResult<ClassHash> {
-        Ok(Pallet::<T>::contract_class_hash_by_address(contract_address))
+        Ok(self
+            .class_hash_update
+            .get(&contract_address)
+            .cloned()
+            .unwrap_or_else(|| Pallet::<T>::contract_class_hash_by_address(contract_address)))
     }
 
     fn get_compiled_contract_class(&self, class_hash: ClassHash) -> StateResult<ContractClass> {
-        Pallet::<T>::contract_class_by_class_hash(class_hash).ok_or(StateError::UndeclaredClassHash(class_hash))
+        match self.contract_class_update.get(&class_hash) {
+            Some(contract_class) => Ok(contract_class.clone()),
+            None => {
+                Pallet::<T>::contract_class_by_class_hash(class_hash).ok_or(StateError::UndeclaredClassHash(class_hash))
+            }
+        }
     }
 
     fn get_compiled_class_hash(&self, class_hash: ClassHash) -> StateResult<CompiledClassHash> {
-        Pallet::<T>::compiled_class_hash_by_class_hash(class_hash).ok_or(StateError::UndeclaredClassHash(class_hash))
+        match self.compiled_class_hash_update.get(&class_hash) {
+            Some(compiled_class_hash) => Ok(compiled_class_hash.clone()),
+            None => Pallet::<T>::compiled_class_hash_by_class_hash(class_hash)
+                .ok_or(StateError::UndeclaredClassHash(class_hash)),
+        }
     }
 }
 
@@ -93,36 +115,32 @@ impl<T: Config> State for BlockifierStateAdapter<T> {
         key: StorageKey,
         value: StarkFelt,
     ) -> StateResult<()> {
-        self.storage_update.insert(contract_address, vec![(key, value)]);
-        let _ = StorageHandler::contract_storage_mut(BlockId::Tag(BlockTag::Latest)).unwrap().insert(
-            &contract_address,
-            &key,
-            value,
-        );
-
+        self.storage_update.insert((contract_address, key), value);
+
         Ok(())
     }
 
     fn increment_nonce(&mut self, contract_address: ContractAddress) -> StateResult<()> {
-        let current_nonce = Pallet::<T>::nonce(contract_address);
-        let current_nonce: FieldElement = Felt252Wrapper::from(current_nonce.0).into();
-        let new_nonce: Nonce = Felt252Wrapper(current_nonce + FieldElement::ONE).into();
-
-        crate::Nonces::<T>::insert(contract_address, new_nonce);
+        let nonce = self
+            .nonce_update
+            .get(&contract_address)
+            .cloned()
+            .unwrap_or_else(|| Pallet::<T>::nonce(contract_address))
+            .try_increment()
+            .map_err(|e| StateError::StarknetApiError(e))?;
+        self.nonce_update.insert(contract_address, nonce);
 
         Ok(())
     }
 
     fn set_class_hash_at(&mut self, contract_address: ContractAddress, class_hash: ClassHash) -> StateResult<()> {
-        self.class_hash_update += 1;
-
-        crate::ContractClassHashes::<T>::insert(contract_address, class_hash);
+        self.class_hash_update.insert(contract_address, class_hash);
 
         Ok(())
     }
 
     fn set_contract_class(&mut self, class_hash: ClassHash, contract_class: ContractClass) -> StateResult<()> {
-        crate::ContractClasses::<T>::insert(class_hash, contract_class);
+        self.contract_class_update.insert(class_hash, contract_class);
 
         Ok(())
     }
@@ -132,8 +150,7 @@ impl<T: Config> State for BlockifierStateAdapter<T> {
         class_hash: ClassHash,
         compiled_class_hash: CompiledClassHash,
     ) -> StateResult<()> {
-        self.compiled_class_hash_update += 1;
-        crate::CompiledClassHashes::<T>::insert(class_hash, compiled_class_hash);
+        self.compiled_class_hash_update.insert(class_hash, compiled_class_hash);
 
         Ok(())
     }

crates/pallets/starknet/src/lib.rs

@@ -48,8 +48,7 @@ pub mod blockifier_state_adapter;
 pub mod genesis_loader;
 /// Simulation, estimations and execution trace logic.
 pub mod simulations;
-/// Transaction validation logic.
-pub mod transaction_validation;
+
 /// The Starknet pallet's runtime custom types.
 pub mod types;
 
@@ -97,7 +96,6 @@ use starknet_api::transaction::{
     L1HandlerTransaction, MessageToL1, Transaction, TransactionHash,
 };
 use starknet_crypto::FieldElement;
-use transaction_validation::TxPriorityInfo;
 
 use crate::alloc::string::ToString;
 use crate::types::{CasmClassHash, ContractStorageKey, SierraClassHash, StorageSlot};