Attack Of The Templates!

src/common.ts

@@ -114,4 +114,5 @@ export function removeFirstArg(args: string) {
     }
     return '(' + args.substring(index).trim();
 }
-export const ORIGINAL_SUFFIX = '_backup_do_not_use';
+export const ORIGINAL_SUFFIX = '_backup_do_not_use';
+export const VTABLE_ADDR_SUFFIX = `_vtable_addr`;

src/method.ts

@@ -1,4 +1,4 @@
-import { INDENT, formatType, getArgNames } from './common';
+import { INDENT, formatType, getArgNames, prependArg } from './common';
 
 export class Method {
     readonly self: string;
@@ -43,11 +43,21 @@ export class Method {
 
         // 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
-            out += `#define __overwrite_helper_for_${this.getName()}(method, original) \\\n`;
+            const helperName = '__create_overwrite_helper_for_' + this.getName();
-            out += `${INDENT}[]${this.args} { \\\n`;
+            out += `template <int>\n`;
-            out += `${INDENT}${INDENT}${returnType.trim() === 'void' ? '' : 'return '}method(${['original'].concat(getArgNames(this.args)).join(', ')}); \\\n`;
+            out += `static ${type} ${helperName}(${overwriteType} method, ${type} original) {\n`;
-            out += `${INDENT}}\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
@@ -58,16 +68,4 @@ export class Method {
     generateDefinition(nameSuffix?: string) {
         return `${this.getType()} ${this.getName()}${nameSuffix !== undefined ? nameSuffix : ''}`;
     }
-
-    // Generate "New Method" Test
-    generateNewMethodTest(parent: string | null, prefix: string, suffix: string) {
-        let out = `#define __is_new_method_${this.getName()}() (`;
-        if (!this.isInherited) {
-            out += 'true';
-        } else {
-            out += `((void *) ${prefix}${this.getName()}${suffix}) != ((void *) ${prefix}${this.getNameWithCustomSelf(parent!)}${suffix})`;
-        }
-        out += ')\n';
-        return out;
-    }
 }

src/struct.ts

@@ -206,7 +206,7 @@ export class Struct {
 
         // VTable
         if (this.#vtable !== null) {
-            out += this.#vtable.generate(this.#directParent);
+            out += this.#vtable.generate();
         }
 
         // Property Typedefs
@@ -266,7 +266,7 @@ export class Struct {
 
         // VTable
         if (this.#vtable !== null) {
-            const vtable = this.#vtable.generateCode();
+            const vtable = this.#vtable.generateCode(this.#directParent);
             out += vtable;
         }
 

src/vtable.ts

@@ -1,7 +1,8 @@
-import { INDENT, ORIGINAL_SUFFIX, POINTER_SIZE, RTTI_SIZE, assertSize, getSelfArg, toHex } from './common';
+import { INDENT, ORIGINAL_SUFFIX, POINTER_SIZE, RTTI_SIZE, VTABLE_ADDR_SUFFIX, assertSize, getSelfArg, toHex } from './common';
 import { Method } from './method';
 import { Property } from './property';
 
+const structuresWithVTableAddress: string[] = [];
 export class VTable {
     readonly #self: string;
     #address: number | null;
@@ -24,6 +25,7 @@ export class VTable {
     // Setters
     setAddress(address: number) {
         this.#address = address;
+        structuresWithVTableAddress.push(this.#self);
     }
     setSize(size: number) {
         this.#size = size;
@@ -84,8 +86,36 @@ export class VTable {
         return out;
     }
 
+    // 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';
+        let ret: string;
+        if (method.isInherited) {
+            if (structuresWithVTableAddress.includes(parent!)) {
+                // Check If Method Differs From Parent
+                out += `${INDENT}void *parent = (void *) *${method.getNameWithCustomSelf(parent!)}${VTABLE_ADDR_SUFFIX};\n`;
+                out += `${INDENT}void *child = (void *) *${method.getName()}${VTABLE_ADDR_SUFFIX};\n`;
+                ret = 'parent != child';
+            } else {
+                // Unable To Determine
+                ret = 'false';
+            }
+        } else {
+            // No Parent
+            ret = 'true';
+        }
+        out += `${INDENT}return ${ret};\n`;
+        out += '}\n';
+        return out;
+    }
+
     // Generate Header Code
-    generate(directParent: string | null) {
+    generate() {
         let out = '';
 
         // Check
@@ -131,8 +161,8 @@ export class VTable {
                     const name = info.getName();
                     const type = info.getType();
                     const pointerType = `${type} *`;
-                    out += `extern ${pointerType}${name}_vtable_addr;\n`;
+                    out += `extern ${pointerType}${name}${VTABLE_ADDR_SUFFIX};\n`;
-                    out += info.generateNewMethodTest(directParent, '*', '_vtable_addr');
+                    out += this.generateOverwritableTestDefinition(info) + ';\n';
                     out += `extern ${info.generateDefinition(ORIGINAL_SUFFIX)};\n`;
                 }
             }
@@ -148,7 +178,7 @@ export class VTable {
     }
 
     // Generate Compiled Code
-    generateCode() {
+    generateCode(directParent: string | null) {
         let out = '';
 
         // Check
@@ -167,8 +197,9 @@ export class VTable {
                     const name = info.getName();
                     const type = info.getType();
                     const pointerType = `${type} *`;
-                    out += `${pointerType}${name}_vtable_addr = (${pointerType}) ${toHex(vtableAddress)};\n`;
+                    out += `${pointerType}${name}${VTABLE_ADDR_SUFFIX} = (${pointerType}) ${toHex(vtableAddress)};\n`;
-                    out += `${info.generateDefinition(ORIGINAL_SUFFIX)} = *${name}_vtable_addr;\n`;
+                    out += this.generateOverwritableTest(info, directParent);
+                    out += `${info.generateDefinition(ORIGINAL_SUFFIX)} = *${name}${VTABLE_ADDR_SUFFIX};\n`;
                 }
             }
         }