Add postInitialize transform step (#1145)

Author: jtenner

cli/asc.js

@@ -590,6 +590,20 @@ exports.main = function main(argv, options, callback) {
   optimizeLevel = Math.min(Math.max(optimizeLevel, 0), 3);
   shrinkLevel = Math.min(Math.max(shrinkLevel, 0), 2);
 
+  try {
+    stats.compileTime += measure(() => {
+      assemblyscript.initializeProgram(program, compilerOptions);
+    });
+  } catch(e) {
+    return callback(e);
+  }
+
+  // Call afterInitialize transform hook
+  {
+    let error = applyTransform("afterInitialize", program);
+    if (error) return callback(error);
+  }
+
   var module;
   stats.compileCount++;
   try {

cli/transform.d.ts

@@ -35,6 +35,9 @@ export abstract class Transform {
   /** Called when parsing is complete, before a program is instantiated from the AST. */
   afterParse?(parser: Parser): void;
 
+  /** Called after the program is instantiated. */
+  afterInitialize?(program: Program): void;
+
   /** Called when compilation is complete, before the module is being validated. */
   afterCompile?(module: Module): void;
 }

src/compiler.ts

@@ -376,14 +376,19 @@ export class Compiler extends DiagnosticEmitter {
     module.setFeatures(featureFlags);
   }
 
+  initializeProgram(): void {
+    // initialize lookup maps, built-ins, imports, exports, etc.
+    this.program.initialize(this.options);
+  }
+
   /** Performs compilation of the underlying {@link Program} to a {@link Module}. */
   compile(): Module {
     var options = this.options;
     var module = this.module;
     var program = this.program;
 
-    // initialize lookup maps, built-ins, imports, exports, etc.
-    program.initialize(options);
+    // check and perform this program initialization if it hasn't been done
+    this.initializeProgram();
 
     // set up the main start function
     var startFunctionInstance = program.makeNativeFunction(BuiltinNames.start, new Signature(program, [], Type.void));

src/index.ts

@@ -214,6 +214,11 @@ export function getDependee(program: Program, file: string): string | null {
 
 // Compiler
 
+/** Initializes the program pre-emptively for transform hooks. */
+export function initializeProgram(program: Program, options: Options): void {
+  program.initialize(options);
+}
+
 /** Compiles the parsed sources to a module. */
 export function compile(program: Program): Module {
   program.parser.finish();

src/program.ts

@@ -491,6 +491,8 @@ export class Program extends DiagnosticEmitter {
   nextClassId: u32 = 0;
   /** Next signature id. */
   nextSignatureId: i32 = 0;
+  /** An indicator if the program has been initialized. */
+  initialized: bool = false;
   /** Constructs a new program, optionally inheriting parser diagnostics. */
   constructor(
     /** Compiler options. */
@@ -646,6 +648,10 @@ export class Program extends DiagnosticEmitter {
 
   /** Initializes the program and its elements prior to compilation. */
   initialize(options: Options): void {
+    // Initialize only once
+    if (this.initialized) return;
+
+    this.initialized = true;
     this.options = options;
 
     // register native types