A babel plugint that turns typescript annotations into runtime values

babel-plugin-reify-typescript-annotations.ts

import transformTypescript from "@babel/plugin-transform-typescript"
import {declare} from "@babel/helper-plugin-utils"
import {types as t, PluginObj, Node, NodePath} from "@babel/core"
import template from "@babel/template"
import {
  TSTypeAliasDeclaration,
  TSType,
  TSTypeParameter,
  TSUnionType,
  TSTypeReference,
  Expression,
  TSLiteralType,
  TSArrayType,
  TSTypeLiteral,
  TSPropertySignature,
  TSIntersectionType,
  TSTupleType,
  CallExpression,
  Identifier,
  Program,
  TSTypeOperator,
  TSInterfaceDeclaration,
  TSInterfaceBody
} from "@babel/types"
import {addNamed} from "@babel/helper-module-imports"

// t.tsImportType()

/**
 * @todo find a way to limit type checking to certain files or maybe certain declarations
 * @note let's skip interfaces for now
 * @note we'll only support references to other types in the root scope
 */

export const plugin = declare((api, options, dirname) => {
  const obj = {
    name: "ts-runtime",
    // inherits: transformTypescript,
    manipulateOptions(opts, parserOpts) {
      parserOpts.plugins.push("typescript")
    },
    visitor: {
      Program(nodePath: NodePath<t.Program>) {
        const globalDeclarations: string[] = []
        nodePath.node.body.forEach(node => {
          if (
            node.type === "TSTypeAliasDeclaration" ||
            node.type === "TSInterfaceDeclaration"
          ) {
            globalDeclarations.push(node.id.name)
          }
        })
        const state: ProgramState = {
          type: "Program",
          program: {
            globalDeclarations,
            importedHelpers: {},
            nodePath,
            resolvedImportedTypes: {}
          }
        }
        nodePath.traverse(rootVisitor, state)
      }
      // ImportDeclaration(a, b) {
      //   debugger
      // },
    }
  }
  return obj
})

type ProgramStateProps = {
  globalDeclarations: string[]
  importedHelpers: {[K in keyof typeof helperMethodsToImportNames]?: Identifier}
  nodePath: NodePath<Program>
  resolvedImportedTypes: Record<string, Expression>
}

type ProgramState = {
  type: "Program"
  program: ProgramStateProps
}

type DeclarationState = {
  type: "Declaration"
  declaration: {
    identifierName: string
    varName: string
    isExported: boolean
    paramNames: string[]
    statementParent: NodePath<t.Statement>
  }
  program: ProgramStateProps
}

type AllStates = ProgramState | DeclarationState

const rootVisitor = {
  TSTypeAliasDeclaration(
    nodePath: NodePath<t.TSTypeAliasDeclaration>,
    programState: ProgramState
  ) {
    const identifierName = nodePath.node.id.name
    const varName = identNameToVarName(identifierName)
    const isExported = nodePath.parent.type === "ExportNamedDeclaration"
    const statementParent = nodePath.getStatementParent()

    if (!t.isProgram(statementParent.scope.block)) {
      return
    }

    const state: DeclarationState = {
      type: "Declaration",
      declaration: {
        identifierName,
        varName,
        isExported,
        paramNames: [],
        statementParent
      },
      program: programState.program
    }

    const aliasDeclaration = generators.declarators.alias(nodePath, state)

    const constDeclaration = isExported
      ? templates.ExportedConst({
          varName,
          value: aliasDeclaration
        })
      : templates.Const({varName, value: aliasDeclaration})

    statementParent.insertAfter(constDeclaration)
  },
  TSInterfaceDeclaration(
    nodePath: NodePath<TSInterfaceDeclaration>,
    programState: ProgramState
  ) {
    const identifierName = nodePath.node.id.name
    const varName = identNameToVarName(identifierName)
    const isExported = nodePath.parent.type === "ExportNamedDeclaration"
    const statementParent = nodePath.getStatementParent()

    if (!t.isProgram(statementParent.scope.block)) {
      return
    }

    const state: DeclarationState = {
      type: "Declaration",
      declaration: {
        identifierName,
        varName,
        isExported,
        paramNames: [],
        statementParent
      },
      program: programState.program
    }

    const interfaceDeclaration = generators.declarators.interface(
      nodePath,
      state
    )

    const constDeclaration = isExported
      ? templates.ExportedConst({
          varName,
          value: interfaceDeclaration
        })
      : templates.Const({varName, value: interfaceDeclaration})

    statementParent.insertAfter(constDeclaration)
  },
  CallExpression(nodePath: NodePath<CallExpression>, state: ProgramState) {
    const {node} = nodePath
    if (node.callee.type === "Identifier") {
      if (helperMethodsToImportNames[node.callee.name]) {
        nodePath.replaceWith(
          generators.fns.helperMethods(
            node.callee.name as $IntentionalAny,
            nodePath,
            state
          )
        )
      }
    }
  }
}

