From 779947520d0aa8fd75deb0f1440f61004299f5cd Mon Sep 17 00:00:00 2001
From: Jan Kassens <jkassens@meta.com>
Date: Thu, 10 Nov 2022 15:45:51 -0500
Subject: [PATCH] [playground] new HIR tab

Instead of hacking into the babel plugin passes, this now takes a new approach
for the HIR tab:

- The different passes are exported from the babel plugin (for simplicity)

- The tab actually runs the compiler steps based on local config in the tab.

- Re-purposed the CompilerFlags component to configure what passes to run. This
is currently mostly causing different errors, but could be useful going forward
as a direction.
---
 .../components/Editor/HIRTabContent.tsx       | 149 ++++++++++++++++++
 .../playground/components/Editor/Output.tsx   |   3 +-
 .../playground/components/Editor/index.tsx    |   8 +-
 compiler/forget/src/CompilerDriver.ts         |   5 -
 compiler/forget/src/CompilerOutputs.ts        |   9 +-
 compiler/forget/src/HIR/DumpHIRPass.ts        |  56 -------
 compiler/forget/src/index.ts                  |  45 ++++++
 7 files changed, 202 insertions(+), 73 deletions(-)
 create mode 100644 compiler/forget/packages/playground/components/Editor/HIRTabContent.tsx
 delete mode 100644 compiler/forget/src/HIR/DumpHIRPass.ts

diff --git a/compiler/forget/packages/playground/components/Editor/HIRTabContent.tsx b/compiler/forget/packages/playground/components/Editor/HIRTabContent.tsx
new file mode 100644
index 0000000000..e164fde0e4
--- /dev/null
+++ b/compiler/forget/packages/playground/components/Editor/HIRTabContent.tsx
@@ -0,0 +1,149 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import clsx from "clsx";
+
+import { HIR } from "babel-plugin-react-forget";
+import { useState } from "react";
+
+const {
+  parseFunctions,
+  Environment,
+  enterSSA,
+  eliminateRedundantPhi,
+  inferReferenceEffects,
+  inferMutableRanges,
+  leaveSSA,
+  lower,
+  printHIR,
+} = HIR;
+
+export default function HIRTabContent({ source }: { source: string }) {
+  const astFunctions = parseFunctions(source);
+
+  const [flags, setFlags] = useState({
+    eliminateRedundantPhi: true,
+    inferReferenceEffects: true,
+    inferMutableRanges: true,
+    leaveSSA: true,
+  });
+
+  return (
+    <div className="h-full flex flex-col">
+      <div className="flex-1">
+        {astFunctions.map((func, index): React.ReactNode => {
+          let body;
+          try {
+            const env = new Environment();
+            const ir = lower(func, env);
+            enterSSA(ir, env);
+            if (flags.eliminateRedundantPhi) {
+              eliminateRedundantPhi(ir);
+            }
+            if (flags.inferReferenceEffects) {
+              inferReferenceEffects(ir);
+            }
+            if (flags.inferMutableRanges) {
+              inferMutableRanges(ir);
+            }
+            if (flags.leaveSSA) {
+              leaveSSA(ir);
+            }
+            const textHIR = printHIR(ir.body);
+            body = <pre>{textHIR}</pre>;
+          } catch (e: any) {
+            body = <div>error: ${e.toString()}</div>;
+          }
+          const name = func.node.id?.name ?? "anonymous";
+          return (
+            <div key={index} className="border-b py-4 px-2">
+              <h3 className="font-medium mb-2">Function: {name}</h3>
+              {body}
+            </div>
+          );
+        })}
+      </div>
+      <CompilerFlagsEditor flags={flags} setFlags={setFlags} />
+    </div>
+  );
+}
+
+type Flags = {
+  eliminateRedundantPhi: boolean;
+  inferReferenceEffects: boolean;
+  inferMutableRanges: boolean;
+  leaveSSA: boolean;
+};
+
+function CompilerFlagsEditor({
+  flags,
+  setFlags,
+}: {
+  flags: Flags;
+  setFlags: (flags: Flags) => void;
+}) {
+  return (
+    <div className="flex flex-col gap-2 px-5 py-4 border-t border-gray-200">
+      <h3 className="font-medium">Compiler Options</h3>
+      <div className="flex flex-col">
+        <LabeledCheckbox
+          label="Eliminate Redundant Phi"
+          value={flags.eliminateRedundantPhi}
+          onChange={(eliminateRedundantPhi) => {
+            setFlags({ ...flags, eliminateRedundantPhi });
+          }}
+        />
+        <LabeledCheckbox
+          label="Infer Reference Effects"
+          value={flags.inferReferenceEffects}
+          onChange={(inferReferenceEffects) => {
+            setFlags({ ...flags, inferReferenceEffects });
+          }}
+        />
+        <LabeledCheckbox
+          label="Infer Mutable Ranges"
+          value={flags.inferMutableRanges}
+          onChange={(inferMutableRanges) => {
+            setFlags({ ...flags, inferMutableRanges });
+          }}
+        />
+        <LabeledCheckbox
+          label="Leave SSA"
+          value={flags.leaveSSA}
+          onChange={(leaveSSA) => {
+            setFlags({ ...flags, leaveSSA });
+          }}
+        />
+      </div>
+    </div>
+  );
+}
+
+function LabeledCheckbox({
+  label,
+  value,
+  onChange,
+}: {
+  label: string;
+  value: boolean;
+  onChange: (value: boolean) => void;
+}) {
+  return (
+    <label className="flex gap-1.5 text-base">
+      <input
+        type="checkbox"
+        checked={value}
+        onChange={(e) => {
+          onChange(e.target.checked);
+        }}
+      />
+      <span className={clsx({ "text-gray-500": !value, "text-black": value })}>
+        {label}
+      </span>
+    </label>
+  );
+}
diff --git a/compiler/forget/packages/playground/components/Editor/Output.tsx b/compiler/forget/packages/playground/components/Editor/Output.tsx
index 11d40cfbbc..efc8125eeb 100644
--- a/compiler/forget/packages/playground/components/Editor/Output.tsx
+++ b/compiler/forget/packages/playground/components/Editor/Output.tsx
@@ -19,6 +19,7 @@ import { getSelectedFile, type Store } from "../../lib/stores";
 import GraphView from "../GraphView";
 import Preview from "../Preview";
 import { monacoOptions } from "./monacoOptions";
