How to transform TypeScript code before type-checking using Compiler API

The Intent

I want to use the TypeScript's Compiler API to experiment with operator overloading in TypeScript code. Specifically, I want to find all instances of x + y and turn them into op_add(x, y). However, I want the language services (eg. IntelliSense in VS Code) to be aware of the transformation and show the correct types.

For example in this code:

interface Vector2 { x: number, y: number }
declare function op_add(x: Vector2, y: Vector2): Vector2
declare let a: Vector2, b: Vector2

let c = a + b

I would expect that when I hover my mouse over c, it would show Vector2.

The Plan

In order to achieve this, I'll have to:

Create a program that exposes the same API as typescript – in the same way that ttypescript does.
Make that program modify the source code before passing it to typescript
Make VS Code (or whatever editor) use my package instead of typescript

The Exectuion

I started by creating a short script called compile.ts that uses the Compiler API to parse a file called sample.ts into AST. Then it directly modifies the AST and changes Binary(x, PlusToken, y) to Call(op_add, x, y). Finally it prints the modified code to console and then tries to emit. This alone isn't enough for an IDE integration, but it is a good start.

compile.ts:

import * as ts from "typescript"
import { possibleChildProperties } from "./visit";

let program = ts.createProgram(['sample.ts'], { target: ts.ScriptTarget.ES5, module: ts.ModuleKind.CommonJS })
let inputFiles = program.getSourceFiles()
const printer = ts.createPrinter({ newLine: ts.NewLineKind.LineFeed })

let outputCode: string

for (let input of inputFiles) {
  if (input.fileName === 'sample.ts') {
    ts.visitNode(input, visitor) // modifies input's AST
    outputCode = printer.printNode(ts.EmitHint.Unspecified, input, input)
    break
  }
}

console.log(outputCode) // works
let emitResult = program.emit() // fails



function visitor(node: ts.Node): ts.Node {
  if (node.kind === ts.SyntaxKind.BinaryExpression) {
    let expr = node as ts.BinaryExpression

    if (expr.operatorToken.kind === ts.SyntaxKind.PlusToken) {
      return ts.createCall(ts.createIdentifier('op_add'), [], [expr.left, expr.right])
    }
  }

  return visitChildren(node, visitor)
}

function visitChildren(node: ts.Node, visitor: ts.Visitor) {
  for (const prop of possibleChildProperties) {
    if (node[prop] !== undefined) {
      if (Array.isArray(node[prop]))
        node[prop] = node[prop].map(visitor)
      else
        node[prop] = visitor(node[prop])
    }
  }

  return node
}

sample.ts:

let a = { a: 4 }
let b = { b: 3 }
let c = a + b

console.log output:

let a = { a: 4 };
let b = { b: 3 };
let c = op_add(a, b);

The Problem

While the code printer works fine and outputs the correct code, calling program.emit() results in an unspecified internal error. This probably means that I'm modifying the AST in an non-supported way.

/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:100920
                throw e;
                ^

Error: start < 0
    at createTextSpan (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:10559:19)
    at Object.createTextSpanFromBounds (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:10568:16)
    at getErrorSpanForNode (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:13914:19)
    at createDiagnosticForNodeInSourceFile (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:13808:20)
    at Object.createDiagnosticForNode (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:13799:16)
    at error (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:35703:22)
    at resolveNameHelper (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:36602:29)
    at resolveName (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:36274:20)
    at getResolvedSymbol (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:52602:21)
    at checkIdentifier (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:54434:26)

The Question

What is the correct way to modify a program's AST before running the type checker? I understand that AST should preferably be read-only, but the standard ts.visitEachChild can be only used after type-checking. And deep-cloning the nodes myself also doesn't seem like an viable option, as there isn't any way to create a Program from code-generated AST.

Updates

EDIT 1: As @jdaz noticed, my sample.ts was missing a declaration for op_add, which could be causing problems. I added this line to the top of the file:

declare function op_add(x: {}, y: {}): string

Now there's a different error – the generation of file diagnostic fails:

/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:100920
                throw e;
                ^

Error: Debug Failure. Expected -2 >= 0
    at Object.createFileDiagnostic (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:17868:18)
    at grammarErrorAtPos (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:69444:36)
    at checkGrammarForAtLeastOneTypeArgument (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:68771:24)
    at checkGrammarTypeArguments (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:68777:17)
    at checkCallExpression (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:59255:18)
    at checkExpressionWorker (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:61687:28)
    at checkExpression (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:61597:38)
    at checkExpressionCached (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:61275:38)
    at checkVariableLikeDeclaration (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:63983:69)
    at checkVariableDeclaration (/home/m93a/Dokumenty/tsc-experiments/node_modules/typescript/lib/typescript.js:64051:20)

Solution

You got close with your code. The first issue you seem to be having is source code file checks that occur, essentially the Debug Failure. Expected -2 >= 0 error is saying that when trying to match the AST to source code it failed.

The second issue is that you need to modify the existing AST tree whereas visitNode is generating a new AST tree. This also must be done as early as possible (before emit is called AFAIK) otherwise the TypeChecker might use the original AST instead of your updated AST.

Below is an example of your visitor function that should solve both problems. Note that this is really hacky and fragile, expect it to break often.

OLD:

function visitor(node: ts.Node): ts.Node {
  if (node.kind === ts.SyntaxKind.BinaryExpression) {
    let expr = node as ts.BinaryExpression

    if (expr.operatorToken.kind === ts.SyntaxKind.PlusToken) {
      return ts.createCall(ts.createIdentifier('op_add'), [], [expr.left, expr.right])
    }
  }

  return visitChildren(node, visitor)
}

NEW:

function visitor(node: ts.Node): ts.Node {
  if (node.kind === ts.SyntaxKind.BinaryExpression) {
    let expr = node as ts.BinaryExpression;

    if (expr.operatorToken.kind === ts.SyntaxKind.PlusToken) {
      const newIdentifierNode = ts.createIdentifier('op_add');
      const newCallNode = ts.createCall(newIdentifierNode, [], [expr.left, expr.right]);
      newCallNode.flags = node.flags;
      newCallNode.pos = node.pos;
      newCallNode.end = node.end;
      newCallNode.parent = node.parent;
      newCallNode.typeArguments = undefined;

      Object.getOwnPropertyNames(node).forEach((prop) => {
          delete node[prop];
      });
      Object.getOwnPropertyNames(newCallNode).forEach((prop) => {
          node[prop] = newCallNode[prop];
      });
      return node;
    }
  }

  return visitChildren(node, visitor);
}