symbol-processor/src/vtable.ts

import { INDENT, POINTER_SIZE, RTTI_SIZE, assertSize, getSelfArg, toHex } from './common';
import { Method } from './method';
import { Property } from './property';

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;
    }
    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;
    }

    // Generate Header Code
    generate(directParent: string | null) {
        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 type = `${info.getType()} *`;
                    out += `extern ${type}${info.getName()}_vtable_addr;\n`;
                    out += `extern ${info.generateDefinition('_non_virtual')}`;
                    out += info.generateNewMethodTest(directParent, '*', '_vtable_addr');
                }
            }
        }

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

        // Return
        return out;
    }

    // Generate Compiled Code
    generateCode() {
        let declarations = '';
        let init = '';

        // Check
        this.#check();

        // Pointers
        if (this.#address !== null) {
            // Base
            init += `${INDENT}${this.#getName()}_base = (${this.#getName()} *) ${toHex(this.#address)};\n`;
            declarations += `${this.#getName()} *${this.#getName()}_base;\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 type = `${info.getType()} *`;
                    init += `${INDENT}${info.getName()}_vtable_addr = (${type}) ${toHex(vtableAddress)};\n`;
                    declarations += `${type}${info.getName()}_vtable_addr;\n`;
                    init += `${INDENT}${info.getName()}_non_virtual = *${info.getName()}_vtable_addr;\n`;
                    declarations += info.generateDefinition('_non_virtual');
                }
            }
        }

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

        // Return
        return {declarations, init};
    }
}