[config] Clean up codegen imports and gating

- Gating checks for debugging / profiling can be moved to within the 
instrumentation or makeReadOnly functions. This simplifies codegen and reduces 
code bloat. 

- Warn if importing conflicting identifiers 

- Aggregate imports from the same source

compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Instrumentation.ts

@@ -11,18 +11,13 @@ import * as t from "@babel/types";
 export function addInstrumentForget(
   fn: NodePath<t.FunctionDeclaration>,
   fnName: string,
-  gatingIdentifierName: string,
   instrumentFnName: string
 ): void {
   const fnBody = fn.get("body");
   // Technically, this is a conditional hook call. However, we expect
   // __DEV__ and gatingIdentifier to be runtime constants
   const testExpr: t.Node = t.ifStatement(
-    t.logicalExpression(
-      "&&",
-      t.identifier("__DEV__"),
-      t.identifier(gatingIdentifierName)
-    ),
+    t.identifier("__DEV__"),
     t.expressionStatement(
       t.callExpression(t.identifier(instrumentFnName), [
         t.stringLiteral(fnName),

compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Options.ts

@@ -55,33 +55,26 @@ export type PluginOptions = {
    */
   gating: ExternalFunction | null;
   /**
-   * Enables instrumentation codegen. This emits a dev-mode only + gated call to
-   * an instrumentation function, for components and hooks that Forget compiles.
+   * Enables instrumentation codegen. This emits a dev-mode only call to an
+   * instrumentation function, for components and hooks that Forget compiles.
    * For example:
    *  instrumentForget: {
-   *    gating: {
-   *      source: 'ReactInstrumentForgetFeatureFlag',
-   *      importSpecifierName: 'isInstrumentForgetEnabled_Pokes',
-   *    },
-   *    instrumentFn: {
-   *      source: 'react-forget-runtime',
-   *      importSpecifierName: 'useRenderCounter',
-   *    }
+   *    source: 'react-forget-runtime',
+   *    importSpecifierName: 'useRenderCounter',
    *  }
    *
    * produces:
-   *  import {isInstrumentForgetEnabled_Pokes} from 'ReactInstrumentForgetFeatureFlag';
-   *  import {useRenderCounter} from 'react-forget-runtime';
+   *  import {useRenderCounter} from 'react-forget-runtime-pokes';
    *
    *  function Component(props) {
-   *    if (__DEV__ && isInstrumentForgetEnabled_Pokes) {
+   *    if (__DEV__) {
    *       useRenderCounter();
    *    }
    *    // ...
    *  }
    *
    */
-  instrumentForget: InstrumentForgetOptions | null;
+  instrumentForget: ExternalFunction | null;
 
   panicOnBailout: boolean;
 

compiler/forget/packages/babel-plugin-react-forget/src/Entrypoint/Program.ts

@@ -16,6 +16,7 @@ import { GeneratedSource } from "../HIR";
 import { addInstrumentForget } from "./Instrumentation";
 import { ExternalFunction, PluginOptions, parsePluginOptions } from "./Options";
 import { compileFn } from "./Pipeline";
+import { getOrInsertDefault } from "../Utils/utils";
 
 export type CompilerPass = {
   opts: PluginOptions;
@@ -81,36 +82,17 @@ export function compileProgram(
           })
         );
         if (pass.opts.instrumentForget != null) {
-          const gatingIdentifierName =
-            pass.opts.instrumentForget.gating.importSpecifierName;
           const instrumentFnName =
-            pass.opts.instrumentForget.instrumentFn.importSpecifierName;
-          addInstrumentForget(
-            fn,
-            originalIdent.name,
-            gatingIdentifierName,
-            instrumentFnName
-          );
-          addInstrumentForget(
-            compiledFn,
-            originalIdent.name,
-            gatingIdentifierName,
-            instrumentFnName
-          );
+            pass.opts.instrumentForget.importSpecifierName;
+          addInstrumentForget(fn, originalIdent.name, instrumentFnName);
+          addInstrumentForget(compiledFn, originalIdent.name, instrumentFnName);
         }
       } else {
         fn.replaceWith(compiled);
         if (pass.opts.instrumentForget != null) {
-          const gatingIdentifierName =
-            pass.opts.instrumentForget.gating.importSpecifierName;
           const instrumentFnName =
-            pass.opts.instrumentForget.instrumentFn.importSpecifierName;
-          addInstrumentForget(
-            fn,
-            originalIdent.name,
-            gatingIdentifierName,
-            instrumentFnName
-          );
+            pass.opts.instrumentForget.importSpecifierName;
+          addInstrumentForget(fn, originalIdent.name, instrumentFnName);
         }
       }
 
@@ -315,29 +297,18 @@ export function compileProgram(
         );
       }
     }
+    const externalFunctions = [];
     // TODO: check for duplicate import specifiers
     if (options.gating != null) {
-      program.unshiftContainer(
-        "body",
-        buildImportForExternalFunction(options.gating)
-      );
+      externalFunctions.push(options.gating);
     }
     if (options.instrumentForget != null) {
-      program.unshiftContainer(
-        "body",
-        buildImportForExternalFunction(options.instrumentForget.gating)
-      );
-      program.unshiftContainer(
-        "body",
-        buildImportForExternalFunction(options.instrumentForget.instrumentFn)
-      );
+      externalFunctions.push(options.instrumentForget);
     }
     if (options.environment?.enableEmitFreeze != null) {
-      program.unshiftContainer(
-        "body",
-        buildImportForExternalFunction(options.environment?.enableEmitFreeze)
-      );
+      externalFunctions.push(options.environment.enableEmitFreeze);
     }
+    addImportsToProgram(program, externalFunctions);
   }
 }
 
