From 6e49be2702ef52da604613cef949f33c8e8a70ee Mon Sep 17 00:00:00 2001
From: Skyth <19259897+blueskythlikesclouds@users.noreply.github.com>
Date: Thu, 17 Oct 2024 17:39:31 +0300
Subject: [PATCH] Allow midasm hook return behavior to be configured.

---
 PowerRecomp/recompiler.cpp        | 58 +++++++++++++++++++++++++++++--
 PowerRecomp/recompiler_config.cpp | 28 +++++++++++++--
 PowerRecomp/recompiler_config.h   |  8 +++++
 3 files changed, 89 insertions(+), 5 deletions(-)

diff --git a/PowerRecomp/recompiler.cpp b/PowerRecomp/recompiler.cpp
index a868b8b..7bfc65c 100644
--- a/PowerRecomp/recompiler.cpp
+++ b/PowerRecomp/recompiler.cpp
@@ -328,7 +328,14 @@ bool Recompiler::Recompile(
     auto midAsmHook = config.midAsmHooks.find(base);
     if (midAsmHook != config.midAsmHooks.end())
     {
-        print("\t{}(", midAsmHook->second.name);
+        bool returnsBool = midAsmHook->second.returnOnFalse || midAsmHook->second.returnOnTrue ||
+            midAsmHook->second.jumpAddressOnFalse != NULL || midAsmHook->second.jumpAddressOnTrue != NULL;
+
+        print("\t");
+        if (returnsBool)
+            print("if (");
+
+        print("{}(", midAsmHook->second.name);
         for (auto& reg : midAsmHook->second.registers)
         {
             if (out.back() != '(')
@@ -367,7 +374,35 @@ bool Recompiler::Recompile(
             }
         }
 
-        println(");");
+        if (returnsBool)
+        {
+            println(")) {{");
+
+            if (midAsmHook->second.returnOnTrue)
+                println("\t\treturn;");
+            else if (midAsmHook->second.jumpAddressOnTrue != NULL)
+                println("\t\tgoto loc_{:X};", midAsmHook->second.jumpAddressOnTrue);
+
+            println("\t}}");
+
+            println("\telse {{");
+
+            if (midAsmHook->second.returnOnFalse)
+                println("\t\treturn;");
+            else if (midAsmHook->second.jumpAddressOnFalse != NULL)
+                println("\t\tgoto loc_{:X};", midAsmHook->second.jumpAddressOnFalse);
+
+            println("\t}}");
+        }
+        else
+        {
+            println(");");
+
+            if (midAsmHook->second.ret)
+                println("\treturn;");
+            else if (midAsmHook->second.jumpAddress != NULL)
+                println("\tgoto loc_{:X};", midAsmHook->second.jumpAddress);
+        }
     }
 
     int id = insn.opcode->id;
@@ -2141,7 +2176,17 @@ bool Recompiler::Recompile(const Function& fn)
         auto midAsmHook = config.midAsmHooks.find(addr);
         if (midAsmHook != config.midAsmHooks.end())
         {
-            print("extern void {}(", midAsmHook->second.name);
+            if (midAsmHook->second.returnOnFalse || midAsmHook->second.returnOnTrue ||
+                midAsmHook->second.jumpAddressOnFalse != NULL || midAsmHook->second.jumpAddressOnTrue != NULL)
+            {
+                print("extern bool ");
+            }
+            else
+            {
+                print("extern void ");
+            }
+
+            print("{}(", midAsmHook->second.name);
             for (auto& reg : midAsmHook->second.registers)
             {
                 if (out.back() != '(')
@@ -2178,6 +2223,13 @@ bool Recompiler::Recompile(const Function& fn)
             }
 
             println(");\n");
+
+            if (midAsmHook->second.jumpAddress != NULL)
+                labels.emplace(midAsmHook->second.jumpAddress);       
+            if (midAsmHook->second.jumpAddressOnTrue != NULL)
+                labels.emplace(midAsmHook->second.jumpAddressOnTrue);    
+            if (midAsmHook->second.jumpAddressOnFalse != NULL)
+                labels.emplace(midAsmHook->second.jumpAddressOnFalse);
         }
     }
 
diff --git a/PowerRecomp/recompiler_config.cpp b/PowerRecomp/recompiler_config.cpp
index 05bb6ef..47a6aa2 100644
--- a/PowerRecomp/recompiler_config.cpp
+++ b/PowerRecomp/recompiler_config.cpp
@@ -82,10 +82,34 @@ void RecompilerConfig::Load(const std::string_view& configFilePath)
 
             RecompilerMidAsmHook midAsmHook;
             midAsmHook.name = *table["name"].value<std::string>();
-            for (auto& reg : *table["registers"].as_array())
+            if (auto registerArray = table["registers"].as_array())
             {
-                midAsmHook.registers.push_back(*reg.value<std::string>());
+                for (auto& reg : *registerArray)
+                    midAsmHook.registers.push_back(*reg.value<std::string>());
             }
+
+            midAsmHook.ret = table["return"].value_or(false);
+            midAsmHook.returnOnTrue = table["return_on_true"].value_or(false);
+            midAsmHook.returnOnFalse = table["return_on_false"].value_or(false);
+
+            midAsmHook.jumpAddress = table["jump_address"].value_or(0u);
+            midAsmHook.jumpAddressOnTrue = table["jump_address_on_true"].value_or(0u);
+            midAsmHook.jumpAddressOnFalse = table["jump_address_on_false"].value_or(0u);
+
+            if ((midAsmHook.ret && midAsmHook.jumpAddress != NULL) ||
+                (midAsmHook.returnOnTrue && midAsmHook.jumpAddressOnTrue != NULL) ||
+                (midAsmHook.returnOnFalse && midAsmHook.jumpAddressOnFalse != NULL))
+            {
+                std::println("{}: can't return and jump at the same time", midAsmHook.name);
+            }
+
+            if ((midAsmHook.ret || midAsmHook.jumpAddress != NULL) &&
+                (midAsmHook.returnOnFalse != NULL || midAsmHook.returnOnTrue != NULL ||
+                    midAsmHook.jumpAddressOnFalse != NULL || midAsmHook.jumpAddressOnTrue != NULL))
+            {
+                std::println("{}: can't mix direct and conditional return/jump", midAsmHook.name);
+            }
+
             midAsmHooks.emplace(*table["address"].value<uint32_t>(), std::move(midAsmHook));
         }
     }
diff --git a/PowerRecomp/recompiler_config.h b/PowerRecomp/recompiler_config.h
index 9bd238d..e20936e 100644
--- a/PowerRecomp/recompiler_config.h
+++ b/PowerRecomp/recompiler_config.h
@@ -10,6 +10,14 @@ struct RecompilerMidAsmHook
 {
     std::string name;
     std::vector<std::string> registers;
+
+    bool ret = false;
+    bool returnOnTrue = false;
+    bool returnOnFalse = false;
+
+    uint32_t jumpAddress = 0;
+    uint32_t jumpAddressOnTrue = 0;
+    uint32_t jumpAddressOnFalse = 0;
 };
 
 struct RecompilerConfig