Codegen from ReactiveFunction (no memoization yet)

Currently codegen operates from HIR using a tree visitor, but for scope 
construction we're converting the HIR (CFG) into a ReactiveFunction (AST-like). 
Our original idea for codegen was that we would convert the ReactiveFunction 
back to HIR, and then codegen from there. However, the ReactiveFunction is 
already in tree form...which makes it very straightforward to generate code 
from. 

So this PR implements codegen from ReactiveFunction. The output is _identical_ 
thanks in large part to reusing as much logic from Codegen.ts as possible. The 
next PR will add memoization logic.

compiler/forget/src/HIR/Codegen.ts

@@ -44,13 +44,13 @@ function withLoc<TNode extends t.Node, T extends (...args: any[]) => TNode>(
   };
 }
 
-const createBinaryExpression = withLoc(t.binaryExpression);
-const createCallExpression = withLoc(t.callExpression);
-const createExpressionStatement = withLoc(t.expressionStatement);
-const createFunctionDeclaration = withLoc(t.functionDeclaration);
-const createLabelledStatement = withLoc(t.labeledStatement);
-const createVariableDeclaration = withLoc(t.variableDeclaration);
-const createWhileStatement = withLoc(t.whileStatement);
+export const createBinaryExpression = withLoc(t.binaryExpression);
+export const createCallExpression = withLoc(t.callExpression);
+export const createExpressionStatement = withLoc(t.expressionStatement);
+export const createFunctionDeclaration = withLoc(t.functionDeclaration);
+export const createLabelledStatement = withLoc(t.labeledStatement);
+export const createVariableDeclaration = withLoc(t.variableDeclaration);
+export const createWhileStatement = withLoc(t.whileStatement);
 
 /**
  * Converts HIR into Babel nodes, which can then be printed into source text.
@@ -87,7 +87,7 @@ export default function codegen(fn: HIRFunction): t.Function {
   );
 }
 
-type Temporaries = Map<IdentifierId, t.Expression>;
+export type Temporaries = Map<IdentifierId, t.Expression>;
 
 class CodegenVisitor
   implements
@@ -189,45 +189,7 @@ class CodegenVisitor
     return codegenInstructionValue(this.temp, value);
   }
   visitInstruction(instr: Instruction, value: t.Expression): t.Statement {
-    if (t.isStatement(value)) {
-      return value;
-    }
-    if (instr.lvalue === null) {
-      return t.expressionStatement(value);
-    }
-    if (
-      instr.lvalue.place.memberPath === null &&
-      instr.lvalue.place.identifier.name === null
-    ) {
-      // temporary
-      this.temp.set(instr.lvalue.place.identifier.id, value);
-      return t.emptyStatement();
-    } else {
-      switch (instr.lvalue.kind) {
-        case InstructionKind.Const: {
-          return createVariableDeclaration(instr.loc, "const", [
-            t.variableDeclarator(codegenLVal(instr.lvalue), value),
-          ]);
-        }
-        case InstructionKind.Let: {
-          return createVariableDeclaration(instr.loc, "let", [
-            t.variableDeclarator(codegenLVal(instr.lvalue), value),
-          ]);
-        }
-        case InstructionKind.Reassign: {
-          return createExpressionStatement(
-            instr.loc,
-            t.assignmentExpression("=", codegenLVal(instr.lvalue), value)
-          );
-        }
-        default: {
-          assertExhaustive(
-            instr.lvalue.kind,
-            `Unexpected instruction kind '${instr.lvalue.kind}'`
-          );
-        }
-      }
-    }
+    return codegenInstruction(this.temp, instr, value);
   }
   visitTerminalId(id: InstructionId): void {}
   visitImplicitTerminal(): t.Statement | null {
@@ -325,11 +287,57 @@ class CodegenVisitor
   }
 }
 
-function codegenLabel(id: BlockId): string {
+export function codegenLabel(id: BlockId): string {
   return `bb${id}`;
 }
 
-function codegenInstructionValue(
+export function codegenInstruction(
+  temp: Temporaries,
+  instr: Instruction,
+  value: t.Expression
+): t.Statement {
+  if (t.isStatement(value)) {
+    return value;
+  }
+  if (instr.lvalue === null) {
+    return t.expressionStatement(value);
+  }
+  if (
+    instr.lvalue.place.memberPath === null &&
+    instr.lvalue.place.identifier.name === null
+  ) {
+    // temporary
+    temp.set(instr.lvalue.place.identifier.id, value);
+    return t.emptyStatement();
+  } else {
+    switch (instr.lvalue.kind) {
+      case InstructionKind.Const: {
+        return createVariableDeclaration(instr.loc, "const", [
+          t.variableDeclarator(codegenLVal(instr.lvalue), value),
+        ]);
+      }
+      case InstructionKind.Let: {
+        return createVariableDeclaration(instr.loc, "let", [
+          t.variableDeclarator(codegenLVal(instr.lvalue), value),
+        ]);
+      }
+      case InstructionKind.Reassign: {
+        return createExpressionStatement(
+          instr.loc,
+          t.assignmentExpression("=", codegenLVal(instr.lvalue), value)
+        );
+      }
+      default: {
+        assertExhaustive(
+          instr.lvalue.kind,
+          `Unexpected instruction kind '${instr.lvalue.kind}'`
+        );
+      }
+    }
+  }
+}
+
+export function codegenInstructionValue(
   temp: Temporaries,
   instrValue: InstructionValue
 ): t.Expression {
@@ -484,7 +492,7 @@ function codegenJsxElement(
   }
 }
 
-function codegenLVal(lval: LValue): t.LVal {
+export function codegenLVal(lval: LValue): t.LVal {
   const expr = convertIdentifier(lval.place.identifier);
   const memberPath = lval.place.memberPath;
   return memberPath == null
@@ -515,7 +523,7 @@ function codegenValue(
   }
 }
 
-function codegenPlace(temp: Temporaries, place: Place): t.Expression {
+export function codegenPlace(temp: Temporaries, place: Place): t.Expression {
   todoInvariant(place.kind === "Identifier", "support scope values");
   if (place.memberPath === null) {
     let tmp = temp.get(place.identifier.id);
@@ -532,7 +540,7 @@ function codegenPlace(temp: Temporaries, place: Place): t.Expression {
   }
 }
 
-function convertIdentifier(identifier: Identifier): t.Identifier {
+export function convertIdentifier(identifier: Identifier): t.Identifier {
   if (identifier.name !== null) {
     return t.identifier(`${identifier.name}$${identifier.id}`);
   }

compiler/forget/src/HIR/CodegenReactiveFunction.ts

@@ -0,0 +1,237 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import * as t from "@babel/types";
+import invariant from "invariant";
+import { assertExhaustive } from "../Common/utils";
+import {
+  codegenInstruction,
+  codegenInstructionValue,
+  codegenLabel,
+  codegenPlace,
+  convertIdentifier,
+  createFunctionDeclaration,
+  Temporaries,
+} from "./Codegen";
+import {
+  Instruction,
+  ReactiveBasicBlock,
+  ReactiveFunction,
+  ReactiveScope,
+  ReactiveTerminal,
+  ReactiveValueBlock,
+} from "./HIR";
+import { todoInvariant } from "./todo";
+
+export function codegenReactiveFunction(fn: ReactiveFunction): t.Function {
+  const cx = new Context();
+  const params = fn.params.map((param) => convertIdentifier(param.identifier));
+  const body = codegenBlock(cx, fn.body);
+  const statements = body.body;
+  if (statements.length !== 0) {
+    const last = statements[statements.length - 1];
+    if (last.type === "ReturnStatement" && last.argument == null) {
+      statements.pop();
+    }
+  }
+  return createFunctionDeclaration(
+    fn.loc,
+    fn.id !== null ? convertIdentifier(fn.id) : null,
+    params,
+    body,
+    fn.generator,
+    fn.async
+  );
+}
+
+class Context {
+  temp: Temporaries = new Map();
+}
+
+function codegenBlock(
+  cx: Context,
+  block: ReactiveBasicBlock
+): t.BlockStatement {
+  const statements: Array<t.Statement> = [];
+  for (const item of block) {
+    switch (item.kind) {
+      case "instruction": {
+        const statement = codegenInstructionNullable(
+          cx.temp,
+          item.instruction,
+          codegenInstructionValue(cx.temp, item.instruction.value)
+        );
+        if (statement !== null) {
+          statements.push(statement);
+        }
+        break;
+      }
+      case "block": {
+        codegenReactiveScope(cx, statements, item.scope, item.instructions);
+        break;
+      }
+      case "terminal": {
+        const statement = codegenTerminal(cx, item.terminal);
+        if (item.label !== null) {
+          statements.push(
+            t.labeledStatement(
+              t.identifier(codegenLabel(item.label)),
+              statement
+            )
+          );
+        } else {
+          statements.push(statement);
+        }
+        break;
+      }
+      default: {
+        assertExhaustive(item, `Unexpected item kind '${(item as any).kind}'`);
+      }
+    }
+  }
+  return t.blockStatement(statements);
+}
+
+function codegenReactiveScope(
+  cx: Context,
+  statements: Array<t.Statement>,
+  scope: ReactiveScope,
+  block: ReactiveBasicBlock
+): void {
+  // TODO @josephsavona: Emit memoized blocks!
+  const body = codegenBlock(cx, block).body;
+  statements.push(...body);
+}
+
+function codegenTerminal(cx: Context, terminal: ReactiveTerminal): t.Statement {
+  switch (terminal.kind) {
+    case "break": {
+      return t.breakStatement(
+        terminal.label !== null
+          ? t.identifier(codegenLabel(terminal.label))
+          : null
+      );
+    }
+    case "continue": {
+      return t.continueStatement(
+        terminal.label !== null
+          ? t.identifier(codegenLabel(terminal.label))
+          : null
+      );
+    }
+    case "for": {
+      return t.forStatement(
+        codegenForInit(cx, terminal.init),
+        codegenValueBlock(cx, terminal.test),
+        codegenValueBlock(cx, terminal.update),
+        codegenBlock(cx, terminal.loop)
+      );
+    }
+    case "if": {
+      return t.ifStatement(
+        codegenPlace(cx.temp, terminal.test),
+        codegenBlock(cx, terminal.consequent),
+        terminal.alternate !== null
+          ? codegenBlock(cx, terminal.alternate)
+          : null
+      );
+    }
+    case "return": {
+      return t.returnStatement(
+        terminal.value !== null ? codegenPlace(cx.temp, terminal.value) : null
+      );
+    }
+    case "switch": {
+      return t.switchStatement(
+        codegenPlace(cx.temp, terminal.test),
+        terminal.cases.map((case_) => {
+          const test =
+            case_.test !== null ? codegenPlace(cx.temp, case_.test) : null;
+          const block = codegenBlock(cx, case_.block!);
+          return t.switchCase(test, [block]);
+        })
+      );
+    }
+    case "throw": {
+      return t.throwStatement(codegenPlace(cx.temp, terminal.value));
+    }
+    case "while": {
+      const test = codegenValueBlock(cx, terminal.test);
+      return t.whileStatement(test, codegenBlock(cx, terminal.loop));
+    }
+    default: {
+      assertExhaustive(
+        terminal,
+        `Unexpected terminal kind '${(terminal as any).kind}'`
+      );
+    }
+  }
+}
+
+export function codegenInstructionNullable(
+  temp: Temporaries,
+  instr: Instruction,
+  value: t.Expression
+): t.Statement | null {
+  const statement = codegenInstruction(temp, instr, value);
+  if (statement.type === "EmptyStatement") {
+    return null;
+  }
+  return statement;
+}
+
+function codegenForInit(
+  cx: Context,
+  init: ReactiveValueBlock
+): t.Expression | t.VariableDeclaration | null {
+  const body = codegenBlock(cx, init.instructions).body;
+  if (init.value !== null) {
+    invariant(
+      body.length === 0,
+      "Expected for init block to produce only temporaries"
+    );
+    return codegenInstructionValue(cx.temp, init.value);
+  } else {
+    invariant(
+      body.length === 1,
+      "Expected for init to have a variable declaration"
+    );
+    const declaration = body[0]!;
+    invariant(
+      declaration.type === "VariableDeclaration",
+      "Expected a variable declaration"
+    );
+    return declaration;
+  }
+}
+
+function codegenValueBlock(
+  cx: Context,
+  block: ReactiveValueBlock
+): t.Expression {
+  const body = codegenBlock(cx, block.instructions).body;
+  const expressions = body.map((stmt) => {
+    if (stmt.type === "ExpressionStatement") {
+      return stmt.expression;
+    } else {
+      todoInvariant(false, `Handle conversion of ${stmt.type} to expression`);
+    }
+  });
+  if (block.value !== null) {
+    const value = codegenInstructionValue(cx.temp, block.value);
+    expressions.push(value);
+  }
+  invariant(
+    expressions.length !== 0,
+    "Expected a value block to produce one or more expressions"
+  );
+  if (expressions.length === 1) {
+    return expressions[0];
+  } else {
+    return t.sequenceExpression(expressions);
+  }
+}

compiler/forget/src/HIR/Pipeline.ts

@@ -12,13 +12,17 @@ import enterSSA from "../HIR/EnterSSA";
 import { Environment } from "../HIR/HIRBuilder";
 import inferReferenceEffects from "../HIR/InferReferenceEffects";
 import { leaveSSA } from "../HIR/LeaveSSA";
-import codegen from "./Codegen";
+import { buildReactiveFunction } from "./BuildReactiveFunction";
+import { codegenReactiveFunction } from "./CodegenReactiveFunction";
+import { flattenReactiveLoops } from "./FlattenReactiveLoops";
 import { HIRFunction } from "./HIR";
 import { inferMutableRanges } from "./InferMutableRanges";
 import { inferReactiveScopes } from "./InferReactiveScopes";
 import { inferReactiveScopeVariables } from "./InferReactiveScopeVariables";
 import { inferTypes } from "./InferTypes";
 import { logHIRFunction } from "./logger";
+import { printReactiveFunction } from "./PrintReactiveFunction";
+import { propagateScopeDependencies } from "./PropagateScopeDependencies";
 
 export type CompilerFlags = {
   eliminateRedundantPhi: boolean;
@@ -35,6 +39,7 @@ export type CompilerFlags = {
 export type CompilerResult = {
   ir: HIRFunction;
   ast: t.Function | null;
+  scopes: string | null;
 };
 
 export default function (
@@ -83,11 +88,17 @@ export default function (
   }
 
   if (flags.codegen) {
+    const reactiveFunction = buildReactiveFunction(ir);
+    flattenReactiveLoops(reactiveFunction);
+    propagateScopeDependencies(reactiveFunction);
+    const scopes = printReactiveFunction(reactiveFunction);
+    const ast = codegenReactiveFunction(reactiveFunction);
     return {
-      ast: codegen(ir),
-      ir: ir,
+      ast,
+      ir,
+      scopes,
     };
   }
 
-  return { ast: null, ir: ir };
+  return { ast: null, scopes: null, ir: ir };
 }

compiler/forget/src/__tests__/hir-test.ts

@@ -14,13 +14,9 @@ import { wasmFolder } from "@hpcc-js/wasm";
 import invariant from "invariant";
 import path from "path";
 import prettier from "prettier";
-import { buildReactiveFunction } from "../HIR/BuildReactiveFunction";
-import { flattenReactiveLoops } from "../HIR/FlattenReactiveLoops";
 import { toggleLogging } from "../HIR/logger";
 import run from "../HIR/Pipeline";
 import { printFunction } from "../HIR/PrintHIR";
-import { printReactiveFunction } from "../HIR/PrintReactiveFunction";
-import { propagateScopeDependencies } from "../HIR/PropagateScopeDependencies";
 import generateTestsFromFixtures from "./test-utils/generateTestsFromFixtures";
 
 function wrapWithTripleBackticks(s: string, ext?: string) {
@@ -137,7 +133,7 @@ function transform(text: string, file: string): Array<TestOutput> {
   traverse(ast, {
     FunctionDeclaration: {
       enter(nodePath) {
-        const { ir, ast } = run(nodePath, {
+        const { ir, scopes, ast } = run(nodePath, {
           eliminateRedundantPhi: true,
           inferReferenceEffects: true,
           inferTypes: true,
@@ -149,14 +145,9 @@ function transform(text: string, file: string): Array<TestOutput> {
           codegen: true,
         });
 
-        const reactiveFunction = buildReactiveFunction(ir);
-        flattenReactiveLoops(reactiveFunction);
-        propagateScopeDependencies(reactiveFunction);
-        const scopes = printReactiveFunction(reactiveFunction);
-
         const textHIR = printFunction(ir);
-
-        invariant(ast !== null, "ast is null when codegen option is enabled");
+        invariant(ast, "Expected an ast");
+        invariant(scopes, "Expected printed scope data");
         const text = prettier.format(generate(ast).code.replace("\n\n", "\n"), {
           semi: true,
           parser: "babel-ts",