Typescript: Compiler/LanguageService API - Modifying AST before type checking

[fis-cz]

Hello,

I would like to know if it is possible to manipulate AST tree of the file before it is type checked. I know it is possible to get single file AST, manipulate and eventually emit it. I didn't find a way how to add manipulated file to program or manipulate file AST added to rootFiles of the program or manipulate / recheck a file already existing in the cache.

I am trying to create typescript preprocessor and include it to the build progress. As it is not possible to pass map files to typescript compiler any text manipulation would require pre/post processing steps in order to fit the TSC output file/map to the input file (and even this would be hard to achieve).

Only the one way I've found (and working) till now is to allow just simple code removals using //#if, #ifdef, #ifndef, #elseif and replacing removed characters by space character. This allows me to keep map files in sync but it is impossible to implement macros and eventually other features.

[RyanCavanaugh]

This isn't possible, but is something we're looking into as part of the plugins API investigation. Follow #16607 for updates

[typescript-bot]

This issue has been marked as 'Question' and has seen no recent activity. It has been automatically closed for house-keeping purposes. If you're still waiting on a response, questions are usually better suited to stackoverflow.

[fis-cz]

Actually, I have found it is possible to manipulate the AST before type checking using the compiler API / LanguageService. It is necessary to implement custom compilerHost getSourceFile method then it is possible to walk through the (incomplete, but sufficient AST) uding standard visitor pattern and modify it.

I wrote simple AST "preprocessor" allowing me to add, remove or replace AST nodes while if this is done in correct way the emitted code exactly matches what I expect it to be and TypeChecking works well too.

I just tested only a few things, for example I have source

let a: string = 1

Without pre-processor this gives error what is correct behavior.

With preprocessor which replaces StringKeyword with NumberKeyword in AST the compilation is without errors and the final emitted result is as expected.

Manipulation needs more work to be done (i.e. renaming the variable name needs to be reflected on more places in AST, such as in SourceFile, FirstStatement and so on).

But it seems to be working fine and it is possible to implement //#ifdef, or //#macro directives.

I'll write updates in this thread.

[fis-cz]

Hi, I have written simple TypeScript preprocessor working over AST tree before it is being type checked (this have multiple advantages than to work over the source text - the biggest one is related to map files).

It is using standard compiler API exclucding AST manipulation (as this is not supported by API at all at this time as far as I know)

Currently, only define, if, else and endif pragmas are supported. I have not found use case for macros yet but it would be possible to be implemented (with some problems with map files in case the macro will be defined in another than the "current" source file).

---

This is definitely not a production quality source code, it is just a proof of concept of the "TypeScript compiler plugin". I can imagine the final source file is synthesized based on original source rather than directly modified as I am doing it now so the idea of immutable AST tree can be kept. But this requires much more work than just a few lines of the code written bellow. And it would be nice if it is somehow directly supported by TS api (like create empty source file, copy node with filters and callbacks for children and so on).

---

When using the code be careful as expressions are evaluated using JS EVAL function.

