From 2af6f2d161f755ca68f8baeb4dc5b69ef58a862c Mon Sep 17 00:00:00 2001
From: Wiseguy <68165316+Mr-Wiseguy@users.noreply.github.com>
Date: Thu, 30 Jan 2025 23:48:20 -0500
Subject: [PATCH] Implement shim function generation (#128)

---
 LiveRecomp/live_generator.cpp        | 26 ++++++++++++++++++++++++++
 include/recomp.h                     |  4 ++++
 include/recompiler/live_recompiler.h | 10 ++++++++++
 3 files changed, 40 insertions(+)

diff --git a/LiveRecomp/live_generator.cpp b/LiveRecomp/live_generator.cpp
index 241dd8d..b93d7ef 100644
--- a/LiveRecomp/live_generator.cpp
+++ b/LiveRecomp/live_generator.cpp
@@ -1896,3 +1896,29 @@ bool N64Recomp::recompile_function_live(LiveGenerator& generator, const Context&
     return recompile_function_custom(generator, context, function_index, output_file, static_funcs_out, tag_reference_relocs);
 }
 
+N64Recomp::ShimFunction::ShimFunction(recomp_func_ext_t* to_shim, uintptr_t value) {
+    sljit_compiler* compiler = sljit_create_compiler(nullptr);
+
+    // Create the function.
+    sljit_label* func_label = sljit_emit_label(compiler);
+    sljit_emit_enter(compiler, 0, SLJIT_ARGS2V(P_R, P_R), 3, 0, 0);
+
+    // Move the provided value into the third argument.
+    sljit_emit_op1(compiler, SLJIT_MOV, SLJIT_R2, 0, SLJIT_IMM, sljit_sw(value));
+
+    // Tail call the provided function.
+    sljit_emit_icall(compiler, SLJIT_CALL | SLJIT_CALL_RETURN, SLJIT_ARGS3V(P, P, W), SLJIT_IMM, sljit_sw(to_shim));
+
+    // Generate the function's code and get the address to the function.
+    code = sljit_generate_code(compiler, 0, nullptr);
+    func = reinterpret_cast<recomp_func_t*>(sljit_get_label_addr(func_label));
+
+    // Cleanup.
+    sljit_free_compiler(compiler);
+}
+
+N64Recomp::ShimFunction::~ShimFunction() {
+    sljit_free_code(code, nullptr);
+    code = nullptr;
+    func = nullptr;
+}
diff --git a/include/recomp.h b/include/recomp.h
index 578d0f9..73f3e73 100644
--- a/include/recomp.h
+++ b/include/recomp.h
@@ -439,7 +439,11 @@ gpr cop0_status_read(recomp_context* ctx);
 void switch_error(const char* func, uint32_t vram, uint32_t jtbl);
 void do_break(uint32_t vram);
 
+// The function signature for all recompiler output functions.
 typedef void (recomp_func_t)(uint8_t* rdram, recomp_context* ctx);
+// The function signature for special functions that need a third argument.
+// These get called via generated shims to allow providing some information about the caller, such as mod id.
+typedef void (recomp_func_ext_t)(uint8_t* rdram, recomp_context* ctx, uintptr_t arg);
 
 recomp_func_t* get_function(int32_t vram);
 
diff --git a/include/recompiler/live_recompiler.h b/include/recompiler/live_recompiler.h
index 6f763ca..b85d894 100644
--- a/include/recompiler/live_recompiler.h
+++ b/include/recompiler/live_recompiler.h
@@ -141,6 +141,16 @@ namespace N64Recomp {
 
     void live_recompiler_init();
     bool recompile_function_live(LiveGenerator& generator, const Context& context, size_t function_index, std::ostream& output_file, std::span<std::vector<uint32_t>> static_funcs_out, bool tag_reference_relocs);
+
+    class ShimFunction {
+    private:
+        void* code;
+        recomp_func_t* func;
+    public:
+        ShimFunction(recomp_func_ext_t* to_shim, uintptr_t value);
+        ~ShimFunction();
+        recomp_func_t* get_func() { return func; }
+    };
 }
 
 #endif
\ No newline at end of file