[compiler] Avoid bailouts when inserting gating (facebook#32598)

This change fixes a coverage hole in rolling out with `gating`. Prior to
this PR, configuring `gating` causes React Compiler to bail out of
optimizing some functions.

This means that it's not entirely safe to cutover from `gating` enabled
for all users (i.e. rolled out 100%) to removing the `gating` config
altogether, as new functions may be opted into compilation when they
stop bailing out due to gating-specific logic.

This is technically slightly slower due to the additional function
indirection. An alternative approach is to recommend running a codemod
to insert `use no memo`s on currently-bailing out functions before
removing the`gating` config.

---
Tested [internally](
https://fburl.com/diff/q982ovua) by enabling on a page that previously
had a few hundred bailouts due to gating + hoisted function declarations
and (1) clicking around locally and (2) running a bunch of e2e tests

Author: mofeiZ

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

@@ -10,6 +10,117 @@ import * as t from '@babel/types';
 import {PluginOptions} from './Options';
 import {CompilerError} from '../CompilerError';
 
+/**
+ * Gating rewrite for function declarations which are referenced before their
+ * declaration site.
+ *
+ * ```js
+ * // original
+ * export default React.memo(Foo);
+ * function Foo() { ... }
+ *
+ * // React compiler optimized + gated
+ * import {gating} from 'myGating';
+ * export default React.memo(Foo);
+ * const gating_result = gating();  <- inserted
+ * function Foo_optimized() {}      <- inserted
+ * function Foo_unoptimized() {}    <- renamed from Foo
+ * function Foo() {                 <- inserted function, which can be hoisted by JS engines
+ *   if (gating_result) return Foo_optimized();
+ *   else return Foo_unoptimized();
+ * }
+ * ```
+ */
+function insertAdditionalFunctionDeclaration(
+  fnPath: NodePath<t.FunctionDeclaration>,
+  compiled: t.FunctionDeclaration,
+  gating: NonNullable<PluginOptions['gating']>,
+): void {
+  const originalFnName = fnPath.node.id;
+  const originalFnParams = fnPath.node.params;
+  const compiledParams = fnPath.node.params;
+  /**
+   * Note that other than `export default function() {}`, all other function
+   * declarations must have a binding identifier. Since default exports cannot
+   * be referenced, it's safe to assume that all function declarations passed
+   * here will have an identifier.
+   * https://tc39.es/ecma262/multipage/ecmascript-language-functions-and-classes.html#sec-function-definitions
+   */
+  CompilerError.invariant(originalFnName != null && compiled.id != null, {
+    reason:
+      'Expected function declarations that are referenced elsewhere to have a named identifier',
+    loc: fnPath.node.loc ?? null,
+  });
+  CompilerError.invariant(originalFnParams.length === compiledParams.length, {
+    reason:
+      'Expected React Compiler optimized function declarations to have the same number of parameters as source',
+    loc: fnPath.node.loc ?? null,
+  });
+
+  const gatingCondition = fnPath.scope.generateUidIdentifier(
+    `${gating.importSpecifierName}_result`,
+  );
+  const unoptimizedFnName = fnPath.scope.generateUidIdentifier(
+    `${originalFnName.name}_unoptimized`,
+  );
+  const optimizedFnName = fnPath.scope.generateUidIdentifier(
+    `${originalFnName.name}_optimized`,
+  );
+  /**
+   * Step 1: rename existing functions
+   */
+  compiled.id.name = optimizedFnName.name;
+  fnPath.get('id').replaceInline(unoptimizedFnName);
+
+  /**
+   * Step 2: insert new function declaration
+   */
+  const newParams: Array<t.Identifier | t.RestElement> = [];
+  const genNewArgs: Array<() => t.Identifier | t.SpreadElement> = [];
+  for (let i = 0; i < originalFnParams.length; i++) {
+    const argName = `arg${i}`;
+    if (originalFnParams[i].type === 'RestElement') {
+      newParams.push(t.restElement(t.identifier(argName)));
+      genNewArgs.push(() => t.spreadElement(t.identifier(argName)));
+    } else {
+      newParams.push(t.identifier(argName));
+      genNewArgs.push(() => t.identifier(argName));
+    }
+  }
+  // insertAfter called in reverse order of how nodes should appear in program
+  fnPath.insertAfter(
+    t.functionDeclaration(
+      originalFnName,
+      newParams,
+      t.blockStatement([
+        t.ifStatement(
+          gatingCondition,
+          t.returnStatement(
+            t.callExpression(
+              compiled.id,
+              genNewArgs.map(fn => fn()),
+            ),
+          ),
+          t.returnStatement(
+            t.callExpression(
+              unoptimizedFnName,
+              genNewArgs.map(fn => fn()),
+            ),
+          ),
+        ),
+      ]),
+    ),
+  );
+  fnPath.insertBefore(
+    t.variableDeclaration('const', [
+      t.variableDeclarator(
+        gatingCondition,
+        t.callExpression(t.identifier(gating.importSpecifierName), []),
+      ),
+    ]),
+  );
+  fnPath.insertBefore(compiled);
+}
 export function insertGatedFunctionDeclaration(
   fnPath: NodePath<
     t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
@@ -21,15 +132,13 @@ export function insertGatedFunctionDeclaration(
   gating: NonNullable<PluginOptions['gating']>,
   referencedBeforeDeclaration: boolean,
 ): void {
-  if (referencedBeforeDeclaration) {
-    const identifier =
-      fnPath.node.type === 'FunctionDeclaration' ? fnPath.node.id : null;
-    CompilerError.invariant(false, {
-      reason: `Encountered a function used before its declaration, which breaks Forget's gating codegen due to hoisting`,
-      description: `Rewrite the reference to ${identifier?.name ?? 'this function'} to not rely on hoisting to fix this issue`,
-      loc: identifier?.loc ?? null,
-      suggestions: null,
+  if (referencedBeforeDeclaration && fnPath.isFunctionDeclaration()) {
+    CompilerError.invariant(compiled.type === 'FunctionDeclaration', {
+      reason: 'Expected compiled node type to match input type',
+      description: `Got ${compiled.type} but expected FunctionDeclaration`,
+      loc: fnPath.node.loc ?? null,
     });
+    insertAdditionalFunctionDeclaration(fnPath, compiled, gating);
   } else {
     const gatingExpression = t.conditionalExpression(
       t.callExpression(t.identifier(gating.importSpecifierName), []),

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/gating/component-syntax-ref-gating.flow.expect.md

@@ -0,0 +1,62 @@
+
+## Input
+
+```javascript
+// @flow @gating
+import {Stringify} from 'shared-runtime';
+import * as React from 'react';
+
+component Foo(ref: React.RefSetter<Controls>) {
+  return <Stringify ref={ref} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: eval('(...args) => React.createElement(Foo, args)'),
+  params: [{ref: React.createRef()}],
+};
+
+```
+
+## Code
+
+```javascript
+import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
+import { c as _c } from "react/compiler-runtime";
+import { Stringify } from "shared-runtime";
+import * as React from "react";
+
+const Foo = React.forwardRef(Foo_withRef);
+const _isForgetEnabled_Fixtures_result = isForgetEnabled_Fixtures();
+function _Foo_withRef_optimized(_$$empty_props_placeholder$$, ref) {
+  const $ = _c(2);
+  let t0;
+  if ($[0] !== ref) {
+    t0 = <Stringify ref={ref} />;
+    $[0] = ref;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+function _Foo_withRef_unoptimized(
+  _$$empty_props_placeholder$$: $ReadOnly<{}>,
+  ref: React.RefSetter<Controls>,
+): React.Node {
+  return <Stringify ref={ref} />;
+}
+function Foo_withRef(arg0, arg1) {
+  if (_isForgetEnabled_Fixtures_result)
+    return _Foo_withRef_optimized(arg0, arg1);
+  else return _Foo_withRef_unoptimized(arg0, arg1);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: eval("(...args) => React.createElement(Foo, args)"),
+  params: [{ ref: React.createRef() }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"ref":null}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/gating/component-syntax-ref-gating.flow.js

@@ -0,0 +1,12 @@
+// @flow @gating
+import {Stringify} from 'shared-runtime';
+import * as React from 'react';
+
+component Foo(ref: React.RefSetter<Controls>) {
+  return <Stringify ref={ref} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: eval('(...args) => React.createElement(Foo, args)'),
+  params: [{ref: React.createRef()}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/gating/error.component-syntax-ref-gating.flow.expect.md

@@ -0,0 +1,61 @@
+
+## Input
+
+```javascript
+// @gating
+import {createRef, forwardRef} from 'react';
+import {Stringify} from 'shared-runtime';
+
+const Foo = forwardRef(Foo_withRef);
+function Foo_withRef(props, ref) {
+  return <Stringify ref={ref} {...props} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: eval('(...args) => React.createElement(Foo, args)'),
+  params: [{prop1: 1, prop2: 2, ref: createRef()}],
+};
+
+```
+
+## Code
+
+```javascript
+import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
+import { c as _c } from "react/compiler-runtime"; // @gating
+import { createRef, forwardRef } from "react";
+import { Stringify } from "shared-runtime";
+
+const Foo = forwardRef(Foo_withRef);
+const _isForgetEnabled_Fixtures_result = isForgetEnabled_Fixtures();
+function _Foo_withRef_optimized(props, ref) {
+  const $ = _c(3);
+  let t0;
+  if ($[0] !== props || $[1] !== ref) {
+    t0 = <Stringify ref={ref} {...props} />;
+    $[0] = props;
+    $[1] = ref;
+    $[2] = t0;
+  } else {
+    t0 = $[2];
+  }
+  return t0;
+}
+function _Foo_withRef_unoptimized(props, ref) {
+  return <Stringify ref={ref} {...props} />;
+}
+function Foo_withRef(arg0, arg1) {
+  if (_isForgetEnabled_Fixtures_result)
+    return _Foo_withRef_optimized(arg0, arg1);
+  else return _Foo_withRef_unoptimized(arg0, arg1);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: eval("(...args) => React.createElement(Foo, args)"),
+  params: [{ prop1: 1, prop2: 2, ref: createRef() }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"0":{"prop1":1,"prop2":2,"ref":{"current":null}},"ref":"[[ cyclic ref *3 ]]"}</div>