[compiler] General-purpose function outlining

Implements general-purpose function outlining. Specifically, anonymous function expressions which have no dependencies/context variables are extracted into named top-level functions. The original function expression is replaced with a `LoadGlobal` of the generated name.

Note that the architecture is designed to allow very general purpose forms of outlining, though we currently are very conservative in what we outline. Specifically, the outlining allows annotating functions with an optional ReactiveFunctionType, which if set will cause the outlined function to get compiled as that type. So we could for example outline a helper hook or helper component, set the type, and then have the hook/component get memoized as well. For now though we just outline with no type set, and generate the function as-is without running it through compilation.

ghstack-source-id: 5b86845c4f32ed94678d3e7b0f1a912525b5e3fa
Pull Request resolved: #30331

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: [],
   });
 }