TypeScriptToLua plugin to automatically add type checks to annotated functions

typecheck.ts

import * as ts from "typescript";
import * as tstl from "@jackmacwindows/typescript-to-lua";

const EXPECT_PLUGIN = "cc.expect";
const EXPECT_METHOD = "expect";
const EXPECT_FUNCTION = "___TS_Expect";

function fillTypeList(args: ts.Expression[], type: ts.TypeNode, context: tstl.TransformationContext): boolean {
    if (ts.isUnionTypeNode(type)) {
        for (let t of type.types) if (!fillTypeList(args, t, context)) return false;
    } else if (ts.isFunctionOrConstructorTypeNode(type)) {
        args.push(ts.factory.createStringLiteral("function"));
    } else if (ts.isArrayTypeNode(type)) {
        args.push(ts.factory.createStringLiteral("table"));
    } else if (ts.isLiteralTypeNode(type)) {
        if (type.literal.kind === ts.SyntaxKind.NullKeyword) args.push(ts.factory.createStringLiteral("nil"));
        else if (type.literal.kind === ts.SyntaxKind.FalseKeyword || type.literal.kind === ts.SyntaxKind.TrueKeyword) args.push(ts.factory.createStringLiteral("boolean"));
        else return false;
    } else if (ts.isParenthesizedTypeNode(type)) {
        return fillTypeList(args, type.type, context);
    } else if (ts.isOptionalTypeNode(type)) {
        if (!fillTypeList(args, type.type, context)) return false;
        args.push(ts.factory.createStringLiteral("nil"));
    } else if (type.kind === ts.SyntaxKind.NullKeyword) {
        args.push(ts.factory.createStringLiteral("nil"));
    } else if (type.kind === ts.SyntaxKind.BooleanKeyword) {
        args.push(ts.factory.createStringLiteral("boolean"));
    } else if (type.kind === ts.SyntaxKind.NumberKeyword) {
        args.push(ts.factory.createStringLiteral("number"));
    } else if (type.kind === ts.SyntaxKind.StringKeyword) {
        args.push(ts.factory.createStringLiteral("string"));
    } else if (type.kind === ts.SyntaxKind.FunctionKeyword) {
        args.push(ts.factory.createStringLiteral("function"));
    } else if (type.kind === ts.SyntaxKind.ObjectKeyword) {
        args.push(ts.factory.createStringLiteral("table"));
    } else if (ts.isTypeReferenceNode(type)) {
        args.push(ts.factory.createStringLiteral(type.typeName.getText()));
    } else {
        context.diagnostics.push({
            category: ts.DiagnosticCategory.Warning,
            code: 0,
            file: type.getSourceFile(),
            start: type.pos,
            length: type.end - type.pos,
            messageText: "Could not construct type name for parameter; no type check will be emitted."
        })
        return false;
    }
    return true;
}

function addTypeChecks(m: ts.MethodDeclaration | ts.FunctionDeclaration, context: tstl.TransformationContext) {
    if (m["jsDoc"]) {
        let jsDoc = m["jsDoc"] as ts.JSDoc[];
        if (jsDoc[0].tags?.find(v => v.tagName.escapedText === "typecheck")) {
            let add: ts.Statement[] = [];
            for (let a in m.parameters) {
                let arg = m.parameters[a];
                if (arg.type && !ts.isThisTypeNode(arg.type)) {
                    let args: ts.Expression[] = [
                        ts.factory.createNumericLiteral(parseInt(a) + 1),
                        ts.factory.createIdentifier(arg.name.getText())
                    ];
                    if (fillTypeList(args, arg.type, context)) {
                        context.program["__usesExpect"] = true;
                        add.push(ts.factory.createExpressionStatement(ts.factory.createCallExpression(
                            ts.factory.createIdentifier(EXPECT_FUNCTION),
                            [
                                ts.factory.createKeywordTypeNode(ts.SyntaxKind.NumberKeyword),
                                ts.factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword),
                                ts.factory.createRestTypeNode(ts.factory.createArrayTypeNode(ts.factory.createKeywordTypeNode(ts.SyntaxKind.StringKeyword)))
                            ], args)));
                    }
                }
            }
            m.body?.statements["unshift"](...add);
        }
    }
}

class TypeCheckPlugin implements tstl.Plugin {
    public visitors = {
        [ts.SyntaxKind.FunctionDeclaration]: (node: ts.FunctionDeclaration, context: tstl.TransformationContext): tstl.Statement[] => {
            addTypeChecks(node, context);
            return context.superTransformStatements(node);
        },
        [ts.SyntaxKind.ClassDeclaration]: (node: ts.ClassDeclaration, context: tstl.TransformationContext): tstl.Statement[] => {
            for (let m of node.members) {
                if (ts.isMethodDeclaration(m)) {
                    addTypeChecks(m, context);
                }
            }
            let stat = context.superTransformStatements(node);
            for (const idx in stat) {
                const s = stat[idx]
                if (tstl.isAssignmentStatement(s) && tstl.isStringLiteral(s.right[0]) && tstl.isTableIndexExpression(s.left[0]) && tstl.isStringLiteral(s.left[0].index) && s.left[0].index.value == "name") {
                    stat.splice(parseInt(idx), 0, tstl.createAssignmentStatement(tstl.createTableIndexExpression(s.left[0].table, tstl.createStringLiteral("__name")), s.right[0]));
                    break;
                }
            }
            return stat;
        }
    }
    public beforeEmit(program: ts.Program, options: tstl.CompilerOptions, emitHost: tstl.EmitHost, result: tstl.EmitFile[]): void | ts.Diagnostic[] {
        if (program["__usesExpect"]) {
            for (const file of result) {
                file.code = `local ___TS_Expect_Temp = require('${EXPECT_PLUGIN}')\nlocal function ${EXPECT_FUNCTION}(this, ...) return ___TS_Expect_Temp.${EXPECT_METHOD}(...) end\n` + file.code;
            }
        }
    }
}

const plugin = new TypeCheckPlugin();
export default plugin;