const identNameToVarName = (identName: string) => "__type_" + identName

const templates = {
  ExportedConst: template(`export const %%varName%% = %%value%%;`),
  Const: template(`const %%varName%% = %%value%%;`)
}

const generators = {
  declarators: {
    alias(nodePath: NodePath<TSTypeAliasDeclaration>, state: DeclarationState) {
      const {node} = nodePath

      state.declaration.paramNames = !node.typeParameters
        ? []
        : node.typeParameters.params.map(p => p.name)

      const annotation = generators.annotations.decideWhichAnnotationToGenerate(
        node.typeAnnotation,
        state
      )

      const parameterDeclarations = !node.typeParameters
        ? []
        : node.typeParameters.params.map((param: TSTypeParameter) =>
            generators.declarators.typeParameter(param, state)
          )

      const parameters = t.arrayExpression(parameterDeclarations)

      return template.expression(`
      {
        sort: "alias",
        name: "${state.declaration.identifierName}",
        parameters: %%parameters%%,
        annotation: %%annotation%%
      }
      `)({
        annotation,
        parameters
      })
    },
    interface(
      nodePath: NodePath<TSInterfaceDeclaration>,
      state: DeclarationState
    ) {
      const {node} = nodePath

      state.declaration.paramNames = !node.typeParameters
        ? []
        : node.typeParameters.params.map(p => p.name)

      const body = generators.declarators.interfaceBody(node.body, state)

      const parameterDeclarations = !node.typeParameters
        ? []
        : node.typeParameters.params.map((param: TSTypeParameter) =>
            generators.declarators.typeParameter(param, state)
          )

      const parameters = t.arrayExpression(parameterDeclarations)

      return template.expression(`
      {
        sort: "interface",
        name: "${state.declaration.identifierName}",
        parameters: %%parameters%%,
        body: %%body%%
      }
      `)({
        body: body,
        parameters
      })
    },
    typeParameter(p: TSTypeParameter, state: DeclarationState) {
      const def =
        p.default === undefined
          ? "undefined"
          : generators.annotations.decideWhichAnnotationToGenerate(
              p.default,
              state
            )

      const constraint =
        p.constraint === undefined
          ? "undefined"
          : generators.annotations.decideWhichAnnotationToGenerate(
              p.constraint,
              state
            )

      return template.expression(`{
        sort: 'declarationParameter',
        name: "${p.name}",
        default: %%default%%,
        constraint: %%constraint%%
      }`)({
        default: def,
        constraint
      })
    },
    interfaceBody(node: TSInterfaceBody, state: DeclarationState) {
      return template.expression("null")()
    }
  },
  annotations: {
    decideWhichAnnotationToGenerate(
      node: TSType,
      state: AllStates
    ): Expression {
      if (node.type.match(/^TS.*Keyword$/)) {
        return generators.annotations.primitive(node)
      } else if (node.type === "TSUnionType") {
        return generators.annotations.union(node, state)
      } else if (node.type === "TSIntersectionType") {
        return generators.annotations.intersection(node, state)
      } else if (node.type === "TSTypeReference") {
        return generators.annotations.decideWhichReferenceToGenerate(
          node,
          state
        )
      } else if (node.type === "TSLiteralType") {
        return generators.annotations.literal(node)
      } else if (node.type === "TSArrayType") {
        return generators.annotations.arrayType(node, state)
      } else if (node.type === "TSTypeLiteral") {
        return generators.annotations.objectLiteral(node, state)
      } else if (node.type === "TSTupleType") {
        return generators.annotations.tuple(node, state)
      } else if (node.type === "TSTypeOperator") {
        return generators.operators.decideWhichOperatorToGenerate(node, state)
      } else {
        throw new Error(
          `TODO implement case "${node.type}" in decideWhichAnnotationToGenerate()`
        )
      }
    },
    literal(node: TSLiteralType) {
      const {literal} = node
      const value =
        literal.type === "StringLiteral"
          ? JSON.stringify(literal.value)
          : String(literal.value)

      return template.expression(`{
        sort: 'literal', value: %%value%%
      }`)({value})
    },
    primitive(node: TSType): Expression {
      // replace TSStringKeyword => String
      const _name = node.type.replace(/^TS/, "").replace(/Keyword$/, "")
      // String => string
      const name = _name[0].toLowerCase() + _name.substr(1, _name.length)
      const which = `"${name}"`

      return template.expression(`{sort: "primitive", which: %%which%%}`)({
        which
      })
    },
    union(node: TSUnionType, state: AllStates): Expression {
      const types = node.types.map((n: TSType) =>
        generators.annotations.decideWhichAnnotationToGenerate(n, state)
      )
      const typesExpression = t.arrayExpression(types)
      return template.expression(`{
        sort: 'union',
        types: %%types%%
      }`)({types: typesExpression})
    },
    intersection(node: TSIntersectionType, state: AllStates): Expression {
      const types = node.types.map((n: TSType) =>
        generators.annotations.decideWhichAnnotationToGenerate(n, state)
      )
      const typesExpression = t.arrayExpression(types)
      return template.expression(`{
        sort: 'intersection',
        types: %%types%%
      }`)({types: typesExpression})
    },
    tuple(node: TSTupleType, state: AllStates): Expression {
      const types = node.elementTypes.map((n: TSType) =>
        generators.annotations.decideWhichAnnotationToGenerate(n, state)
      )
      const typesExpression = t.arrayExpression(types)
      return template.expression(`{
        sort: 'tuple',
        types: %%types%%
      }`)({types: typesExpression})
    },
    decideWhichReferenceToGenerate(
      node: TSTypeReference,
      state: DeclarationState | ProgramState
    ): Expression {
      const targetTypeNode = node.typeName
      const targetTypeName = (targetTypeNode as $LibBugAny).name as string
      const kind =
        // if it's an Array, or one of typescript's native types (like Pick, Exclude, etc)
        standardReferenceIdentifiers[targetTypeName]
          ? "standardLib"
          : // if it's a type parameter
          state.type === "Declaration" &&
            state.declaration.paramNames.some(n => n === targetTypeName)
          ? "typeParameterReference"
          : // otherwise, it's a reference outside of the current declaration
            "reference"

      let ref: t.Expression
      if (kind === "reference") {
        const referencedTypeVarName = identNameToVarName(targetTypeName)
        if (state.program.globalDeclarations.some(n => n === targetTypeName)) {
          ref = template.expression(`{
            sort: 'referenceToTypeInModule',
            referencedTypeName: %%referencedTypeName%%,
            referencedType: () => %%referencedTypeVarName%%,
          }`)({
            referencedTypeName: `"${targetTypeName}"`,
            referencedTypeVarName: referencedTypeVarName
          })
        } else {
          ref = template.expression(`{
            sort: 'referenceToExternalType',
            referencedTypeName: %%referencedTypeName%%,
            referencedType: () => %%referencedTypeVarName%%,
          }`)({
            referencedTypeName: `"${targetTypeName}"`,
            referencedTypeVarName: typeReferenceToVar(node, state)
          })
        }
      } else if (kind === "typeParameterReference") {
        ref = template.expression(`{
          sort: "typeParameterReference",
          name: %%name%%,
        }`)({name: `"${targetTypeName}"`})
      } else {
        ref = template.expression(`{
          sort: 'standardLib',
          which: %%which%%,
        }`)({which: `"${targetTypeName}"`})
      }

      if (!node.typeParameters || node.typeParameters.params.length === 0) {
        return ref
      } else {
        const params = t.arrayExpression(
          node.typeParameters.params.map(param =>
            generators.annotations.decideWhichAnnotationToGenerate(param, state)
          )
        )
        return template.expression(`{
          sort: 'genericApplication',
          params: %%params%%,
          ref: %%ref%%
          }
        `)({params, ref})
      }
    },
    arrayType(node: TSArrayType, state: AllStates): Expression {
      const params = t.arrayExpression([
        generators.annotations.decideWhichAnnotationToGenerate(
          node.elementType,
          state
        )
      ])

      return template.expression(`{
        sort: 'genericApplication',
        params: %%params%%,
        ref: {
          sort: 'standardLib',
          which: 'Array',
        },
      }`)({params})
    },
    objectLiteral(node: TSTypeLiteral, state: AllStates): Expression {
      const props = []
      const members = t.objectExpression(
        node.members.map((member: TSPropertySignature) => {
          if (member.typeAnnotation.type !== "TSTypeAnnotation") {
            throw new Error("@todo Unexpected member type. Investigate.")
          }
          if (member.key.type !== "Identifier") {
            throw new Error("@todo Unexpected key type. Investigate.")
          }
          const memberExpression = template.expression(`{
            sort: 'objectMember',
            optional: %%optional%%,
            annotation: %%annotation%%
          }`)({
            // @todo
            optional: member.optional ? 'true' : 'false',
            annotation: generators.annotations.decideWhichAnnotationToGenerate(
              member.typeAnnotation.typeAnnotation,
              state
            )
          })
          return t.objectProperty(
            t.stringLiteral(member.key.name),
            memberExpression
          )
        })
      )
      return template.expression(`{
        sort: 'objectLiteral',
        members: %%members%%
      }`)({members})
    }
  },
  operators: {
    decideWhichOperatorToGenerate(
      node: TSTypeOperator,
      state: AllStates
    ): Expression {
      if (node.operator === "keyof") {
        return generators.operators.keyof(node, state)
      } else {
        throw new Error(`@todo Unsupported operator "${node.operator}"`)
      }
    },
    keyof(node: TSTypeOperator, state: AllStates): Expression {
      // const a : Exclude<>
      return template.expression(`{
        sort: "keyof",
        target: %%target%%
        }`)({
        target: generators.annotations.decideWhichAnnotationToGenerate(
          node.typeAnnotation,
          state
        )
      })
    }
  },
  fns: {
    helperMethods(
      which: "refine" | "is" | "assertType",
      nodePath: NodePath<CallExpression>,
      state: ProgramState
    ): Expression {
      const {node} = nodePath
      const {typeParameters} = node

      if (
        !typeParameters ||
        typeParameters.type !== "TSTypeParameterInstantiation"
      ) {
        throw nodePath.buildCodeFrameError(
          `A call to ${which}() must include a type parameter, like ${which}<A>(fn)`
        )
      }
      if (typeParameters.params.length !== 1) {
        throw nodePath.buildCodeFrameError(
          `A call to ${which}() must include exactly one type parameter, like ${which}<A>(fn)`
        )
      }

      const paramNode = typeParameters.params[0]
      // debugger
      const typeVarName = generators.annotations.decideWhichAnnotationToGenerate(
        paramNode,
        state
      )

      let fnName = state.program.importedHelpers[which]
      if (!fnName) {
        fnName = state.program.importedHelpers[which] = addNamed(
          nodePath,
          helperMethodsToImportNames[which],
          "ts-runtime"
        )
      }

      const args = node.arguments
      // debugger
      return template.expression(`
        %%fnName%%(%%typeVarName%%)(%%args%%)
      `)({
        typeVarName,
        args,
        // @ts-ignore
        fnName: t.cloneNode(fnName)
      })
    }
  }
}

