From 2484265f27f1f35874dac3e0c9a06df8127aab9d Mon Sep 17 00:00:00 2001
From: Yair Bakalchuk <yair@starkware.co>
Date: Fri, 29 Sep 2023 09:23:30 +0300
Subject: [PATCH] fix(execution): fetch the class from the right block

---
 crates/papyrus_rpc/src/v0_4_0/api/mod.rs | 41 ++++++++++++------------
 1 file changed, 21 insertions(+), 20 deletions(-)

diff --git a/crates/papyrus_rpc/src/v0_4_0/api/mod.rs b/crates/papyrus_rpc/src/v0_4_0/api/mod.rs
index ecf7718985..a77e00af8b 100644
--- a/crates/papyrus_rpc/src/v0_4_0/api/mod.rs
+++ b/crates/papyrus_rpc/src/v0_4_0/api/mod.rs
@@ -312,16 +312,7 @@ pub(crate) fn stored_txn_to_executable_txn(
             let class_hash = value.class_hash;
             Ok(ExecutableTransactionInput::DeclareV0(
                 value,
-                storage_txn
-                    .get_state_reader()
-                    .map_err(internal_server_error)?
-                    .get_deprecated_class_definition_at(state_number, &class_hash)
-                    .map_err(internal_server_error)?
-                    .ok_or_else(|| {
-                        internal_server_error(format!(
-                            "Missing deprecated class definition of {class_hash}."
-                        ))
-                    })?,
+                get_deprecated_class_for_re_execution(storage_txn, state_number, class_hash)?,
             ))
         }
         starknet_api::transaction::Transaction::Declare(
@@ -331,16 +322,7 @@ pub(crate) fn stored_txn_to_executable_txn(
             let class_hash = value.class_hash;
             Ok(ExecutableTransactionInput::DeclareV1(
                 value,
-                storage_txn
-                    .get_state_reader()
-                    .map_err(internal_server_error)?
-                    .get_deprecated_class_definition_at(state_number, &class_hash)
-                    .map_err(internal_server_error)?
-                    .ok_or_else(|| {
-                        internal_server_error(format!(
-                            "Missing deprecated class definition of {class_hash}."
-                        ))
-                    })?,
+                get_deprecated_class_for_re_execution(storage_txn, state_number, class_hash)?,
             ))
         }
         starknet_api::transaction::Transaction::Declare(
@@ -391,6 +373,25 @@ pub(crate) fn stored_txn_to_executable_txn(
     }
 }
 
+// For re-execution (traceTransaction, traceBlockTransactions) we need to get the class definition
+// of declare transactions from the storage before the execution. They are stored in the state after
+// the block in which they appeared, so we need to get it from the state after given block.
+fn get_deprecated_class_for_re_execution(
+    storage_txn: &StorageTxn<'_, RO>,
+    state_number: StateNumber,
+    class_hash: ClassHash,
+) -> Result<starknet_api::deprecated_contract_class::ContractClass, ErrorObjectOwned> {
+    let state_number_after_block = StateNumber::right_after_block(state_number.block_after());
+    storage_txn
+        .get_state_reader()
+        .map_err(internal_server_error)?
+        .get_deprecated_class_definition_at(state_number_after_block, &class_hash)
+        .map_err(internal_server_error)?
+        .ok_or_else(|| {
+            internal_server_error(format!("Missing deprecated class definition of {class_hash}."))
+        })
+}
+
 impl TryFrom<BroadcastedDeclareTransaction> for ExecutableTransactionInput {
     type Error = ErrorObjectOwned;
     fn try_from(value: BroadcastedDeclareTransaction) -> Result<Self, Self::Error> {