symbol-processor/src/vtable.ts

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

export class VTable implements Property {
    readonly #name: string;
    #address: number | null;
    readonly #size: Size;
    readonly #methods: Method[];

    // Constructor
    constructor(name: string) {
        this.#name = name;
        this.#address = null;
        this.#size = new Size(true);
        this.#methods = [];
        // Add Destructors (https://stackoverflow.com/a/17960941)
        const destructor_return = `${name} *`;
        const destructor_args = `(${getSelfArg(name)})`;
        this.add(new Method(name, 'destructor_complete', destructor_return, destructor_args, 0x0));
        this.add(new Method(name, 'destructor_deleting', destructor_return, destructor_args, 0x4));
    }

    // Property Information
    propertyOffset() {
        return 0;
    }
    propertySize() {
        return POINTER_SIZE;
    }
    propertyType() {
        return this.#getName() + ' *';
    }
    propertyName() {
        return 'vtable';
    }
    propertyAlignment() {
        return 4;
    }

    // Setters
    setAddress(address: number) {
        this.#address = address;
    }
    setSize(size: number, isExact: boolean) {
        this.#size.set(size, isExact);
    }

    // Add To VTable
    add(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 (this.#methods[index]) {
            throw new Error(`Duplicate Virtual Method At Offset: ${toHex(offset)}`);
        }
        this.#methods[index] = method;
    }

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

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

    // Generate Code
    generate() {
        let out = '';

        // Check
        this.#check();

        // Method Prototypes
        for (let i = 0; i < this.#methods.length; i++) {
            const info = this.#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 < this.#methods.length; i++) {
            let name = `unknown${i}`;
            let type = 'void *';
            const info = this.#methods[i];
            if (info) {
                name = info.shortName;
                type = info.getType() + ' ';
            }
            out += `${INDENT}${type}${name};\n`;
        }
        out += `};\n`;

        // Sanity Check Size
        const isSizeDefined = this.#size.isExact();
        const size = isSizeDefined ? this.#size.get() : (this.#methods.length * POINTER_SIZE);
        out += assertSize(this.#getName(), size, isSizeDefined);

        // Pointers
        if (this.#address !== null) {
            // Base
            out += `extern ${this.#getName()} *${this.#getName()}_base;\n`;
            // Methods
            for (let i = 0; i < this.#methods.length; i++) {
                const info = this.#methods[i];
                if (info) {
                    const type = `${info.getType()} *`;
                    out += `extern ${type}${info.getName()}_vtable_addr;\n`;
                    out += `extern ${info.generateDefinition('_non_virtual')}`;
                }
            }
        }

        // Duplication Method
        if (isSizeDefined) {
            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
            for (let i = 0; i < this.#methods.length; i++) {
                const info = this.#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.isExact()) {
            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};
    }
}