[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: ea1ef4108d97298a1cc59b41ed64f988e808f3c7
Pull Request resolved: #30331

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,9 @@ function* runWithEnvironment(
   inferReactiveScopeVariables(hir);
   yield log({ kind: "hir", name: "InferReactiveScopeVariables", value: hir });
 
+  outlineFunctions(hir);
+  yield log({ kind: "hir", name: "OutlineFunctions", value: hir });
+
   alignMethodCallScopes(hir);
   yield log({
     kind: "hir",
@@ -480,6 +484,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

@@ -88,6 +88,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;
 };
@@ -267,6 +273,7 @@ export function compileProgram(
   const lintError = suppressionsToCompilerError(suppressions);
   let hasCriticalError = lintError != null;
   const queue: Array<{
+    kind: "original" | "outlined";
     fn: BabelFn;
     fnType: ReactFunctionType;
   }> = [];
@@ -286,7 +293,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
@@ -396,7 +403,26 @@ 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) {
+        queue.push({
+          kind: "outlined",
+          fn,
+          fnType: outlined.type,
+        });
+      }
+    }
     compiledFns.push({
+      kind: current.kind,
       compiledFn: compiled,
       originalFn: current.fn,
     });
@@ -467,10 +493,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,6 +23,9 @@ import {
   BuiltInType,
   Effect,
   FunctionType,
+  GeneratedSource,
+  HIRFunction,
+  Identifier,
   IdentifierId,
   NonLocalBinding,
   PolyType,
@@ -31,7 +34,10 @@ import {
   ValueKind,
   makeBlockId,
   makeIdentifierId,
+  makeIdentifierName,
+  makeInstructionId,
   makeScopeId,
+  makeType,
 } from "./HIR";
 import {
   BuiltInMixedReadonlyId,
@@ -42,6 +48,7 @@ import {
   addHook,
 } from "./ObjectShape";
 import { Scope as BabelScope } from "@babel/traverse";
+import { CodegenFunction } from "../ReactiveScopes";
 
 export const ExternalFunctionSchema = z.object({
   // Source for the imported module that exports the `importSpecifierName` functions
@@ -506,6 +513,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 +610,32 @@ export class Environment {
     return this.#hoistedIdentifiers.has(node);
   }
 
+  generateUniqueGlobalIdentifier(name: string | null): Identifier {
+    const id = this.nextIdentifierId;
+    const identifierNode = this.#scope.generateUidIdentifier(name ?? undefined);
+    return {
+      id,
+      loc: GeneratedSource,
+      mutableRange: { start: makeInstructionId(0), end: makeInstructionId(0) },
+      name: makeIdentifierName(identifierNode.name),
+      scope: null,
+      type: { kind: "Poly" },
+    };
+  }
+
+  outlineFunction(fn: HIRFunction, type: ReactFunctionType | null): void {
+    this.#outlinedFunctions.push({ fn, type });
+  }
+
+  takeOutlinedFunctions(): Array<{
+    fn: HIRFunction;
+    type: ReactFunctionType | null;
+  }> {
+    const outlined = this.#outlinedFunctions;
+    this.#outlinedFunctions = [];
+    return outlined;
+  }
+
   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,52 @@
+/**
+ * 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 { CompilerError } from "..";
+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.generateUniqueGlobalIdentifier(loweredFunc.id);
+        CompilerError.invariant(id.name !== null && id.name.kind === "named", {
+          reason: "Expected to generate a named variable",
+          loc: value.loc,
+        });
+        loweredFunc.id = id.name.value;
+
+        fn.env.outlineFunction(loweredFunc, null);
+        instr.value = {
+          kind: "LoadGlobal",
+          binding: {
+            kind: "Global",
+            name: id.name.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,38 @@ export function codegenFunction(
     compiled.body.body.unshift(test);
   }
 
+  const outlined: CodegenFunction["outlined"] = [];
+  for (const { fn: outlinedFunction, type } of cx.env.takeOutlinedFunctions()) {
+    const reactiveFunction = buildReactiveFunction(outlinedFunction);
+    pruneUnusedLabels(reactiveFunction);
+    pruneUnusedLValues(reactiveFunction);
+    pruneHoistedContexts(reactiveFunction);
+
+    // TODO: temporary function params (due to destructuring) should alwasy 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 +351,7 @@ function codegenReactiveFunction(
     memoValues: countMemoBlockVisitor.memoValues,
     prunedMemoBlocks: countMemoBlockVisitor.prunedMemoBlocks,
     prunedMemoValues: countMemoBlockVisitor.prunedMemoValues,
+    outlined: [],
   });
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-global-mutation-in-effect-indirect-usecallback.expect.md

@@ -40,57 +40,52 @@ import { useCallback, useEffect, useState } from "react";
 let someGlobal = {};
 
 function Component() {
-  const $ = _c(7);
+  const $ = _c(6);
   const [state, setState] = useState(someGlobal);
+
+  const setGlobal = _temp;
   let t0;
+  let t1;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     t0 = () => {
-      someGlobal.value = true;
-    };
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const setGlobal = t0;
-  let t1;
-  let t2;
-  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = () => {
       setGlobal();
     };
-    t2 = [];
+    t1 = [];
+    $[0] = t0;
     $[1] = t1;
-    $[2] = t2;
   } else {
+    t0 = $[0];
     t1 = $[1];
-    t2 = $[2];
   }
-  useEffect(t1, t2);
+  useEffect(t0, t1);
+  let t2;
   let t3;
-  let t4;
-  if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
-    t3 = () => {
+  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
+    t2 = () => {
       setState(someGlobal.value);
     };
-    t4 = [someGlobal];
+    t3 = [someGlobal];
+    $[2] = t2;
     $[3] = t3;
-    $[4] = t4;
   } else {
+    t2 = $[2];
     t3 = $[3];
-    t4 = $[4];
   }
-  useEffect(t3, t4);
+  useEffect(t2, t3);
 
-  const t5 = String(state);
-  let t6;
-  if ($[5] !== t5) {
-    t6 = <div>{t5}</div>;
+  const t4 = String(state);
+  let t5;
+  if ($[4] !== t4) {
+    t5 = <div>{t4}</div>;
+    $[4] = t4;
     $[5] = t5;
-    $[6] = t6;
   } else {
-    t6 = $[6];
+    t5 = $[5];
   }
-  return t6;
+  return t5;
+}
+function _temp() {
+  someGlobal.value = true;
 }
 
 export const FIXTURE_ENTRYPOINT = {