perf: add caching to optimize tx trace creation (#1422)

Co-authored-by: fishseabowl <[email protected]>

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Next release
 
+- dev: optimize tx trace creation
 - dev: make Madara std compatible
 - CI: fix taplo version
 - chore: add cache usage for `getEvents` and `getTransactionReceipt`

crates/client/rpc/src/trace_api.rs

@@ -1,3 +1,5 @@
+use std::collections::HashMap;
+
 use blockifier::execution::entry_point::CallInfo;
 use blockifier::transaction::errors::TransactionExecutionError;
 use blockifier::transaction::objects::TransactionExecutionInfo;
@@ -17,6 +19,7 @@ use sc_transaction_pool_api::TransactionPool;
 use sp_api::ProvideRuntimeApi;
 use sp_blockchain::HeaderBackend;
 use sp_runtime::traits::Block as BlockT;
+use starknet_api::api_core::ContractAddress;
 use starknet_core::types::{
     BlockId, BroadcastedTransaction, DeclareTransactionTrace, DeployAccountTransactionTrace, ExecuteInvocation,
     FeeEstimate, InvokeTransactionTrace, RevertedInvocation, SimulatedTransaction, SimulationFlag, TransactionTrace,
@@ -153,14 +156,17 @@ fn try_get_funtion_invocation_from_call_info<B: BlockT>(
     storage_override: &dyn StorageOverride<B>,
     substrate_block_hash: B::Hash,
     call_info: &CallInfo,
+    class_hash_cache: &mut HashMap<ContractAddress, FieldElement>,
 ) -> Result<starknet_core::types::FunctionInvocation, TryFuntionInvocationFromCallInfoError> {
     let messages = collect_call_info_ordered_messages(call_info);
     let events = blockifier_to_starknet_rs_ordered_events(&call_info.execution.events);
 
     let inner_calls = call_info
         .inner_calls
         .iter()
-        .map(|call| try_get_funtion_invocation_from_call_info(storage_override, substrate_block_hash, call))
+        .map(|call| {
+            try_get_funtion_invocation_from_call_info(storage_override, substrate_block_hash, call, class_hash_cache)
+        })
         .collect::<Result<_, _>>()?;
 
     call_info.get_sorted_l2_to_l1_payloads_length()?;
@@ -186,12 +192,18 @@ fn try_get_funtion_invocation_from_call_info<B: BlockT>(
     // address". We have to do this decution ourselves here
     let class_hash = if let Some(class_hash) = call_info.call.class_hash {
         class_hash.0.into()
+    } else if let Some(cached_hash) = class_hash_cache.get(&call_info.call.storage_address) {
+        *cached_hash
     } else {
-        let class_hash = storage_override
+        // Compute and cache the class hash
+        let computed_hash = storage_override
             .contract_class_hash_by_address(substrate_block_hash, call_info.call.storage_address)
             .ok_or_else(|| TryFuntionInvocationFromCallInfoError::ContractNotFound)?;
 
-        FieldElement::from_byte_slice_be(class_hash.0.bytes()).unwrap()
+        let computed_hash = FieldElement::from_byte_slice_be(computed_hash.0.bytes()).unwrap();
+        class_hash_cache.insert(call_info.call.storage_address, computed_hash);
+
+        computed_hash
     };
 
     Ok(starknet_core::types::FunctionInvocation {
@@ -218,6 +230,7 @@ fn tx_execution_infos_to_simulated_transactions<B: BlockT>(
     >,
 ) -> Result<Vec<SimulatedTransaction>, ConvertCallInfoToExecuteInvocationError> {
     let mut results = vec![];
+    let mut class_hash_cache: HashMap<ContractAddress, FieldElement> = HashMap::new();
     for (tx_type, res) in tx_types.iter().zip(transaction_execution_results.iter()) {
         match res {
             Ok(tx_exec_info) => {
@@ -227,7 +240,12 @@ fn tx_execution_infos_to_simulated_transactions<B: BlockT>(
                     .validate_call_info
                     .as_ref()
                     .map(|call_info| {
-                        try_get_funtion_invocation_from_call_info(storage_override, substrate_block_hash, call_info)
+                        try_get_funtion_invocation_from_call_info(
+                            storage_override,
+                            substrate_block_hash,
+                            call_info,
+                            &mut class_hash_cache,
+                        )
                     })
                     .transpose()?;
                 // If simulated with `SimulationFlag::SkipFeeCharge` this will be `None`
@@ -236,7 +254,12 @@ fn tx_execution_infos_to_simulated_transactions<B: BlockT>(
                     .fee_transfer_call_info
                     .as_ref()
                     .map(|call_info| {
-                        try_get_funtion_invocation_from_call_info(storage_override, substrate_block_hash, call_info)
+                        try_get_funtion_invocation_from_call_info(
+                            storage_override,
+                            substrate_block_hash,
+                            call_info,
+                            &mut class_hash_cache,
+                        )
                     })
                     .transpose()?;
 
@@ -251,6 +274,7 @@ fn tx_execution_infos_to_simulated_transactions<B: BlockT>(
                                 substrate_block_hash,
                                 // Safe to unwrap because is only `None`  for `Declare` txs
                                 tx_exec_info.execute_call_info.as_ref().unwrap(),
+                                &mut class_hash_cache,
                             )?)
                         },
                         fee_transfer_invocation,
@@ -271,6 +295,7 @@ fn tx_execution_infos_to_simulated_transactions<B: BlockT>(
                                 substrate_block_hash,
                                 // Safe to unwrap because is only `None`  for `Declare` txs
                                 tx_exec_info.execute_call_info.as_ref().unwrap(),
+                                &mut class_hash_cache,
                             )?,
                             fee_transfer_invocation,
                             // TODO(#1291): Compute state diff correctly