[be] Extract gating logic to its own module

Just moving code around to another module.

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

@@ -0,0 +1,85 @@
+import { NodePath } from "@babel/core";
+import * as t from "@babel/types";
+import { CompilerError } from "../CompilerError";
+import { GeneratedSource } from "../HIR";
+import { ExternalFunction, PluginOptions } from "./Options";
+
+type GatingTestOptions = {
+  originalFnDecl: NodePath<t.FunctionDeclaration>;
+  compiledIdent: t.Identifier;
+  originalIdent: t.Identifier;
+  gating: ExternalFunction;
+};
+function buildGatingTest({
+  originalFnDecl,
+  compiledIdent,
+  originalIdent,
+  gating,
+}: GatingTestOptions): t.Node | t.Node[] {
+  const testVarDecl = t.variableDeclaration("const", [
+    t.variableDeclarator(
+      originalIdent,
+      t.conditionalExpression(
+        t.callExpression(t.identifier(gating.importSpecifierName), []),
+        compiledIdent,
+        originalFnDecl.node.id!
+      )
+    ),
+  ]);
+
+  // Re-export new declaration
+  const parent = originalFnDecl.parentPath;
+  if (t.isExportDefaultDeclaration(parent)) {
+    // Re-add uncompiled function
+    parent.replaceWith(originalFnDecl)[0].skip();
+
+    // Add test and synthesize new export
+    return [testVarDecl, t.exportDefaultDeclaration(originalIdent)];
+  } else if (t.isExportNamedDeclaration(parent)) {
+    // Re-add uncompiled function
+    parent.replaceWith(originalFnDecl)[0].skip();
+
+    // Add and export test
+    return t.exportNamedDeclaration(testVarDecl);
+  }
+
+  // Just add the test, no need for re-export
+  return testVarDecl;
+}
+
+function addSuffix(id: t.Identifier, suffix: string): t.Identifier {
+  return t.identifier(`${id.name}${suffix}`);
+}
+
+export function insertGatedFunctionDeclaration(
+  fnPath: NodePath<t.FunctionDeclaration>,
+  compiled: t.FunctionDeclaration,
+  originalIdent: t.Identifier,
+  gating: NonNullable<PluginOptions["gating"]>
+): NodePath<t.FunctionDeclaration> {
+  // Rename existing function
+  fnPath.node.id = addSuffix(originalIdent, "_uncompiled");
+
+  // Rename and append compiled function
+  CompilerError.invariant(compiled.id != null, {
+    reason: "FunctionDeclaration must produce a name",
+    description: null,
+    loc: fnPath.node.loc ?? GeneratedSource,
+    suggestions: null,
+  });
+  compiled.id = addSuffix(compiled.id, "_forget");
+  const compiledFn = fnPath.insertAfter(compiled)[0];
+  compiledFn.skip();
+
+  // Build and append gating test
+  compiledFn.insertAfter(
+    buildGatingTest({
+      originalFnDecl: fnPath,
+      compiledIdent: compiled.id,
+      originalIdent,
+      gating,
+    })
+  );
+
+  return compiledFn;
+}

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

@@ -16,6 +16,7 @@ import {
 import { GeneratedSource } from "../HIR";
 import { isComponentDeclaration } from "../Utils/ComponentDeclaration";
 import { getOrInsertDefault } from "../Utils/utils";
+import { insertGatedFunctionDeclaration } from "./Gating";
 import { addInstrumentForget } from "./Instrumentation";
 import { ExternalFunction, PluginOptions, parsePluginOptions } from "./Options";
 import { compileFn } from "./Pipeline";
@@ -115,39 +116,6 @@ function compileAndInsertNewFunctionDeclaration(
   return false;
 }
 
-function insertGatedFunctionDeclaration(
-  fnPath: NodePath<t.FunctionDeclaration>,
-  compiled: t.FunctionDeclaration,
-  originalIdent: t.Identifier,
-  gating: NonNullable<PluginOptions["gating"]>
-): NodePath<t.FunctionDeclaration> {
-  // Rename existing function
-  fnPath.node.id = addSuffix(originalIdent, "_uncompiled");
-
-  // Rename and append compiled function
-  CompilerError.invariant(compiled.id != null, {
-    reason: "FunctionDeclaration must produce a name",
-    description: null,
-    loc: fnPath.node.loc ?? GeneratedSource,
-    suggestions: null,
-  });
-  compiled.id = addSuffix(compiled.id, "_forget");
-  const compiledFn = fnPath.insertAfter(compiled)[0];
-  compiledFn.skip();
-
-  // Build and append gating test
-  compiledFn.insertAfter(
-    buildGatingTest({
-      originalFnDecl: fnPath,
-      compiledIdent: compiled.id,
-      originalIdent,
-      gating,
-    })
-  );
-
-  return compiledFn;
-}
-
 export function compileProgram(
   program: NodePath<t.Program>,
   pass: CompilerPass
@@ -508,53 +476,6 @@ function addImportsToProgram(
   path.unshiftContainer("body", stmts);
 }
 
-type GatingTestOptions = {
-  originalFnDecl: NodePath<t.FunctionDeclaration>;
-  compiledIdent: t.Identifier;
-  originalIdent: t.Identifier;
-  gating: ExternalFunction;
-};
-function buildGatingTest({
-  originalFnDecl,
-  compiledIdent,
-  originalIdent,
-  gating,
-}: GatingTestOptions): t.Node | t.Node[] {
-  const testVarDecl = t.variableDeclaration("const", [
-    t.variableDeclarator(
-      originalIdent,
-      t.conditionalExpression(
-        t.callExpression(t.identifier(gating.importSpecifierName), []),
-        compiledIdent,
-        originalFnDecl.node.id!
-      )
-    ),
-  ]);
-
-  // Re-export new declaration
-  const parent = originalFnDecl.parentPath;
-  if (t.isExportDefaultDeclaration(parent)) {
-    // Re-add uncompiled function
-    parent.replaceWith(originalFnDecl)[0].skip();
-
-    // Add test and synthesize new export
-    return [testVarDecl, t.exportDefaultDeclaration(originalIdent)];
-  } else if (t.isExportNamedDeclaration(parent)) {
-    // Re-add uncompiled function
-    parent.replaceWith(originalFnDecl)[0].skip();
-
-    // Add and export test
-    return t.exportNamedDeclaration(testVarDecl);
-  }
-
-  // Just add the test, no need for re-export
-  return testVarDecl;
-}
-
-function addSuffix(id: t.Identifier, suffix: string): t.Identifier {
-  return t.identifier(`${id.name}${suffix}`);
-}
-
 /**
  * Matches `import { ... } from 'react';`
  * but not `import * as React from 'react';`

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

@@ -1,4 +1,5 @@
-export * from "./Program";
+export * from "./Gating";
 export * from "./Instrumentation";
 export * from "./Options";
 export * from "./Pipeline";
+export * from "./Program";