From 49bf144b0d27870d5a08c9d6d0539acf748ea32b Mon Sep 17 00:00:00 2001
From: LittleCube <littlecubehax@gmail.com>
Date: Sat, 4 Jan 2025 22:10:29 -0500
Subject: [PATCH] Add `TRACE_RETURN` (#117)

---
 LiveRecomp/live_generator.cpp        |  2 +-
 include/recompiler/generator.h       |  4 ++--
 include/recompiler/live_recompiler.h |  2 +-
 src/cgenerator.cpp                   |  5 ++++-
 src/recompilation.cpp                | 10 +++++-----
 5 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/LiveRecomp/live_generator.cpp b/LiveRecomp/live_generator.cpp
index 76b2115..cb733f1 100644
--- a/LiveRecomp/live_generator.cpp
+++ b/LiveRecomp/live_generator.cpp
@@ -1590,7 +1590,7 @@ void N64Recomp::LiveGenerator::emit_switch_close() const {
     // Nothing to do here, the jump table is built in emit_switch.
 }
 
-void N64Recomp::LiveGenerator::emit_return() const {
+void N64Recomp::LiveGenerator::emit_return(const Context& context) const {
     sljit_emit_return_void(compiler);
 }
 
diff --git a/include/recompiler/generator.h b/include/recompiler/generator.h
index 0ffde0b..4d36f6c 100644
--- a/include/recompiler/generator.h
+++ b/include/recompiler/generator.h
@@ -48,7 +48,7 @@ namespace N64Recomp {
         virtual void emit_case(int case_index, const std::string& target_label) const = 0;
         virtual void emit_switch_error(uint32_t instr_vram, uint32_t jtbl_vram) const = 0;
         virtual void emit_switch_close() const = 0;
-        virtual void emit_return() const = 0;
+        virtual void emit_return(const Context& context) const = 0;
         virtual void emit_check_fr(int fpr) const = 0;
         virtual void emit_check_nan(int fpr, bool is_double) const = 0;
         virtual void emit_cop0_status_read(int reg) const = 0;
@@ -85,7 +85,7 @@ namespace N64Recomp {
         void emit_case(int case_index, const std::string& target_label) const final;
         void emit_switch_error(uint32_t instr_vram, uint32_t jtbl_vram) const final;
         void emit_switch_close() const final;
-        void emit_return() const final;
+        void emit_return(const Context& context) const final;
         void emit_check_fr(int fpr) const final;
         void emit_check_nan(int fpr, bool is_double) const final;
         void emit_cop0_status_read(int reg) const final;
diff --git a/include/recompiler/live_recompiler.h b/include/recompiler/live_recompiler.h
index 1b92d95..cc7670f 100644
--- a/include/recompiler/live_recompiler.h
+++ b/include/recompiler/live_recompiler.h
@@ -109,7 +109,7 @@ namespace N64Recomp {
         void emit_case(int case_index, const std::string& target_label) const final;
         void emit_switch_error(uint32_t instr_vram, uint32_t jtbl_vram) const final;
         void emit_switch_close() const final;
-        void emit_return() const final;
+        void emit_return(const Context& context) const final;
         void emit_check_fr(int fpr) const final;
         void emit_check_nan(int fpr, bool is_double) const final;
         void emit_cop0_status_read(int reg) const final;
diff --git a/src/cgenerator.cpp b/src/cgenerator.cpp
index 596ad60..1ac63a8 100644
--- a/src/cgenerator.cpp
+++ b/src/cgenerator.cpp
@@ -476,7 +476,10 @@ void N64Recomp::CGenerator::emit_switch_error(uint32_t instr_vram, uint32_t jtbl
     fmt::print(output_file, "default: switch_error(__func__, 0x{:08X}, 0x{:08X});\n", instr_vram, jtbl_vram);
 }
 
-void N64Recomp::CGenerator::emit_return() const {
+void N64Recomp::CGenerator::emit_return(const Context& context) const {
+    if (context.trace_mode) {
+        fmt::print(output_file, "TRACE_RETURN()\n    ");
+    }
     fmt::print(output_file, "return;\n");
 }
 
diff --git a/src/recompilation.cpp b/src/recompilation.cpp
index 8ea9bdd..5761b79 100644
--- a/src/recompilation.cpp
+++ b/src/recompilation.cpp
@@ -238,7 +238,7 @@ bool process_instruction(GeneratorType& generator, const N64Recomp::Context& con
             return false;
         }
         print_indent();
-        generator.emit_return();
+        generator.emit_return(context);
         print_link_branch();
         return true;
     };
@@ -363,7 +363,7 @@ bool process_instruction(GeneratorType& generator, const N64Recomp::Context& con
                     return false;
                 }
                 print_indent();
-                generator.emit_return();
+                generator.emit_return(context);
                 // TODO check if this branch close should exist.
                 // print_indent();
                 // generator.emit_branch_close();
@@ -512,7 +512,7 @@ bool process_instruction(GeneratorType& generator, const N64Recomp::Context& con
                     return false;
                 }
                 print_indent();
-                generator.emit_return();
+                generator.emit_return(context);
             }
             else {
                 fmt::print(stderr, "Unhandled branch in {} at 0x{:08X} to 0x{:08X}\n", func.name, instr_vram, branch_target);
@@ -552,7 +552,7 @@ bool process_instruction(GeneratorType& generator, const N64Recomp::Context& con
             fmt::print("[Info] Indirect tail call in {}\n", func.name);
             print_func_call_by_register(rs);
             print_indent();
-            generator.emit_return();
+            generator.emit_return(context);
             break;
         }
         break;
@@ -561,7 +561,7 @@ bool process_instruction(GeneratorType& generator, const N64Recomp::Context& con
         generator.emit_syscall(instr_vram);
         // syscalls don't link, so treat it like a tail call
         print_indent();
-        generator.emit_return();
+        generator.emit_return(context);
         break;
     case InstrId::cpu_break:
         print_indent();