symbol-processor/src/vtable.ts

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;
    #size: number | null;
    readonly #methods: Method[];
    readonly property: Property;
    #destructorOffset: number;

    // Constructor
    constructor(self: string) {
        this.#self = self;
        this.#address = null;
        this.#size = null;
        this.#methods = [];
        this.#destructorOffset = 0;
        // Create Property
        this.property = new Property(0, this.#getName() + ' *', 'vtable', this.#self);
    }

    // Setters
    setAddress(address: number) {
        this.#address = address;
        structuresWithVTableAddress.push(this.#self);
    }
    setSize(size: number) {
        this.#size = size;
    }
    setDestructorOffset(destructorOffset: number) {
        this.#destructorOffset = destructorOffset;
    }

    // Add To VTable
    #add(target: Method[], method: Method) {
        // Check Offset
        const offset = method.address;
        if ((offset % POINTER_SIZE) !== 0) {
            throw new Error(`Invalid VTable Offset: ${toHex(offset)}`);
        }
        // Add
        const index = offset / POINTER_SIZE;
        if (target[index]) {
            throw new Error(`Duplicate Virtual Method At Offset: ${toHex(offset)}`);
        }
        target[index] = method;
    }
    add(method: Method) {
        this.#add(this.#methods, method);
    }

    // Get Structure Name
    #getName() {
        return this.#self + '_vtable';
    }

    // Check
    #check() {
        // Check Size
        if (this.#size !== null) {
            const size = this.#size;
            const maxMethodCount = size / POINTER_SIZE;
            if (maxMethodCount < this.#methods.length) {
                throw new Error(`VTable Size Too Small: ${toHex(size)}`);
            }
            if (maxMethodCount > this.#methods.length) {
                this.#methods.length = maxMethodCount;
            }
        }
    }

    // Get Full Methods Table
    getMethods() {
        // Copy Array
        const out = [];
        out.push(...this.#methods);
        // Add Destructors (https://stackoverflow.com/a/17960941)
        const destructor_return = `${this.#self} *`;
        const destructor_args = `(${getSelfArg(this.#self)})`;
        this.#add(out, new Method(this.#self, 'destructor_complete', destructor_return, destructor_args, 0x0 + this.#destructorOffset, false, false));
        this.#add(out, new Method(this.#self, 'destructor_deleting', destructor_return, destructor_args, 0x4 + this.#destructorOffset, false, false));
        // Return
        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() {
        let out = '';

        // Check
        this.#check();

        // Method Prototypes
        const methods = this.getMethods();
        for (let i = 0; i < methods.length; i++) {
            const info = methods[i];
            if (info) {
                out += info.generateTypedef();
            }
        }

        // Structure
        out += `typedef struct ${this.#getName()} ${this.#getName()};\n`;
        out += `struct ${this.#getName()} {\n`;
        for (let i = 0; i < methods.length; i++) {
            let name = `unknown${i}`;
            let type = 'void *';
            const info = methods[i];
            if (info) {
                name = info.shortName;
                type = info.getType() + ' ';
            }
            out += `${INDENT}${type}${name};\n`;
        }
        out += `};\n`;

        // Sanity Check Size
        if (this.#size !== null) {
            out += assertSize(this.#getName(), this.#size);
        }

        // Pointers
        if (this.#address !== null) {
            // Base
            out += `extern ${this.#getName()} *${this.#getName()}_base;\n`;
            // Methods
            for (let i = 0; i < methods.length; i++) {
                const info = methods[i];
                if (info) {
                    const name = info.getName();
                    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`;
                }
            }
        }

        // Duplication Method
        if (this.#size !== null) {
            out += `${this.#getName()} *dup_${this.#getName()}(${this.#getName()} *vtable);\n`;
        }

        // Return
        return out;
    }

    // Generate Compiled Code
    generateCode(directParent: string | null) {
        let out = '';

        // Check
        this.#check();

        // Pointers
        if (this.#address !== null) {
            // Base
            out += `${this.#getName()} *${this.#getName()}_base = (${this.#getName()} *) ${toHex(this.#address)};\n`;
            // Methods
            const methods = this.getMethods();
            for (let i = 0; i < methods.length; i++) {
                const info = methods[i];
                if (info) {
                    const vtableAddress = this.#address + (i * POINTER_SIZE);
                    const name = info.getName();
                    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`;
                }
            }
        }

        // Duplication Method
        if (this.#size !== null) {
            out += `${this.#getName()} *dup_${this.#getName()}(${this.#getName()} *vtable) {\n`;
            out += `${INDENT}uchar *real_vtable = (uchar *) vtable;\n`;
            out += `${INDENT}real_vtable -= ${RTTI_SIZE};\n`;
            out += `${INDENT}size_t real_vtable_size = sizeof(${this.#getName()}) + ${RTTI_SIZE};\n`;
            out += `${INDENT}uchar *new_vtable = (uchar *) ::operator new(real_vtable_size);\n`;
            out += `${INDENT}if (new_vtable == NULL) {\n`;
            out += `${INDENT}${INDENT}return NULL;\n`;
            out += `${INDENT}}\n`;
            out += `${INDENT}memcpy((void *) new_vtable, (void *) real_vtable, real_vtable_size);\n`;
            out += `${INDENT}new_vtable += ${RTTI_SIZE};\n`;
            out += `${INDENT}return (${this.#getName()} *) new_vtable;\n`;
            out += '}\n';
        }

        // Return
        return out;
    }
}