cc/src/transformer.ts

import * as ts from 'typescript'
import {} from 'ts-expose-internals'

const interfaceDeclarations = [] as ts.InterfaceDeclaration[]
const zodifiedIds = [] as string[]

const identifierProgram = (program: ts.Program, options: any) => {
    return (context: ts.TransformationContext) => {
        const factory = context.factory

        return (sourceFile: ts.SourceFile) => {
            console.log(`TT: ${sourceFile.fileName} : ${interfaceDeclarations.length}`)  // eslint-disable-line no-console
            const visitor = (node: ts.Node): ts.Node => {
                if ( ts.isNewExpression(node) ) {
                    const exp = node as ts.NewExpression
                    console.log(`New Expression: ${exp.expression.getText()}`)  // eslint-disable-line no-console

                    const schema = [] as number[]
                    for ( const type of exp.typeArguments ?? [] ) {
                        console.log(`Type symbol: ${type.localSymbol?.name} : ${type.id} : ${type.getText()}`)  // eslint-disable-line no-console

                        for ( const decl of interfaceDeclarations ) {
                            if ( decl?.name?.symbol === (type as any)?.typeName?.symbol ) {
                                console.log(`Match to interface: ${decl.name.text}`)  // eslint-disable-line no-console
                                console.log(decl)  // eslint-disable-line no-console

                                const id = parseInt(zodifiedIds[interfaceDeclarations.indexOf(decl)], 10)
                                if ( !isNaN(id) ) {
                                    schema.push(id)
                                }
                            }
                        }
                    }

                    console.log(schema)  // eslint-disable-line no-console

                    const call = factory.createCallExpression(
                        factory.createParenthesizedExpression(factory.createArrowFunction(
                            undefined,
                            undefined,
                            [],
                            undefined,
                            factory.createToken(ts.SyntaxKind.EqualsGreaterThanToken),
                            factory.createBlock(
                                [
                                    factory.createVariableStatement(
                                        undefined,
                                        factory.createVariableDeclarationList(
                                            [factory.createVariableDeclaration(
                                                factory.createIdentifier('vI'),
                                                undefined,
                                                undefined,
                                                node,
                                            )],
                                            ts.NodeFlags.Const,
                                        ),
                                    ),
                                    factory.createExpressionStatement(factory.createBinaryExpression(
                                        factory.createPropertyAccessExpression(
                                            factory.createParenthesizedExpression(factory.createAsExpression(
                                                factory.createIdentifier('vI'),
                                                factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword),
                                            )),
                                            factory.createIdentifier('__exZodifiedSchemata'),
                                        ),
                                        factory.createToken(ts.SyntaxKind.EqualsToken),
                                        factory.createArrayLiteralExpression(
                                            schema.map(x => factory.createNumericLiteral(x)),
                                            false,
                                        ),
                                    )),
                                    factory.createReturnStatement(factory.createIdentifier('vI')),
                                ],
                                true,
                            ),
                        )),
                        undefined,
                        [],
                    )

                    return call
                }

                return ts.visitEachChild(node, visitor, context)
            }

            return ts.visitEachChild(sourceFile, visitor, context)
        }
    }
}

function transformAst(this: typeof ts, context: ts.TransformationContext) {
    // eslint-disable-next-line @typescript-eslint/no-this-alias,no-invalid-this
    const tsInstance = this

    return (sourceFile: ts.SourceFile) => {
        if ( sourceFile.fileName.endsWith('.d.ts') ) {
            return sourceFile
        }

        console.log(`ID: ${sourceFile.fileName}`)  // eslint-disable-line no-console
        const visit = (node: ts.Node): ts.Node => {
            if ( ts.isInterfaceDeclaration(node) ) {
                console.log(`Interface Declaration: ${(node as ts.InterfaceDeclaration).name}`)  // eslint-disable-line no-console

                if ( sourceFile.identifiers.has('exZodifiedSchema') ) {
                    const id = [...(sourceFile as any).statements].reverse()[0].expression.arguments[0].text
                    console.log('Zod-ID: ' + id)  // eslint-disable-line no-console

                    interfaceDeclarations.push(node)
                    zodifiedIds.push(id)
                }
            }

            return tsInstance.visitEachChild(node, visit, context)
        }

        return tsInstance.visitEachChild(sourceFile, visit, context)
    }
}

/**
 * Patches existing Compiler Host (or creates new one) to allow feeding updated file content from cache
 */
function getPatchedHost(
    maybeHost: ts.CompilerHost | undefined,
    tsInstance: typeof ts,
    compilerOptions: ts.CompilerOptions,
): ts.CompilerHost & { fileCache: Map<string, ts.SourceFile> } {
    const fileCache = new Map()
    const compilerHost = maybeHost ?? tsInstance.createCompilerHost(compilerOptions, true)
    const originalGetSourceFile = compilerHost.getSourceFile

    return Object.assign(compilerHost, {
        getSourceFile(fileName: string, languageVersion: ts.ScriptTarget) {
            fileName = tsInstance.normalizePath(fileName)
            if (fileCache.has(fileName)) {
                return fileCache.get(fileName)
            }

            // eslint-disable-next-line prefer-spread,prefer-rest-params
            const sourceFile = originalGetSourceFile.apply(void 0, Array.from(arguments) as any)
            fileCache.set(fileName, sourceFile)

            return sourceFile
        },
        fileCache,
    })
}

