From 7fff40c59b77cece7c1a6ae68f673b7f6b57da8c Mon Sep 17 00:00:00 2001
From: Ammar Arif <evergreenkary@gmail.com>
Date: Mon, 25 Mar 2024 22:28:55 +0900
Subject: [PATCH] fix(katana): add missing contract address in `CallInfo`

---
 .../katana/executor/src/implementation/blockifier/utils.rs  | 1 +
 crates/katana/executor/src/implementation/sir/utils.rs      | 1 +
 crates/katana/primitives/src/trace.rs                       | 2 ++
 crates/katana/rpc/rpc-types/src/trace.rs                    | 6 ++----
 4 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/crates/katana/executor/src/implementation/blockifier/utils.rs b/crates/katana/executor/src/implementation/blockifier/utils.rs
index 2c5e5fc3d2..7016583079 100644
--- a/crates/katana/executor/src/implementation/blockifier/utils.rs
+++ b/crates/katana/executor/src/implementation/blockifier/utils.rs
@@ -533,6 +533,7 @@ fn to_call_info(call_info: CallInfo) -> katana_primitives::trace::CallInfo {
     };
 
     katana_primitives::trace::CallInfo {
+        contract_address: to_address(call_info.call.storage_address),
         caller_address: to_address(call_info.call.caller_address),
         call_type: match call_info.call.call_type {
             CallType::Call => katana_primitives::trace::CallType::Call,
diff --git a/crates/katana/executor/src/implementation/sir/utils.rs b/crates/katana/executor/src/implementation/sir/utils.rs
index 9db7edb6b2..9b296ac68f 100644
--- a/crates/katana/executor/src/implementation/sir/utils.rs
+++ b/crates/katana/executor/src/implementation/sir/utils.rs
@@ -567,6 +567,7 @@ fn from_sir_call_info(call_info: CallInfo) -> katana_primitives::trace::CallInfo
     };
 
     katana_primitives::trace::CallInfo {
+        contract_address: to_address(&call_info.contract_address),
         caller_address: to_address(&call_info.caller_address),
         call_type: match call_info.call_type {
             Some(CallType::Call) => katana_primitives::trace::CallType::Call,
diff --git a/crates/katana/primitives/src/trace.rs b/crates/katana/primitives/src/trace.rs
index 85d78c03f1..1349e10358 100644
--- a/crates/katana/primitives/src/trace.rs
+++ b/crates/katana/primitives/src/trace.rs
@@ -54,6 +54,8 @@ pub struct CallInfo {
     pub caller_address: ContractAddress,
     /// The call type.
     pub call_type: CallType,
+    /// The contract address.
+    pub contract_address: ContractAddress,
     /// The address where the code is being executed.
     /// Optional, since there is no address to the code implementation in a delegate call.
     pub code_address: Option<ContractAddress>,
diff --git a/crates/katana/rpc/rpc-types/src/trace.rs b/crates/katana/rpc/rpc-types/src/trace.rs
index f49b1305f4..3bea0d6edb 100644
--- a/crates/katana/rpc/rpc-types/src/trace.rs
+++ b/crates/katana/rpc/rpc-types/src/trace.rs
@@ -33,8 +33,7 @@ impl From<CallInfo> for FunctionInvocation {
                 order: m.order,
                 payload: m.payload,
                 to_address: m.to_address,
-                // TODO: fix missing storage address in `TxExecInfo`
-                from_address: Default::default(),
+                from_address: m.from_address.into(),
             })
             .collect();
 
@@ -65,9 +64,8 @@ impl From<CallInfo> for FunctionInvocation {
             execution_resources,
             result: info.retdata,
             calldata: info.calldata,
-            // TODO: fix missing storage address in `CallInfo`
-            contract_address: Default::default(),
             caller_address: info.caller_address.into(),
+            contract_address: info.contract_address.into(),
             entry_point_selector: info.entry_point_selector,
             // See <https://github.com/starkware-libs/blockifier/blob/cb464f5ac2ada88f2844d9f7d62bd6732ceb5b2c/crates/blockifier/src/execution/call_info.rs#L220>
             class_hash: info.class_hash.expect("Class hash mut be set after execution"),