function typeReferenceToVar(
  node: TSTypeReference,
  state: AllStates
): Expression {
  const globalScope = state.program.nodePath.scope
  const refName = (node.typeName as $LibBugAny).name as string
  if (state.program.resolvedImportedTypes[refName]) {
    return state.program.resolvedImportedTypes[refName]
  }

  if (!globalScope.hasReference(refName)) {
    return template.expression(`{
      sort: "unresolvable",
      name: %%name%%
    }`)({name: refName})
  }

  const binding = globalScope.getBinding(refName)
  const importSpecifier = binding.path
  if (importSpecifier.type !== "ImportSpecifier") {
    throw new Error(`@todo unexpected import`)
  }
  const importDeclaration = importSpecifier.parentPath
  if (importDeclaration.type !== "ImportDeclaration") {
    throw new Error(`@todo unexpected import`)
  }

  const localizedVarName = "__imported_type__" + refName
  const exportedVarName = identNameToVarName(refName)

  // @ts-ignore
  importDeclaration.node.specifiers.push(
    t.importSpecifier(
      t.identifier(localizedVarName),
      t.identifier(exportedVarName)
    )
  )

  const refToVar = t.identifier(localizedVarName)
  state.program.resolvedImportedTypes[refName] = refToVar

  return refToVar
}

const helperMethodsToImportNames = {
  refine: "applyRefinement",
  is: "valueConformsToType",
  getTypeErrors: "getTypeErrors",
  assertType: "assertType"
}

const standardReferenceIdentifiers = {
  Array: "array",
  Partial: "Partial",
  Required: "Required",
  Readonly: "Readonly",
  Pick: "Pick",
  Record: "Record",
  Exclude: "Exclude",
  Extract: "Extract",
  Omit: "Omit",
  NonNullable: "NonNullable"
}