From 10026e9a0419ff1cea1ddb80f5ccddf6a53818d3 Mon Sep 17 00:00:00 2001
From: TheBrokenRail <connor24nolan@live.com>
Date: Wed, 17 Jul 2024 06:40:04 -0400
Subject: [PATCH] Finishing Touches!

---
 data/function.cpp | 10 +---------
 data/function.h   | 41 ++++++++++++++++++++++++++---------------
 data/out.cpp      |  7 +------
 src/index.ts      |  5 +----
 src/struct.ts     | 25 +------------------------
 5 files changed, 30 insertions(+), 58 deletions(-)

diff --git a/data/function.cpp b/data/function.cpp
index 10cb655..2da4222 100644
--- a/data/function.cpp
+++ b/data/function.cpp
@@ -29,12 +29,4 @@ __Function<Ret(Args...)>::__Function(const char *const name_, __Function<Ret(Arg
     thunk(thunk_) {}
 
 // Thunks
-template <typename Ret, typename... Args>
-void __Function<Ret(Args...)>::enable_thunk(const thunk_enabler_t &thunk_enabler) {
-    if (enabled) {
-        ptr_type real_thunk = (ptr_type) thunk_enabler((void *) get(), (void *) thunk);
-        if (!is_virtual) {
-            func = real_thunk;
-        }
-    }
-}
\ No newline at end of file
+thunk_enabler_t thunk_enabler;
\ No newline at end of file
diff --git a/data/function.h b/data/function.h
index 7e2dc6a..ea66971 100644
--- a/data/function.h
+++ b/data/function.h
@@ -18,7 +18,7 @@ class __VirtualFunctionInfo {
 
 // Thunks
 typedef void *(*thunk_enabler_t)(void *target, void *thunk);
-void enable_all_thunks(const thunk_enabler_t &thunk_enabler);
+extern thunk_enabler_t thunk_enabler;
 
 // Function Information
 template <typename Ret, typename... Args>
@@ -40,25 +40,29 @@ public:
         if (!enabled) {
             return false;
         }
+        // Enable Thunk
+        enable_thunk();
         // Overwrite
         type original = get_thunk_target();
         thunk_target = [original, target](Args... args) {
-            return target(original, args...);
+            return target(original, std::forward<Args>(args)...);
         };
         return true;
     }
 
     // Getters
-    [[nodiscard]] ptr_type get_backup() const {
-        return backup;
-    }
-    [[nodiscard]] ptr_type get() const {
+    [[nodiscard]] ptr_type get(bool result_will_be_stored) {
         if (!enabled) {
             return nullptr;
-        } else if (is_virtual) {
-            return *get_vtable_addr();
         } else {
-            return std::get<ptr_type>(func);
+            if (result_will_be_stored) {
+                enable_thunk();
+            }
+            if (is_virtual) {
+                return *get_vtable_addr();
+            } else {
+                return std::get<ptr_type>(func);
+            }
         }
     }
     [[nodiscard]] ptr_type *get_vtable_addr() const {
@@ -68,14 +72,11 @@ public:
             return nullptr;
         }
     }
-    [[nodiscard]] const char *get_name() const {
-        return name;
-    }
     [[nodiscard]] type get_thunk_target() const {
         if (thunk_target) {
             return thunk_target;
         } else {
-            return get_backup();
+            return backup;
         }
     }
 
@@ -84,6 +85,7 @@ private:
     const bool is_virtual;
     std::variant<ptr_type, __VirtualFunctionInfo<ptr_type>> func;
 
+public:
     // State
     const bool enabled;
     const char *const name;
@@ -91,9 +93,18 @@ private:
     // Backup Of Original Function Pointer
     const ptr_type backup;
 
+private:
     // Thunk
     const ptr_type thunk;
     type thunk_target;
-    void enable_thunk(const thunk_enabler_t &thunk_enabler);
-    friend void enable_all_thunks(const thunk_enabler_t &);
+    bool thunk_enabled = false;
+    void enable_thunk() {
+        if (!thunk_enabled && enabled) {
+            ptr_type real_thunk = (ptr_type) thunk_enabler((void *) backup, (void *) thunk);
+            if (!is_virtual) {
+                func = real_thunk;
+            }
+            thunk_enabled = true;
+        }
+    }
 };
\ No newline at end of file
diff --git a/data/out.cpp b/data/out.cpp
index 7db593e..c7260cb 100644
--- a/data/out.cpp
+++ b/data/out.cpp
@@ -12,9 +12,4 @@ struct __Thunk<Ret(Args...)> {
     }
 };
 
-{{ main }}
-
-// Enable All Thunks
-void enable_all_thunks(const thunk_enabler_t &thunk_enabler) {
-{{ enableThunks }}
-}
\ No newline at end of file
+{{ main }}
\ No newline at end of file
diff --git a/src/index.ts b/src/index.ts
index 1f50abe..61f3463 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -157,25 +157,22 @@ function makeCompiledCode(output: string) {
 
     // Generate
     let declarations = '';
-    let enableThunks = '';
     for (const structure of structureObjects) {
         const name = structure.getName();
         declarations += `// ${name}\n`;
         try {
             declarations += structure.generateCode();
-            enableThunks += structure.generateEnableThunks();
         } catch (e) {
             console.log(`Error Generating Code: ${name}: ${e instanceof Error ? e.stack : e}`);
             process.exit(1);
         }
         declarations += '\n';
     }
-    enableThunks = enableThunks.slice(0, -1); // Remove Last Newline
 
     // Write
     const headerPath = fs.realpathSync(headerOutput);
     const main = declarations.trim();
-    const result = formatFile('out.cpp', {headerPath, main, enableThunks, data: getDataDir()});
+    const result = formatFile('out.cpp', {headerPath, main, data: getDataDir()});
     fs.writeFileSync(output, result);
 }
 makeCompiledCode(sourceOutput);
diff --git a/src/struct.ts b/src/struct.ts
index b19337c..8545ce0 100644
--- a/src/struct.ts
+++ b/src/struct.ts
@@ -144,7 +144,7 @@ export class Struct {
         if (isVirtual) {
             call += `this->vtable->${method.shortName}`;
         } else {
-            call += `${method.getName()}->get()`;
+            call += `${method.getName()}->get(false)`;
         }
         call += '(';
         if (!method.isStatic) {
@@ -269,27 +269,4 @@ export class Struct {
     setDirectParent(directParent: string) {
         this.#directParent = directParent;
     }
-
-    // Generate Part Of enable_all_thunks()
-    generateEnableThunks() {
-        // Get All Methods
-        const allMethods: Method[] = [];
-        for (const method of this.#methods) {
-            allMethods.push(method);
-        }
-        if (this.#vtable !== null && this.#vtable.canGenerateWrappers()) {
-            const virtualMethods = this.#vtable.getMethods();
-            for (const method of virtualMethods) {
-                if (method) {
-                    allMethods.push(method);
-                }
-            }
-        }
-        // Generate
-        let out = '';
-        for (const method of allMethods) {
-            out += `${INDENT}${method.getName()}->enable_thunk(thunk_enabler);\n`;
-        }
-        return out;
-    }
 }
\ No newline at end of file