const transformProgram = (program: ts.Program, host: ts.CompilerHost | undefined, options: any, { ts: tsInstance }: any) => {
    const compilerOptions = program.getCompilerOptions()
    const compilerHost = getPatchedHost(host, tsInstance, compilerOptions)
    const rootFileNames = program.getRootFileNames().map(tsInstance.normalizePath)

    // transform ast
    const transformedSource = tsInstance.transform(
        [...program.getSourceFiles()],
        [transformAst.bind(tsInstance)],
        compilerOptions,
    ).transformed

    return program

    /* /!* Render modified files and create new SourceFiles for them to use in host's cache *!/
    const { printFile } = tsInstance.createPrinter()
    for (const sourceFile of transformedSource) {
        const { fileName, languageVersion } = sourceFile
        const updatedSourceFile = tsInstance.createSourceFile(fileName, printFile(sourceFile), languageVersion)
        compilerHost.fileCache.set(fileName, updatedSourceFile)
    }

    /!* Re-create Program instance *!/
    return tsInstance.createProgram(rootFileNames, compilerOptions, compilerHost)*/
}

const identifierProgramOld = (program: ts.Program) => {
    const transformerFactory: ts.TransformerFactory<ts.SourceFile> = context => {
        const factory = context.factory
        return sourceFile => {
            console.log(`ID Source: ${sourceFile.fileName}`)  // eslint-disable-line no-console
            const visitor = (node: ts.Node): ts.Node => {
                if (ts.isInterfaceDeclaration(node)) {
                    console.log('Found Interface: ' + (node as ts.InterfaceDeclaration).name.text)  // eslint-disable-line no-console
                    interfaceDeclarations.push(node)
                }

                return ts.visitEachChild(node, visitor, context)
            }

            return ts.visitNode(sourceFile, visitor)
        }
    }

    return transformerFactory
}

/* const transformerProgram = (program: ts.Program) => {
    // const typeChecker = program.getTypeChecker()

    // Create array of found symbols
    // const foundSymbols = new Array<ts.Symbol>()
    const typeReferences = [] as ts.TypeReference[]
    const newExpressions = [] as ts.NewExpression[]

    const transformerFactory: ts.TransformerFactory<ts.SourceFile> = context => {
        const factory = context.factory
        return sourceFile => {
            console.log(`TT Source: ${sourceFile.fileName} : ${interfaceDeclarations.length}`)  // eslint-disable-line no-console
            const visitor = (node: ts.Node): ts.Node => {
                if ( ts.isNewExpression(node) ) {
                    const exp = node as ts.NewExpression
                    console.log('New Expression: ' + exp.expression.getText() + ' w/ type args: ' + exp.typeArguments?.[0].getText())  // eslint-disable-line no-console
                    // newExpressions.push(node)

                    const call = factory.createCallExpression(
                        factory.createParenthesizedExpression(factory.createArrowFunction(
                            undefined,
                            undefined,
                            [],
                            undefined,
                            factory.createToken(ts.SyntaxKind.EqualsGreaterThanToken),
                            factory.createBlock(
                                [
                                    factory.createVariableStatement(
                                        undefined,
                                        factory.createVariableDeclarationList(
                                            [factory.createVariableDeclaration(
                                                factory.createIdentifier('vI'),
                                                undefined,
                                                undefined,
                                                exp,
                                            )],
                                            ts.NodeFlags.Const,
                                        ),
                                    ),
                                    factory.createExpressionStatement(factory.createBinaryExpression(
                                        factory.createPropertyAccessExpression(
                                            factory.createParenthesizedExpression(factory.createAsExpression(
                                                factory.createIdentifier('vI'),
                                                factory.createKeywordTypeNode(ts.SyntaxKind.AnyKeyword),
                                            )),
                                            factory.createIdentifier('__exSchema'),
                                        ),
                                        factory.createToken(ts.SyntaxKind.EqualsToken),
                                        factory.createObjectLiteralExpression(
                                            [],
                                            false,
                                        ),
                                    )),
                                    factory.createReturnStatement(factory.createIdentifier('vI')),
                                ],
                                true,
                            ),
                        )),
                        undefined,
                        [],
                    )

                    console.log('returned call expression')  // eslint-disable-line no-console
                    return call
                }

                return ts.visitEachChild(node, visitor, context)
            }

            return ts.visitNode(sourceFile, visitor)
        }
    }

    return transformerFactory
}*/

export {identifierProgram, transformProgram}