All defines are shared through all files as they are going parsed through the compiler (so the root files from TS config are is first then all referenced or imported files follows. It means that defines from root file are available in all referenced files.

All AST nodes in the if-else-endif pragma block are removed in case the directive evalutes as false
All references and imports are also removed.
If the name if the identifier inside of the code matches the defined "constant" the identifier AST node will be replaced with appropriate value (currently string, number and boolean data types are supported so no arrays or objects)

---

Examples:

//TEST will be internally set to 'true'
//#define TEST

//#define PI 3.14159265359

//#define R 5

// Computes the result based on previously defined values
//#define CIRCUMFERENCE 2 * PI * R

//#define STR1 "Hello"

//#define STR2 "World"

// STR will contain "Hello World!"
//#define STR STR1 + STR2 + " " + STR3 + "!"

// CIRCUMFERENCE will be replaced by value from CIRCUMFERENCE define
let a: number = CIRCUMFERENCE;

// CIRCUMFERENCE will be replaced by value from CIRCUMFERENCE define
// Type error will be thrown by the compiler
let b: string = CIRCUMFERENCE;

// Files referenced based on defined value

//#define REF true

//#if REF
/// <reference path="./referencedA.ts" />
/// <reference types="./typesA" />
/// <reference lib="es5" />
//#else
/// <reference path="./referencedB.ts" />
/// <reference types="./typesB" />
/// <reference lib="es2020" />
//#endif


// Imports
//#if REF
import * as lib from "./libA";
//#else
import * as lib from "./libB";
//#endif

---

Current sources follows.
Comments are welcome.
Its a standard code of standalone TypeScript compiler for NodeJS using TS API - nothing special.
Only the one dependency is the TypeScript compiler from NPM.

---

src/index.ts

import * as path from "path";
import * as ts from "typescript";
import * as compilerHost from "./typescript/compilerHost";
import * as astPreprocessor from "./typescript/astPreprocessor";

let configFileName = "./tsconfig.json";

const host: ts.ParseConfigFileHost = <any>ts.sys;
const parsedCmd = ts.getParsedCommandLineOfConfigFile(configFileName, undefined, host);

const { options, fileNames } = parsedCmd;

const chost = compilerHost.create(options);
astPreprocessor.program();

console.log("CREATING PROGRAM");
const program = ts.createProgram(fileNames, options, chost);
console.log("PROGRAM CREATED");

console.log("CREATING DIAGNOSTICS");
let diag = ts.getPreEmitDiagnostics(program);
console.log("DIAGNOSTICS CREATED");

for (const m of diag) {
   const lineChar: ts.LineAndCharacter = m.file.getLineAndCharacterOfPosition(m.start);
   const rfn: string = path.relative(process.cwd(), m.file.fileName);
   const err = `${rfn}:${lineChar.line}:${lineChar.character} - ${ts.DiagnosticCategory[m.category]} ${m.code}: ${m.messageText}`;
   console.log(err);
}

const emit = program.emit();

console.log("DONE");

---

src/typescript/compilerHost.ts

import * as ts from "typescript";
import * as astPreprocessor from "./astPreprocessor";

export function create(options: ts.CompilerOptions): ts.CompilerHost {

   function getSourceFile(fileName: string, languageVersion: ts.ScriptTarget, onError?: (message: string) => void, shouldCreateNewSourceFile?: boolean): ts.SourceFile | undefined {

      // console.log(`Creating source file '${fileName}'`);
      const sourceText = ts.sys.readFile(fileName);
      const sourceFile = sourceText !== undefined ? ts.createSourceFile(fileName, sourceText, languageVersion) : undefined;

      if (sourceFile && !sourceFile.isDeclarationFile) {
         astPreprocessor.preprocessSourceFile(sourceFile);
      }

      return sourceFile;

   }

   const compilerHost = ts.createCompilerHost(options, false);
   compilerHost.getSourceFile = getSourceFile;
   return compilerHost;

}

---

src/typescript/astPreprocessor.ts

import * as ts from "typescript";

interface Define {
   [name: string]: number | string | boolean;
}

enum CommentType {
   Inline = 0,
   Block = 1
}

let defines: Define = {};
let currentSource: ts.SourceFile = undefined;

let evalDirectives: boolean[];
let lastParsedDirectiveEndPos: number;
let nodesToDelete: { parent: ts.Node, node: ts.Node }[];

export function program(defs?: Define): void {
   defines = defs || {};
   defines["TRUE"] = true;
   defines["FALSE"] = false;
   currentSource = undefined;
}

export function preprocessSourceFile(sourceFile: ts.SourceFile): void {

   currentSource = sourceFile;
   evalDirectives = [ true ];
   lastParsedDirectiveEndPos = 0;
   nodesToDelete = [];

   ts.forEachChild(currentSource, (node: ts.Node) => { processNode(node, sourceFile); });

   for (const ntd of nodesToDelete) {
      deleteNode(ntd.parent, ntd.node);
   }

   console.log("Processing file done!");
}

function processNode(node: ts.Node, parent: ts.Node): void {

   const text = currentSource.text.substr(node.pos, node.end - node.pos);

   console.log(`Processing node: ${ ts.SyntaxKind[node.kind] } ::: ${text.replace(/\n/g, " ")}`);
   console.log("");

   parseDirectives(text, node.pos);

   if (!evalDirectives[evalDirectives.length - 1]) {
      nodesToDelete.push({ parent, node });
   } else {
      if (node.kind === ts.SyntaxKind.Identifier) {
         const iName = text.trim();
         if (defines.hasOwnProperty(iName)) {
            if (typeof(defines[iName]) === "string") {
               const newNode = ts.createStringLiteral(defines[iName].toString());
               newNode.pos = node.pos;
               replaceNode(parent, node, newNode);
            }
            if (typeof(defines[iName]) === "number") {
               const newNode = ts.createNumericLiteral(defines[iName].toString());
               newNode.pos = node.pos;
               replaceNode(parent, node, newNode);
            }
            if (typeof(defines[iName]) === "boolean") {
               const newNode = ts.createNode(defines[iName] ? ts.SyntaxKind.TrueKeyword : ts.SyntaxKind.FalseKeyword, node.pos, -1);
               replaceNode(parent, node, newNode);
            }
         }
      }

      ts.forEachChild(node, (child: ts.Node) => { processNode(child, node); });
   }

}

function deleteReferencedFile(fileName: string): void {
   let i = 0;
   while (i < currentSource.referencedFiles.length) {
      if (currentSource.referencedFiles[i].fileName === fileName) {
         (<Array<any>>currentSource.referencedFiles).splice(i, 1);
      } else {
         i++;
      }
   }
}

function deleteTypeReferenceDirective(fileName: string): void {
   let i = 0;
   while (i < currentSource.typeReferenceDirectives.length) {
      if (currentSource.typeReferenceDirectives[i].fileName === fileName) {
         (<Array<any>>currentSource.typeReferenceDirectives).splice(i, 1);
      } else {
         i++;
      }
   }
}

function deleteLibReferenceDirective(fileName: string): void {
   let i = 0;
   while (i < currentSource.libReferenceDirectives.length) {
      if (currentSource.libReferenceDirectives[i].fileName === fileName) {
         (<Array<any>>currentSource.libReferenceDirectives).splice(i, 1);
      } else {
         i++;
      }
   }
}


function deleteNode(parent: ts.Node, node: ts.Node): void {

   for (const key in parent) {
      if ((<any>parent)[key] instanceof Array) {
         for (let i = 0; i < (<any>parent)[key].length; i++) {
            if ((<any>parent)[key][i] === node) {
               (<any>parent)[key].splice(i, 1);
               return;
            }
         }
      } else if ((<any>parent)[key] === node) {
         delete (<any>parent)[key];
         return;
      }
   }

   releaseNodeTree(node);

}

function replaceNode(parent: ts.Node, oldNode: ts.Node, newNode: ts.Node): void {
   for (const key in parent) {
      if ((<any>parent)[key] instanceof Array) {
         for (let i = 0; i < (<any>parent)[key].length; i++) {
            if ((<any>parent)[key][i] === oldNode) {
               (<any>parent)[key][i] = newNode;
               return;
            }
         }
      } else if ((<any>parent)[key] === oldNode) {
         (<any>parent)[key] = newNode;
         return;
      }
   }

   releaseNodeTree(oldNode);

}

function releaseNodeTree(node: ts.Node): void {

}


function parseDirectives(text: string, nodePos: number): void {

   let pos: number = 0;
   let inComment: boolean = false;
   let blockCounter: number = 0;
   let commentType: CommentType = CommentType.Inline;
   let readingDirective: boolean = false;
   let directive: string = "";

   while (pos < text.length) {

      if (!inComment && text.substr(pos, 3) === "///") {

         let tsd = "";
         pos += 3;
         while (text[pos] !== "\n" && text[pos] !== "\r") {
            tsd += text[pos];
            pos++;
         }

         if (!evalDirectives[evalDirectives.length - 1]) {

            tsd = tsd.trim();

            let rxa: RegExpMatchArray;

            rxa = tsd.match(/<reference\s*path[\s?]*=[\s?]*"(.*)"[\s?]*\/>/);
            if (rxa && rxa.length === 2) { deleteReferencedFile(rxa[1]); }
            rxa = tsd.match(/<reference\s*types[\s?]*=[\s?]*"(.*)"[\s?]*\/>/);
            if (rxa && rxa.length === 2) { deleteTypeReferenceDirective(rxa[1]); }
            rxa = tsd.match(/<reference\s*lib[\s?]*=[\s?]*"(.*)"[\s?]*\/>/);
            if (rxa && rxa.length === 2) { deleteLibReferenceDirective(rxa[1]); }

         }

      }

      if (!inComment && text.substr(pos, 2) === "//") {
         inComment = true;
         commentType = CommentType.Inline
         pos++;
      }

      if (!inComment && text.substr(pos, 2) === "/*") {
         inComment = true;
         commentType = CommentType.Block
         pos++;
      }

      if (!inComment && text[pos] === "{") {
         blockCounter++;
      }

      if (!inComment && text[pos] === "}") {
         blockCounter--;
      }

      if (inComment && commentType === CommentType.Inline && (text[pos] === "\r" || text[pos] === "\n")) {
         inComment = false;
      }

      if (inComment && commentType === CommentType.Block && text.substr(pos, 2) === "*/") {
         inComment = false;
      }

      if (!blockCounter && inComment && commentType === CommentType.Inline && text[pos] === "#" && nodePos + pos >= lastParsedDirectiveEndPos) {
         directive = "";
         readingDirective = true;
         pos++;
      }

      if (readingDirective) {
         if (text[pos] === "\r" || text[pos] ==="\n") {
            readingDirective = false;
            lastParsedDirectiveEndPos = nodePos + pos;
            parseDirective(directive);
         } else {
            directive += text[pos];
         }
      }

      pos++;

   }

}

function parseDirective(directive: string): void {

   interface Directives {
      [key: string]: (expression?: string) => void;
   }

   const directives: Directives = {
      "define\\s(.*)$": (expression: string) => {
         if (evalDirectives[evalDirectives.length - 1]) {
            evalDefineExpression(expression);
         }
      },

      "if\\s(.*)$": (expression: string) => {
         if (evalDirectives[evalDirectives.length - 1]) {
            evalDirectives.push(evalIfExpression(expression));
         } else {
            evalDirectives.push(false);
         }
      },

      "else\s*$": () => {
         if (evalDirectives.length === 0 || evalDirectives[evalDirectives.length - 2]) evalDirectives[evalDirectives.length - 1] = !evalDirectives[evalDirectives.length - 1];
      },

      "endif\s*$": () => {
         if (evalDirectives.length > 1) {
            evalDirectives.pop();
         } else {
            console.error("Missing if directive!");
         }
      }
   }

   for (const dk in directives) {
      const rx = new RegExp(dk);
      const rxa = directive.match(rx);
      if (rxa) {
         const expr = rxa.length === 2 ? rxa[1] : undefined;
         directives[dk](expr);
      }
   }

}

function tokenizeExpression(expression: string): any[] {

   let tokenized: any[] = [];

   let pos = 0
   let token = "";
   let inString = false;

   function addToken() {
      if (token !== "") {
         tokenized.push(token);
         token = "";
      }
   }

   while(pos < expression.length) {

      const ch: string = expression[pos];
      const ch2: string = expression.substr(pos, 2);

      // space, tab
      if (!inString && /\s/.test(ch)) {
         addToken();
      }

      // double char operators <= >= == != && ||
      else if (!inString && /\<\=|\>\=|\=\=|\!\=|\&\&|\|\|/.test(ch2) ) {
         addToken();
         tokenized.push(ch2);
         pos++;
      }


      // single char operators and expression characters
      else if (!inString && /\+|\-|\*|\/|\\|\^|\~|\!|\%|\&|\(|\)|\[|\]/.test(ch)) {
         addToken();
         tokenized.push(ch);
      }

      else if (!inString && ch === "\"") {
         inString = true;
         token += ch;
      }

      else if (inString && ch === "\"") {
         inString = false;
         token += ch;
      }

      else {
         token += ch;
      }

      pos++;

   }

   addToken();

   return tokenized;

}

function evalDefineExpression(expression: string): void {

   const expr: any[] = tokenizeExpression(expression);

   if (expr.length < 1) {
      console.error("Invalid define expression")
   }

   if (defines.hasOwnProperty(expr[0])) {
      console.log(`Warning: Preprocessor variable ${expr[0]} defined already`);
   }

   let key = expr[0]
   let value = true;

   if (expr.length > 1) {

      for (let i = 1; i < expr.length; i++) {
         if (defines.hasOwnProperty(expr[i])) {
            if (typeof(defines[expr[i]]) === "string") {
               expr[i] = `"${defines[expr[i]]}"`;
            } else {
               expr[i] = defines[expr[i]];
            }
         }
      }

      expr.splice(0, 1);
      const evalExpr = "value = " + expr.join(" ");

      try {
         eval(evalExpr);
      } catch (e) {
         console.error(e);
      }

   }

   defines[key] = value;

}

function evalIfExpression(expression: string): boolean {

   const expr: any[] = tokenizeExpression(expression);

   let value;

   if (expr.length > 0) {

      for (let i = 0; i < expr.length; i++) {
         if (defines.hasOwnProperty(expr[i])) {
            if (typeof(defines[expr[i]]) === "string") {
               expr[i] = `"${defines[expr[i]]}"`;
            } else {
               expr[i] = defines[expr[i]].toString();
            }
         }
      }

      const evalExpr = "value = " + expr.join(" ");

      try {
         eval(evalExpr);
      } catch (e) {
         console.error(e);
      }

   }

   return value;

}

program();

---

I have also idea to have following plugins which most of them should be possible to implement relatively simply using the TS API. It would be great if we can run them from tsconfig such as we can do with LanguageService plugins.

• Source code preprocessor (text preprocessor) is easy to implement but it has many limitations as the map file will not match original source code if something will be added to the source text (removal can be done using replacing removed text with appropriate number of correct whitespacing)
• AST code preprocessor (done here) - mainly for conditional compilation - it can help to minimize the final code footprint while it makes it possible to include additional code during debugging. In easy way, not like I will need special tool for every new project. I will also try to implement LanduageService plugin based on the code presented here so the Intellisense and correct conditional compilation can be seen in the editor too, including defines and their values.
• AST code postprocessor - especially for code optimization and minification. Based on knowledge of the whole program including data types, exports, imports and so on it should be possible to perform the best possible minification and optimization of the TS code.
• Emitter... TypeScript emiter? Babel like emitter? So many options here. For me it is not necessary to write any as I usually go with esnext profile and I am using Babel if I need something lower.