typocomath/src/plugins/typeInferPlugin.ts

import ts from 'typescript'

const transformer: ts.TransformerFactory<ts.SourceFile> = context => {
  // TODO: get a reference to the program instance that the plugin is running in instead of creating a new program?
  const program = ts.createProgram([],  {})
  const checker = program.getTypeChecker()

  return sourceFile => {
    const visitor = (node: ts.Node): ts.Node => {
      // we're looking for a function call like $reflect(deps => ...)
      // @ts-ignore
      if (ts.isCallExpression(node) && ts.isIdentifier(node.expression) && node.expression.escapedText === '$reflect') {
        console.log('PLUGIN: FOUND AN OCCURRENCE OF $reflect')
        // console.log('PARENT')
        // console.log(node)

        // TODO: validate that argNode is an ArrowFunction
        // @ts-ignore
        const argNode = node.arguments[0]
        // @ts-ignore
        const returnType = argNode.type.getText(sourceFile)
        console.log('PLUGIN: RETURN TYPE')
        console.log(returnType)
        // (a: number) => number

        // @ts-ignore
        const paramNode = argNode.parameters[0]
        const paramTypeSrc = paramNode.type.getText(sourceFile)
        console.log('PLUGIN: PARAM TYPE SRC', paramTypeSrc)
        // {
        //   multiply: (a: number, b: number) => number,
        //   unaryMinus: (x: number) => number, // just for the experiment
        // }

        const type = checker.getTypeAtLocation(paramNode)
        const paramType = checker.typeToString(type, paramNode, ts.TypeFormatFlags.InTypeAlias)
        console.log('PLUGIN: PARAM TYPE STRING', paramType)
        // { multiply: (a: number, b: number) => number; unaryMinus: (x: number) => number; }

        const depsAndReturnType = `{ deps: ${paramType}; return: ${returnType} }`

        // Now we insert a second argument to the $reflect function call: a string with the types
        // @ts-ignore
        node.arguments.push(ts.factory.createStringLiteral(depsAndReturnType))

        return node
      }

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

    return ts.visitNode(sourceFile, visitor)
  }
}

export default transformer