@@ -453,6 +424,49 @@ function buildBlockStatement(
   return body.node;
 }
 
+function addImportsToProgram(
+  path: NodePath<t.Program>,
+  importList: Array<ExternalFunction>
+): void {
+  const identifiers: Set<string> = new Set();
+  const sortedImports: Map<string, Array<string>> = new Map();
+  for (const { importSpecifierName, source } of importList) {
+    // Codegen currently does not rename import specifiers, so we do additional
+    // validation here
+    if (identifiers.has(importSpecifierName)) {
+      CompilerError.invalidInput(
+        `[InvalidConfig] Encountered conflicting import specifier for ${importSpecifierName} in Forget config.`,
+        GeneratedSource
+      );
+    }
+    if (path.scope.hasBinding(importSpecifierName)) {
+      CompilerError.invalidInput(
+        `[InvalidConfig] Encountered conflicting import specifiers for ${importSpecifierName} in generated program.`,
+        GeneratedSource
+      );
+    }
+    identifiers.add(importSpecifierName);
+
+    const importSpecifierNameList = getOrInsertDefault(
+      sortedImports,
+      source,
+      []
+    );
+    importSpecifierNameList.push(importSpecifierName);
+  }
+
+  const stmts: Array<t.ImportDeclaration> = [];
+  for (const [source, importSpecifierNameList] of sortedImports) {
+    const importSpecifiers = importSpecifierNameList.map((name) => {
+      const id = t.identifier(name);
+      return t.importSpecifier(id, id);
+    });
+
+    stmts.push(t.importDeclaration(importSpecifiers, t.stringLiteral(source)));
+  }
+  path.unshiftContainer("body", stmts);
+}
+
 type GatingTestOptions = {
   originalFnDecl: NodePath<t.FunctionDeclaration>;
   compiledIdent: t.Identifier;
@@ -469,7 +483,7 @@ function buildGatingTest({
     t.variableDeclarator(
       originalIdent,
       t.conditionalExpression(
-        t.callExpression(buildSpecifierIdent(gating), []),
+        t.callExpression(t.identifier(gating.importSpecifierName), []),
         compiledIdent,
         originalFnDecl.node.id!
       )
@@ -500,20 +514,6 @@ function addSuffix(id: t.Identifier, suffix: string): t.Identifier {
   return t.identifier(`${id.name}${suffix}`);
 }
 
-function buildImportForExternalFunction(
-  gating: ExternalFunction
-): t.ImportDeclaration {
-  const specifierIdent = buildSpecifierIdent(gating);
-  return t.importDeclaration(
-    [t.importSpecifier(specifierIdent, specifierIdent)],
-    t.stringLiteral(gating.source)
-  );
-}
-
-function buildSpecifierIdent(gating: ExternalFunction): t.Identifier {
-  return t.identifier(gating.importSpecifierName);
-}
-
 /**
  * Matches `import { ... } from 'react';`
  * but not `import * as React from 'react';`

compiler/forget/packages/babel-plugin-react-forget/src/Utils/utils.ts

@@ -44,3 +44,16 @@ export function retainWhere<T>(
   }
   array.length = writeIndex;
 }
+
+export function getOrInsertDefault<U, V>(
+  m: Map<U, V>,
+  key: U,
+  defaultValue: V
+): V {
+  if (m.has(key)) {
+    return m.get(key) as V;
+  } else {
+    m.set(key, defaultValue);
+    return defaultValue;
+  }
+}

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.expect.md

@@ -0,0 +1,35 @@
+
+## Input
+
+```javascript
+// @enableEmitFreeze @instrumentForget
+
+function useFoo(props) {
+  return foo(props.x);
+}
+
+```
+
+## Code
+
+```javascript
+import { useRenderCounter, makeReadOnly } from "react-forget-runtime";
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableEmitFreeze @instrumentForget
+
+function useFoo(props) {
+  if (__DEV__) useRenderCounter("useFoo");
+  const $ = useMemoCache(2);
+  const c_0 = $[0] !== props.x;
+  let t0;
+  if (c_0) {
+    t0 = foo(props.x);
+    $[0] = props.x;
+    $[1] = __DEV__ ? makeReadOnly(t0, "useFoo") : t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+```
+      

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-emit-imports-same-source.js

@@ -0,0 +1,5 @@
+// @enableEmitFreeze @instrumentForget
+
+function useFoo(props) {
+  return foo(props.x);
+}

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-emit-make-read-only.expect.md

@@ -19,7 +19,7 @@ function MyComponentName(props) {
 ## Code
 
 ```javascript
-import { makeReadOnly } from "react-forget-runtime-emit-freeze";
+import { makeReadOnly } from "react-forget-runtime";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @enableEmitFreeze true
 
 function MyComponentName(props) {

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-gating-test.expect.md

@@ -23,18 +23,17 @@ function Foo(props) {
 ## Code
 
 ```javascript
-import { useRenderCounter } from "react-forget-runtime";
-import { isInstrumentForgetEnabled_Fixtures } from "ReactInstrumentForgetFeatureFlag";
 import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
+import { useRenderCounter } from "react-forget-runtime";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @instrumentForget @forgetDirective @gating
 
 function Bar_uncompiled(props) {
   "use forget";
-  if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Bar");
+  if (__DEV__) useRenderCounter("Bar");
   return <div>{props.bar}</div>;
 }
 function Bar_forget(props) {
-  if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Bar");
+  if (__DEV__) useRenderCounter("Bar");
   const $ = useMemoCache(2);
   const c_0 = $[0] !== props.bar;
   let t0;
@@ -55,11 +54,11 @@ function NoForget(props) {
 
 function Foo_uncompiled(props) {
   "use forget";
-  if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Foo");
+  if (__DEV__) useRenderCounter("Foo");
   return <Foo>{props.bar}</Foo>;
 }
 function Foo_forget(props) {
-  if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Foo");
+  if (__DEV__) useRenderCounter("Foo");
   const $ = useMemoCache(2);
   const c_0 = $[0] !== props.bar;
   let t0;

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/codegen-instrument-forget-test.expect.md

@@ -24,11 +24,10 @@ function Foo(props) {
 
 ```javascript
 import { useRenderCounter } from "react-forget-runtime";
-import { isInstrumentForgetEnabled_Fixtures } from "ReactInstrumentForgetFeatureFlag";
 import { unstable_useMemoCache as useMemoCache } from "react"; // @instrumentForget @forgetDirective
 
 function Bar(props) {
-  if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Bar");
+  if (__DEV__) useRenderCounter("Bar");
   const $ = useMemoCache(2);
   const c_0 = $[0] !== props.bar;
   let t0;
@@ -47,7 +46,7 @@ function NoForget(props) {
 }
 
 function Foo(props) {
-  if (__DEV__ && isInstrumentForgetEnabled_Fixtures) useRenderCounter("Foo");
+  if (__DEV__) useRenderCounter("Foo");
   const $ = useMemoCache(2);
   const c_0 = $[0] !== props.bar;
   let t0;

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.codegen-error-on-conflicting-imports.expect.md

@@ -0,0 +1,21 @@
+
+## Input
+
+```javascript
+// @enableEmitFreeze @instrumentForget
+
+let makeReadOnly = "conflicting identifier";
+function useFoo(props) {
+  return foo(props.x);
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidInput: [InvalidConfig] Encountered conflicting import specifiers for makeReadOnly in generated program.
+```
+          
+      

compiler/forget/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.codegen-error-on-conflicting-imports.js

@@ -0,0 +1,6 @@
+// @enableEmitFreeze @instrumentForget
+
+let makeReadOnly = "conflicting identifier";
+function useFoo(props) {
+  return foo(props.x);
+}

compiler/forget/packages/snap/src/compiler-worker.ts

@@ -109,14 +109,8 @@ export async function compile(
     }
     if (firstLine.indexOf("@instrumentForget") !== -1) {
       instrumentForget = {
-        gating: {
-          source: "ReactInstrumentForgetFeatureFlag",
-          importSpecifierName: "isInstrumentForgetEnabled_Fixtures",
-        },
-        instrumentFn: {
-          source: "react-forget-runtime",
-          importSpecifierName: "useRenderCounter",
-        },
+        source: "react-forget-runtime",
+        importSpecifierName: "useRenderCounter",
       };
     }
     if (firstLine.indexOf("@panicOnBailout false") !== -1) {
@@ -139,7 +133,7 @@ export async function compile(
     }
     if (firstLine.indexOf("@enableEmitFreeze") !== -1) {
       enableEmitFreeze = {
-        source: "react-forget-runtime-emit-freeze",
+        source: "react-forget-runtime",
         importSpecifierName: "makeReadOnly",
       };
     }