The Compiler Strikes Back!

src/common.ts

@@ -116,3 +116,4 @@ export function removeFirstArg(args: string) {
 }
 export const ORIGINAL_SUFFIX = '_backup_do_not_use';
 export const VTABLE_ADDR_SUFFIX = `_vtable_addr`;
+export const OVERWRITE_TEST_SUFFIX = '__is_overwritable_';

src/index.ts

@@ -141,10 +141,6 @@ function makeMainHeader(output: string) {
     result += '#ifndef __arm__\n';
     result += '#error "Symbols Are ARM-Only"\n';
     result += '#endif\n';
-    result += '\n// Shortcuts\n';
-    result += 'typedef unsigned char uchar;\n';
-    result += 'typedef unsigned short ushort;\n';
-    result += 'typedef unsigned int uint;\n';
     result += '\n// Forward Declarations\n';
     for (const name in STRUCTURE_FILES) {
         result += `typedef struct ${name} ${name};\n`;
@@ -158,6 +154,12 @@ function makeMainHeader(output: string) {
     result += '#include <string>\n';
     result += '#include <vector>\n';
     result += '#include <map>\n';
+    result += `#include <type_traits>\n`;
+    result += `#include <functional>\n`;
+    result += '\n// Shortcuts\n';
+    result += 'typedef unsigned char uchar;\n';
+    result += 'typedef unsigned short ushort;\n';
+    result += 'typedef unsigned int uint;\n';
     result += '\n// Warnings\n';
     result += '#pragma GCC diagnostic push\n';
     result += '#pragma GCC diagnostic ignored "-Winvalid-offsetof"\n';

src/method.ts

@@ -1,4 +1,4 @@
-import { INDENT, formatType, getArgNames, prependArg } from './common';
+import { INDENT, OVERWRITE_TEST_SUFFIX, VTABLE_ADDR_SUFFIX, formatType, getArgNames, prependArg } from './common';
 
 export class Method {
     readonly self: string;
@@ -29,41 +29,89 @@ export class Method {
     getName() {
         return this.getNameWithCustomSelf(this.self);
     }
-    getType() {
+    #getBaseType() {
         return `${this.getName()}_t`;
     }
+    getType() {
+        return `__raw_${this.#getBaseType()}`;
+    }
 
     // Generate Type Definition
     generateTypedef() {
         let out = '';
 
         // Normal Definition
-        const returnType = formatType(this.returnType);
+        out += `typedef ${formatType(this.returnType)}(*${this.getType()})${this.args};\n`;
-        out += `typedef ${returnType}(*${this.getType()})${this.args};\n`;
-
-        // Fancy Overwrite Does Not Support Varargs
-        if (!this.hasVarargs) {
-            // Add Overwrite Typedef
-            const type = this.getType();
-            const overwriteType = `__overwrite_${type}`;
-            out += `typedef ${returnType}(*${overwriteType})${prependArg(this.args, type + ' original')};\n`;
-            // Overwrite Helper
-            const helperName = '__create_overwrite_helper_for_' + this.getName();
-            out += `template <int>\n`;
-            out += `static ${type} ${helperName}(${overwriteType} method, ${type} original) {\n`;
-            out += `${INDENT}static ${overwriteType} stored_method = method;\n`;
-            out += `${INDENT}static ${type} stored_original = original;\n`;
-            out += `${INDENT}return []${this.args} { \\\n`;
-            out += `${INDENT}${INDENT}${returnType.trim() === 'void' ? '' : 'return '}stored_method(${['stored_original'].concat(getArgNames(this.args)).join(', ')}); \\\n`;
-            out += `${INDENT}};\n`;
-            out += `}\n`;
-            out += `#define ${helperName} ${helperName}<__COUNTER__>\n`;
-        }
 
         // Return
         return out;
     }
 
+    // Overwrite Helper
+    #hasOverwriteHelper() {
+        // Fancy Overwrite Does Not Support Varargs
+        return !this.hasVarargs;
+    }
+    generateOverwriteHelper(code: boolean, isVirtual: boolean) {
+        let out = '';
+        if (this.#hasOverwriteHelper()) {
+            // Add Additional Typedefs
+            const baseType = this.#getBaseType();
+            const overwriteType = `__overwrite_${baseType}`;
+            const functionType = baseType;
+            const rawType = this.getType();
+            if (!code) {
+                out += `typedef std::function<std::remove_pointer<${rawType}>::type> ${functionType};\n`;
+                out += `typedef std::function<${this.returnType}${prependArg(this.args, functionType + ' original')}> ${overwriteType};\n`;
+            }
+            // Thunk
+            const name = this.getName();
+            const currentOverwriteName = '__current_overwrite_for_' + name;
+            const thunkName = '__overwrite_thunk_for_' + name;
+            if (code) {
+                out += `static ${functionType} ${currentOverwriteName};\n`;
+                out += `static ${formatType(this.returnType)}${thunkName}${this.args} {\n`;
+                out += `${INDENT}${this.returnType.trim() === 'void' ? '' : 'return '}${currentOverwriteName}(${getArgNames(this.args).join(', ')});\n`;
+                out += `}\n`;
+            }
+            // Overwrite Helper
+            out += `std::string __set_overwrite_for_${name}(const ${overwriteType} &method, const std::function<void *(void *, void *)> &overwriter)`;
+            if (code) {
+                out += ' {\n';
+                if (isVirtual) {
+                    // Check If Method Can Be Overwritten
+                    out += `${INDENT}if (!${OVERWRITE_TEST_SUFFIX}${name}()) {\n`;
+                    out += `${INDENT}${INDENT}return "Unable To Overwrite Virtual Method";\n`;
+                    out += `${INDENT}}\n`;
+                }
+                // Redirect Method Calls To Thunk
+                out += `${INDENT}static bool enabled = false;\n`;
+                out += `${INDENT}if (!enabled) {\n`;
+                if (isVirtual) {
+                    // Get Address Of Virtual Function
+                    out += `${INDENT}${INDENT}${rawType} ${name} = *${name}${VTABLE_ADDR_SUFFIX};\n`;
+                }
+                out += `${INDENT}${INDENT}${rawType} &original = ${name};\n`;
+                out += `${INDENT}${INDENT}${currentOverwriteName} = original;\n`;
+                out += `${INDENT}${INDENT}original = (${rawType}) overwriter((void *) original, (void *) ${thunkName});\n`;
+                out += `${INDENT}${INDENT}enabled = true;\n`;
+                out += `${INDENT}}\n`;
+                // Bind Original Parameter
+                out += `${INDENT}${functionType} original = ${currentOverwriteName};\n`;
+                out += `${INDENT}${currentOverwriteName} = [method, original]${this.args} {\n`;
+                out += `${INDENT}${INDENT}${this.returnType.trim() === 'void' ? '' : 'return '}method(${['original'].concat(getArgNames(this.args)).join(', ')});\n`;
+                out += `${INDENT}};\n`;
+                // Return Success
+                out += `${INDENT}return "";\n`;
+                out += `}`;
+            } else {
+                out += ';';
+            }
+            out += '\n';
+        }
+        return out;
+    }
+
     // Generate Variable Definition
     generateDefinition(nameSuffix?: string) {
         return `${this.getType()} ${this.getName()}${nameSuffix !== undefined ? nameSuffix : ''}`;

src/struct.ts

@@ -201,6 +201,7 @@ export class Struct {
                 out += method.generateTypedef();
                 out += `extern ${method.generateDefinition()};\n`;
                 out += `extern ${method.generateDefinition(ORIGINAL_SUFFIX)};\n`;
+                out += method.generateOverwriteHelper(false, false);
             }
         }
 
@@ -261,6 +262,7 @@ export class Struct {
             if (!method.isInherited) {
                 out += `${method.generateDefinition()} = (${method.getType()}) ${toHex(method.address)};\n`;
                 out += `${method.generateDefinition(ORIGINAL_SUFFIX)} = ${method.getName()};\n`;
+                out += method.generateOverwriteHelper(true, false);
             }
         }
 

src/vtable.ts

@@ -1,4 +1,4 @@
-import { INDENT, ORIGINAL_SUFFIX, POINTER_SIZE, RTTI_SIZE, VTABLE_ADDR_SUFFIX, assertSize, getSelfArg, toHex } from './common';
+import { INDENT, ORIGINAL_SUFFIX, OVERWRITE_TEST_SUFFIX, POINTER_SIZE, RTTI_SIZE, VTABLE_ADDR_SUFFIX, assertSize, getSelfArg, toHex } from './common';
 import { Method } from './method';
 import { Property } from './property';
 
@@ -87,13 +87,10 @@ export class VTable {
     }
 
     // Overwritable Test (Only For Virtual Methods)
-    generateOverwritableTestDefinition(method: Method) {
-        return `bool __is_overwritable_${method.getName()}()`;
-    }
     generateOverwritableTest(method: Method, parent: string | null) {
         // Test If Overwriting Method Would Also Modify Parent
         let out = '';
-        out += this.generateOverwritableTestDefinition(method) + ' {\n';
+        out += `static inline bool ${OVERWRITE_TEST_SUFFIX}${method.getName()}() {\n`;
         let ret: string;
         if (method.isInherited) {
             if (structuresWithVTableAddress.includes(parent!)) {
@@ -162,8 +159,8 @@ export class VTable {
                     const type = info.getType();
                     const pointerType = `${type} *`;
                     out += `extern ${pointerType}${name}${VTABLE_ADDR_SUFFIX};\n`;
-                    out += this.generateOverwritableTestDefinition(info) + ';\n';
                     out += `extern ${info.generateDefinition(ORIGINAL_SUFFIX)};\n`;
+                    out += info.generateOverwriteHelper(false, true);
                 }
             }
         }
@@ -198,8 +195,9 @@ export class VTable {
                     const type = info.getType();
                     const pointerType = `${type} *`;
                     out += `${pointerType}${name}${VTABLE_ADDR_SUFFIX} = (${pointerType}) ${toHex(vtableAddress)};\n`;
-                    out += this.generateOverwritableTest(info, directParent);
                     out += `${info.generateDefinition(ORIGINAL_SUFFIX)} = *${name}${VTABLE_ADDR_SUFFIX};\n`;
+                    out += this.generateOverwritableTest(info, directParent);
+                    out += info.generateOverwriteHelper(true, true);
                 }
             }
         }