From f814b2b42fbace83e5365f058f3c6744e9ed1d5d Mon Sep 17 00:00:00 2001
From: orizi <104711814+orizi@users.noreply.github.com>
Date: Wed, 28 Jun 2023 10:34:25 +0300
Subject: [PATCH] Properly restoring contract addresses on revert. (#3534)

---
 crates/cairo-lang-runner/src/casm_run/mod.rs | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/crates/cairo-lang-runner/src/casm_run/mod.rs b/crates/cairo-lang-runner/src/casm_run/mod.rs
index e748ae5bcae..5fe6aeaf2a3 100644
--- a/crates/cairo-lang-runner/src/casm_run/mod.rs
+++ b/crates/cairo-lang-runner/src/casm_run/mod.rs
@@ -613,6 +613,9 @@ impl HintProcessor for CairoHintProcessor<'_> {
                         // Call the function.
                         let Some(entry_point) = contract_info.externals.get(&selector) else
                         {
+                            // Restore the contract addresses in the context.
+                            self.starknet_state.exec_info.caller_address = old_caller_address;
+                            self.starknet_state.exec_info.contract_address = old_contract_address;
                             fail_syscall!(b"ENTRYPOINT_NOT_FOUND");
                         };
                         let function = runner
@@ -636,6 +639,10 @@ impl HintProcessor for CairoHintProcessor<'_> {
                                 read_array_result_as_vec(&res.memory, &value)
                             }
                             RunResultValue::Panic(mut panic_data) => {
+                                // Restore the contract addresses in the context.
+                                self.starknet_state.exec_info.caller_address = old_caller_address;
+                                self.starknet_state.exec_info.contract_address =
+                                    old_contract_address;
                                 fail_syscall!(panic_data, b"ENTRYPOINT_FAILED");
                             }
                         };