+import HIRTabContent from "./HIRTabContent";
 
 const MemoizedOutput = memo(Output);
 
@@ -43,7 +44,7 @@ function Output({ store, updateDiagnostics }: Props) {
       tabs={{
         Preview: <Preview store={store} />,
         IR: <TextTabContent outputs={outputs} kind={OutputKind.IR} />,
-        HIR: <TextTabContent outputs={outputs} kind={OutputKind.HIR} />,
+        HIR: <HIRTabContent source={getSelectedFile(store).content} />,
         CFG: <TextTabContent outputs={outputs} kind={OutputKind.CFG} />,
         ValGraph: (
           <GraphTabContent outputs={outputs} kind={OutputKind.ValGraph} />
diff --git a/compiler/forget/packages/playground/components/Editor/index.tsx b/compiler/forget/packages/playground/components/Editor/index.tsx
index 1e45c80a5d..be4657098f 100644
--- a/compiler/forget/packages/playground/components/Editor/index.tsx
+++ b/compiler/forget/packages/playground/components/Editor/index.tsx
@@ -1,5 +1,8 @@
-/*
- * Copyright (c) Facebook, Inc. and its affiliates.
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
  */
 
 import { useDrag } from "@use-gesture/react";
@@ -89,7 +92,6 @@ export default function Editor() {
           })}
         >
           <Output store={deferredStore} updateDiagnostics={updateDiagnostics} />
-          <CompilerFlagsEditor />
         </div>
       </div>
       <button
diff --git a/compiler/forget/src/CompilerDriver.ts b/compiler/forget/src/CompilerDriver.ts
index 324aabb86a..2f83f05660 100644
--- a/compiler/forget/src/CompilerDriver.ts
+++ b/compiler/forget/src/CompilerDriver.ts
@@ -10,7 +10,6 @@ import * as t from "@babel/types";
 import * as BE from "./BackEnd";
 import { CompilerContext } from "./CompilerContext";
 import { CompilerOptions } from "./CompilerOptions";
-import DumpHIRPass from "./HIR/DumpHIRPass";
 import * as ME from "./MiddleEnd";
 import { PassManager } from "./PassManager";
 import * as Validation from "./Validation";
@@ -38,10 +37,6 @@ export function createCompilerDriver(
     compile() {
       const passManager = new PassManager(program, context);
 
-      // New architecture pass to print HIR, unrelated to the
-      // other passes here and just here for the playground.
-      passManager.addPass(DumpHIRPass);
-
       // Syntax Analysis and IR Generation.
       passManager.addPass(ME.ReactFuncsInfer);
       passManager.addPass(ME.ParamAnalysis);
diff --git a/compiler/forget/src/CompilerOutputs.ts b/compiler/forget/src/CompilerOutputs.ts
index 9accbf52b4..45a7e93c6b 100644
--- a/compiler/forget/src/CompilerOutputs.ts
+++ b/compiler/forget/src/CompilerOutputs.ts
@@ -16,7 +16,6 @@ import { SCCGraphSnapshot, ValGraphSnapshot } from "./DepGraph";
  */
 export enum OutputKind {
   IR = "DumpIR",
-  HIR = "DumpHIR",
   CFG = "DumpCFG",
   ValGraph = "DumpValGraph",
   SCCGraph = "DumpSCCGraph",
@@ -34,7 +33,6 @@ export function isSyntacticallyValidJS(kind: OutputKind) {
 
 export interface CompilerOutputs {
   [OutputKind.IR]: string;
-  [OutputKind.HIR]: string;
   [OutputKind.CFG]: string;
   [OutputKind.ValGraph]: ValGraphSnapshot[];
   [OutputKind.SCCGraph]: SCCGraphSnapshot[];
@@ -46,7 +44,6 @@ export interface CompilerOutputs {
 export function createCompilerOutputs(): CompilerOutputs {
   return {
     [OutputKind.IR]: "",
-    [OutputKind.HIR]: "",
     [OutputKind.CFG]: "",
     [OutputKind.ValGraph]: [],
     [OutputKind.SCCGraph]: [],
@@ -67,7 +64,6 @@ export function stringifyCompilerOutputs(
 ): PrettyCompilerOutputs {
   const outputs: PrettyCompilerOutputs = {
     [OutputKind.IR]: undefined,
-    [OutputKind.HIR]: undefined,
     [OutputKind.CFG]: undefined,
     [OutputKind.ValGraph]: undefined,
     [OutputKind.SCCGraph]: undefined,
@@ -95,10 +91,7 @@ export function stringifyCompilerOutputs(
             console.warn("Error formatting output:", e);
           }
         }
-      } else if (
-        outputKind === OutputKind.CFG ||
-        outputKind === OutputKind.HIR
-      ) {
+      } else if (outputKind === OutputKind.CFG) {
         outputs[outputKind] = output;
       } else {
         invariant(Array.isArray(output), "Graph outputs are arrays.");
diff --git a/compiler/forget/src/HIR/DumpHIRPass.ts b/compiler/forget/src/HIR/DumpHIRPass.ts
deleted file mode 100644
index 8628d02c61..0000000000
--- a/compiler/forget/src/HIR/DumpHIRPass.ts
+++ /dev/null
@@ -1,56 +0,0 @@
-/**
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- */
-
-import type { NodePath } from "@babel/traverse";
-import * as t from "@babel/types";
-import type { CompilerContext } from "../CompilerContext";
-import { OutputKind } from "../CompilerOutputs";
-import { PassKind, PassName } from "../Pass";
-
-import { lower } from "../HIR/BuildHIR";
-import { eliminateRedundantPhi } from "../HIR/EliminateRedundantPhi";
-import enterSSA from "../HIR/EnterSSA";
-import { inferMutableRanges } from "../HIR/InferMutableLifetimes";
-import inferReferenceEffects from "../HIR/InferReferenceEffects";
-import printHIR from "../HIR/PrintHIR";
-import { Environment } from "./HIRBuilder";
-import leaveSSA from "./LeaveSSA";
-
-export default {
-  name: PassName.DumpHIR,
-  kind: PassKind.Prog as const,
-  run,
-};
-
-/**
- * A step in the first approach compiler pipeline that outputs
- * the new architecture's HIR for the playground.
- */
-export function run(program: NodePath<t.Program>, context: CompilerContext) {
-  if (!context.opts.outputKinds.includes(OutputKind.HIR)) {
-    return;
-  }
-  const results: string[] = [];
-  program.traverse({
-    Function(func) {
-      try {
-        const env = new Environment();
-        const ir = lower(func, env);
-        enterSSA(ir, env);
-        eliminateRedundantPhi(ir);
-        inferReferenceEffects(ir);
-        inferMutableRanges(ir);
-        leaveSSA(ir);
-        const textHIR = printHIR(ir.body);
-        results.push(textHIR);
-      } catch (e) {
-        results.push(`Error: ${e}`);
-      }
-    },
-  });
-  context.outputs[OutputKind.HIR] = results.join("\n\n");
-}
diff --git a/compiler/forget/src/index.ts b/compiler/forget/src/index.ts
index e9ca894cba..31e8f3c971 100644
--- a/compiler/forget/src/index.ts
+++ b/compiler/forget/src/index.ts
@@ -21,4 +21,49 @@ export * from "./Diagnostic";
 export * from "./Logger";
 export { NoUseBeforeDefineRule } from "./Validation";
 
+import { lower } from "./HIR/BuildHIR";
+import { eliminateRedundantPhi } from "./HIR/EliminateRedundantPhi";
+import enterSSA from "./HIR/EnterSSA";
+import { inferMutableRanges } from "./HIR/InferMutableLifetimes";
+import inferReferenceEffects from "./HIR/InferReferenceEffects";
+import printHIR from "./HIR/PrintHIR";
+import { Environment } from "./HIR/HIRBuilder";
+import leaveSSA from "./HIR/LeaveSSA";
+import traverse, { NodePath } from "@babel/traverse";
+import * as t from "@babel/types";
+import { parse } from "@babel/parser";
+
+function parseFunctions(
+  source: string
+): Array<NodePath<t.FunctionDeclaration>> {
+  try {
+    const ast = parse(source, {
+      plugins: ["typescript", "jsx"],
+    });
+    const items: Array<NodePath<t.FunctionDeclaration>> = [];
+    traverse(ast, {
+      FunctionDeclaration: {
+        enter(nodePath) {
+          items.push(nodePath);
+        },
+      },
+    });
+    return items;
+  } catch (e) {
+    return [];
+  }
+}
+
+export const HIR = {
+  parseFunctions,
+  lower,
+  eliminateRedundantPhi,
+  enterSSA,
+  inferMutableRanges,
+  inferReferenceEffects,
+  printHIR,
+  Environment,
+  leaveSSA,
+};
+
 export default BabelPlugin;