Pass environment options through babel plugin

Updates the babel plugin so that environment options — including custom hook 
definitions — can be passed in through the plugin: 

* Renames `CompilerFlags` => `PluginOptions` since they are specific to the 
babel plugin, and are no longer just flags. 

* Moves the definition of `useFreeze()` out of the builtin hook list and instead 
passes it when our unit tests configure the plugin.

compiler/forget/src/Babel/BabelPlugin.ts

@@ -12,11 +12,11 @@ import generate from "@babel/generator";
 import jsx from "@babel/plugin-syntax-jsx";
 import * as t from "@babel/types";
 import prettier from "prettier";
-import { CompilerFlags, parseCompilerFlags } from "../CompilerFlags";
 import { compile } from "../CompilerPipeline";
+import { parsePluginOptions, PluginOptions } from "./PluginOptions";
 
 type BabelPluginPass = {
-  opts: CompilerFlags;
+  opts: PluginOptions;
 };
 
 /**
@@ -47,7 +47,7 @@ export default function ReactForgetBabelPlugin(
       if (fn.scope.getProgramParent() !== fn.scope.parent) {
         return;
       }
-      const ast = compile(fn, null);
+      const ast = compile(fn, pass.opts.environment);
 
       // We are generating a new FunctionDeclaration node, so we must skip over it or this
       // traversal will loop infinitely.
@@ -73,7 +73,7 @@ export default function ReactForgetBabelPlugin(
       // prior to B, if A does not have a Program visitor and B does, B will run first. We always
       // want Forget to run true to source as possible.
       Program(path, pass) {
-        const flags = parseCompilerFlags(pass.opts);
+        const flags = parsePluginOptions(pass.opts);
         path.traverse(visitor, {
           ...pass,
           opts: { ...pass.opts, ...flags },

compiler/forget/src/Babel/PluginOptions.ts

@@ -0,0 +1,45 @@
+/**
+ * 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 { EnvironmentOptions } from "../HIR/Environment";
+
+export type PluginOptions = {
+  /**
+   * Enable to make Forget only compile functions containing the 'use forget' directive.
+   */
+  enableOnlyOnUseForgetDirective: boolean;
+
+  environment: EnvironmentOptions | null;
+};
+
+export const defaultOptions: PluginOptions = {
+  enableOnlyOnUseForgetDirective: false,
+  environment: null,
+} as const;
+
+export function parsePluginOptions(obj: unknown): PluginOptions {
+  if (obj == null || typeof obj !== "object") {
+    return defaultOptions;
+  }
+  const invalidOptions: Array<string> = [];
+  let parsedOptions: Partial<PluginOptions> = Object.create(null);
+  for (const [key, value] of Object.entries(obj)) {
+    if (isCompilerFlag(key)) {
+      parsedOptions[key] = value;
+    } else {
+      invalidOptions.push(key);
+    }
+  }
+  if (invalidOptions.length > 0) {
+    console.error(`Unexpected React Forget compiler flags: ${invalidOptions}`);
+  }
+  return { ...defaultOptions, ...parsedOptions };
+}
+
+function isCompilerFlag(s: string): s is keyof typeof defaultOptions {
+  return Object.prototype.hasOwnProperty.call(defaultOptions, s);
+}

compiler/forget/src/Babel/RunReactForgetBabelPlugin.ts

@@ -4,6 +4,7 @@ import * as parser from "@babel/parser";
 import invariant from "invariant";
 import prettier from "prettier";
 import ReactForgetBabelPlugin from "./BabelPlugin";
+import { PluginOptions } from "./PluginOptions";
 
 type ReactForgetBabelPluginResult = {
   ast: BabelCore.BabelFileResult["ast"];
@@ -14,7 +15,8 @@ type ReactForgetBabelPluginResult = {
 export default function runReactForgetBabelPlugin(
   text: string,
   file: string,
-  language: "flow" | "typescript"
+  language: "flow" | "typescript",
+  options: PluginOptions | null
 ): ReactForgetBabelPluginResult {
   const ast = parser.parse(text, {
     sourceFilename: file,
@@ -25,7 +27,7 @@ export default function runReactForgetBabelPlugin(
     filename: file,
     highlightCode: false,
     retainLines: true,
-    plugins: [ReactForgetBabelPlugin],
+    plugins: [[ReactForgetBabelPlugin, options]],
     sourceType: "module",
   });
   invariant(

compiler/forget/src/CompilerFlags.ts

@@ -1,40 +0,0 @@
-/**
- * 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.
- */
-
-export type CompilerFlags = {
-  /**
-   * Enable to make Forget only compile functions containing the 'use forget' directive.
-   */
-  enableOnlyOnUseForgetDirective: boolean;
-};
-
-export const defaultFlags: CompilerFlags = {
-  enableOnlyOnUseForgetDirective: false,
-} as const;
-
-export function parseCompilerFlags(obj: unknown): CompilerFlags {
-  if (obj == null || typeof obj !== "object") {
-    return defaultFlags;
-  }
-  const invalidFlags: Array<string> = [];
-  let parsedFlags: Partial<CompilerFlags> = Object.create(null);
-  for (const [key, value] of Object.entries(obj)) {
-    if (isCompilerFlag(key)) {
-      parsedFlags[key] = value;
-    } else {
-      invalidFlags.push(key);
-    }
-  }
-  if (invalidFlags.length > 0) {
-    console.error(`Unexpected React Forget compiler flags: ${invalidFlags}`);
-  }
-  return { ...defaultFlags, ...parsedFlags };
-}
-
-function isCompilerFlag(s: string): s is keyof typeof defaultFlags {
-  return Object.prototype.hasOwnProperty.call(defaultFlags, s);
-}

compiler/forget/src/HIR/Hooks.ts

@@ -26,15 +26,6 @@ export const BUILTIN_HOOKS: Map<string, Hook> = new Map([
       valueKind: ValueKind.Mutable,
     },
   ],
-  [
-    "useFreeze",
-    {
-      kind: "Ref",
-      name: "useFreeze",
-      effectKind: Effect.Freeze,
-      valueKind: ValueKind.Frozen,
-    },
-  ],
   [
     "useMemo",
     {

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

@@ -10,6 +10,7 @@
 import { wasmFolder } from "@hpcc-js/wasm";
 import path from "path";
 import runReactForgetBabelPlugin from "../Babel/RunReactForgetBabelPlugin";
+import { Effect, ValueKind } from "../HIR";
 import { toggleLogging } from "../Utils/logger";
 import generateTestsFromFixtures from "./test-utils/generateTestsFromFixtures";
 
@@ -60,7 +61,22 @@ describe("React Forget (HIR version)", () => {
       }
       try {
         items.push({
-          js: runReactForgetBabelPlugin(input, file, language).code,
+          js: runReactForgetBabelPlugin(input, file, language, {
+            enableOnlyOnUseForgetDirective: false,
+            environment: {
+              customHooks: new Map([
+                [
+                  "useFreeze",
+                  {
+                    name: "useFreeze",
+                    kind: "Custom",
+                    valueKind: ValueKind.Frozen,
+                    effectKind: Effect.Freeze,
+                  },
+                ],
+              ]),
+            },
+          }).code,
         });
       } catch (e) {
         error = e;