[compiler] General-purpose function outlining

compiler/packages/babel-plugin-react-compiler/scripts/jest/makeTransform.ts

@@ -28,6 +28,7 @@ import { basename } from "path";
 const e2eTransformerCacheKey = 1;
 const forgetOptions: EnvironmentConfig = validateEnvironmentConfig({
   enableAssumeHooksFollowRulesOfReact: true,
+  enableFunctionOutlining: false,
 });
 const debugMode = process.env["DEBUG_FORGET_COMPILER"] != null;
 

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -97,6 +97,7 @@ import {
   validateUseMemo,
 } from "../Validation";
 import { validateLocalsNotReassignedAfterRender } from "../Validation/ValidateLocalsNotReassignedAfterRender";
+import { outlineFunctions } from "../Optimization/OutlineFunctions";
 
 export type CompilerPipelineValue =
   | { kind: "ast"; name: string; value: CodegenFunction }
@@ -242,6 +243,11 @@ function* runWithEnvironment(
   inferReactiveScopeVariables(hir);
   yield log({ kind: "hir", name: "InferReactiveScopeVariables", value: hir });
 
+  if (env.config.enableFunctionOutlining) {
+    outlineFunctions(hir);
+    yield log({ kind: "hir", name: "OutlineFunctions", value: hir });
+  }
+
   alignMethodCallScopes(hir);
   yield log({
     kind: "hir",
@@ -480,6 +486,9 @@ function* runWithEnvironment(
 
   const ast = codegenFunction(reactiveFunction, uniqueIdentifiers).unwrap();
   yield log({ kind: "ast", name: "Codegen", value: ast });
+  for (const outlined of ast.outlined) {
+    yield log({ kind: "ast", name: "Codegen (outlined)", value: outlined.fn });
+  }
 
   /**
    * This flag should be only set for unit / fixture tests to check

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

@@ -87,6 +87,12 @@ export type BabelFn =
   | NodePath<t.ArrowFunctionExpression>;
 
 export type CompileResult = {
+  /**
+   * Distinguishes existing functions that were compiled ('original') from
+   * functions which were outlined. Only original functions need to be gated
+   * if gating mode is enabled.
+   */
+  kind: "original" | "outlined";
   originalFn: BabelFn;
   compiledFn: CodegenFunction;
 };
@@ -266,6 +272,7 @@ export function compileProgram(
   const lintError = suppressionsToCompilerError(suppressions);
   let hasCriticalError = lintError != null;
   const queue: Array<{
+    kind: "original" | "outlined";
     fn: BabelFn;
     fnType: ReactFunctionType;
   }> = [];
@@ -285,7 +292,7 @@ export function compileProgram(
     ALREADY_COMPILED.add(fn.node);
     fn.skip();
 
-    queue.push({ fn, fnType });
+    queue.push({ kind: "original", fn, fnType });
   };
 
   // Main traversal to compile with Forget
@@ -395,7 +402,33 @@ export function compileProgram(
     if (compiled === null) {
       continue;
     }
+    for (const outlined of compiled.outlined) {
+      CompilerError.invariant(outlined.fn.outlined.length === 0, {
+        reason: "Unexpected nested outlined functions",
+        loc: outlined.fn.loc,
+      });
+      const fn = current.fn.insertAfter(
+        createNewFunctionNode(current.fn, outlined.fn)
+      )[0]!;
+      fn.skip();
+      ALREADY_COMPILED.add(fn.node);
+      if (outlined.type !== null) {
+        CompilerError.throwTodo({
+          reason: `Implement support for outlining React functions (components/hooks)`,
+          loc: outlined.fn.loc,
+        });
+        /*
+         * Above should be as simple as the following, but needs testing:
+         * queue.push({
+         *   kind: "outlined",
+         *   fn,
+         *   fnType: outlined.type,
+         * });
+         */
+      }
+    }
     compiledFns.push({
+      kind: current.kind,
       compiledFn: compiled,
       originalFn: current.fn,
     });
@@ -466,10 +499,10 @@ export function compileProgram(
    * error elsewhere in the file, regardless of bailout mode.
    */
   for (const result of compiledFns) {
-    const { originalFn, compiledFn } = result;
+    const { kind, originalFn, compiledFn } = result;
     const transformedFn = createNewFunctionNode(originalFn, compiledFn);
 
-    if (gating != null) {
+    if (gating != null && kind === "original") {
       insertGatedFunctionDeclaration(originalFn, transformedFn, gating);
     } else {
       originalFn.replaceWith(transformedFn);

compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts

@@ -23,14 +23,17 @@ import {
   BuiltInType,
   Effect,
   FunctionType,
+  HIRFunction,
   IdentifierId,
   NonLocalBinding,
   PolyType,
   ScopeId,
   Type,
+  ValidatedIdentifier,
   ValueKind,
   makeBlockId,
   makeIdentifierId,
+  makeIdentifierName,
   makeScopeId,
 } from "./HIR";
 import {
@@ -284,6 +287,12 @@ const EnvironmentConfigSchema = z.object({
    */
   enableInstructionReordering: z.boolean().default(false),
 
+  /**
+   * Enables function outlinining, where anonymous functions that do not close over
+   * local variables can be extracted into top-level helper functions.
+   */
+  enableFunctionOutlining: z.boolean().default(true),
+
   /*
    * Enables instrumentation codegen. This emits a dev-mode only call to an
    * instrumentation function, for components and hooks that Forget compiles.
@@ -506,6 +515,10 @@ export class Environment {
   #nextBlock: number = 0;
   #nextScope: number = 0;
   #scope: BabelScope;
+  #outlinedFunctions: Array<{
+    fn: HIRFunction;
+    type: ReactFunctionType | null;
+  }> = [];
   logger: Logger | null;
   filename: string | null;
   code: string | null;
@@ -599,6 +612,24 @@ export class Environment {
     return this.#hoistedIdentifiers.has(node);
   }
 
+  generateGloballyUniqueIdentifierName(
+    name: string | null
+  ): ValidatedIdentifier {
+    const identifierNode = this.#scope.generateUidIdentifier(name ?? undefined);
+    return makeIdentifierName(identifierNode.name);
+  }
+
+  outlineFunction(fn: HIRFunction, type: ReactFunctionType | null): void {
+    this.#outlinedFunctions.push({ fn, type });
+  }
+
+  getOutlinedFunctions(): Array<{
+    fn: HIRFunction;
+    type: ReactFunctionType | null;
+  }> {
+    return this.#outlinedFunctions;
+  }
+
   getGlobalDeclaration(binding: NonLocalBinding): Global | null {
     if (this.config.hookPattern != null) {
       const match = new RegExp(this.config.hookPattern).exec(binding.name);

compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineFunctions.ts

@@ -0,0 +1,47 @@
+/**
+ * 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 { HIRFunction } from "../HIR";
+
+export function outlineFunctions(fn: HIRFunction): void {
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      const { value } = instr;
+
+      if (
+        value.kind === "FunctionExpression" ||
+        value.kind === "ObjectMethod"
+      ) {
+        // Recurse in case there are inner functions which can be outlined
+        outlineFunctions(value.loweredFunc.func);
+      }
+
+      if (
+        value.kind === "FunctionExpression" &&
+        value.loweredFunc.dependencies.length === 0 &&
+        value.loweredFunc.func.context.length === 0 &&
+        // TODO: handle outlining named functions
+        value.loweredFunc.func.id === null
+      ) {
+        const loweredFunc = value.loweredFunc.func;
+
+        const id = fn.env.generateGloballyUniqueIdentifierName(loweredFunc.id);
+        loweredFunc.id = id.value;
+
+        fn.env.outlineFunction(loweredFunc, null);
+        instr.value = {
+          kind: "LoadGlobal",
+          binding: {
+            kind: "Global",
+            name: id.value,
+          },
+          loc: value.loc,
+        };
+      }
+    }
+  }
+}

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -7,7 +7,12 @@
 
 import * as t from "@babel/types";
 import { createHmac } from "crypto";
-import { pruneHoistedContexts, pruneUnusedLValues, pruneUnusedLabels } from ".";
+import {
+  pruneHoistedContexts,
+  pruneUnusedLValues,
+  pruneUnusedLabels,
+  renameVariables,
+} from ".";
 import { CompilerError, ErrorSeverity } from "../CompilerError";
 import { Environment, EnvironmentConfig, ExternalFunction } from "../HIR";
 import {
@@ -36,6 +41,7 @@ import {
   ValidIdentifierName,
   getHookKind,
   makeIdentifierName,
+  promoteTemporary,
 } from "../HIR/HIR";
 import { printIdentifier, printPlace } from "../HIR/PrintHIR";
 import { eachPatternOperand } from "../HIR/visitors";
@@ -45,6 +51,8 @@ import { assertExhaustive } from "../Utils/utils";
 import { buildReactiveFunction } from "./BuildReactiveFunction";
 import { SINGLE_CHILD_FBT_TAGS } from "./MemoizeFbtAndMacroOperandsInSameScope";
 import { ReactiveFunctionVisitor, visitReactiveFunction } from "./visitors";
+import { ReactFunctionType } from "../HIR/Environment";
+import { logReactiveFunction } from "../Utils/logger";
 
 export const MEMO_CACHE_SENTINEL = "react.memo_cache_sentinel";
 export const EARLY_RETURN_SENTINEL = "react.early_return_sentinel";
@@ -85,6 +93,11 @@ export type CodegenFunction = {
    * because they were part of a pruned memo block.
    */
   prunedMemoValues: number;
+
+  outlined: Array<{
+    fn: CodegenFunction;
+    type: ReactFunctionType | null;
+  }>;
 };
 
 export function codegenFunction(
@@ -258,6 +271,40 @@ export function codegenFunction(
     compiled.body.body.unshift(test);
   }
 
+  const outlined: CodegenFunction["outlined"] = [];
+  for (const { fn: outlinedFunction, type } of cx.env.getOutlinedFunctions()) {
+    const reactiveFunction = buildReactiveFunction(outlinedFunction);
+    pruneUnusedLabels(reactiveFunction);
+    pruneUnusedLValues(reactiveFunction);
+    pruneHoistedContexts(reactiveFunction);
+
+    /*
+     * TODO: temporary function params (due to destructuring) should always be
+     * promoted so that they can be renamed
+     */
+    for (const param of reactiveFunction.params) {
+      const place = param.kind === "Identifier" ? param : param.place;
+      if (place.identifier.name === null) {
+        promoteTemporary(place.identifier);
+      }
+    }
+    const identifiers = renameVariables(reactiveFunction);
+    logReactiveFunction("Outline", reactiveFunction);
+    const codegen = codegenReactiveFunction(
+      new Context(
+        cx.env,
+        reactiveFunction.id ?? "[[ anonymous ]]",
+        identifiers
+      ),
+      reactiveFunction
+    );
+    if (codegen.isErr()) {
+      return codegen;
+    }
+    outlined.push({ fn: codegen.unwrap(), type });
+  }
+  compiled.outlined = outlined;
+
   return compileResult;
 }
 
@@ -306,6 +353,7 @@ function codegenReactiveFunction(
     memoValues: countMemoBlockVisitor.memoValues,
     prunedMemoBlocks: countMemoBlockVisitor.prunedMemoBlocks,
     prunedMemoValues: countMemoBlockVisitor.prunedMemoValues,
+    outlined: [],
   });
 }