[compiler] Avoid bailouts when inserting gating

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
This commit is contained in:
Mofei Zhang
2025-03-13 19:26:28 -04:00
parent 93b61fc4ec
commit 74b351dc8a
17 changed files with 526 additions and 96 deletions
@@ -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), []),
@@ -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>
@@ -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()}],
};
@@ -1,24 +0,0 @@
## Input
```javascript
// @flow @gating
component Foo(ref: React.RefSetter<Controls>) {
return <Bar ref={ref} />;
}
```
## Error
```
1 | // @flow @gating
> 2 | component Foo(ref: React.RefSetter<Controls>) {
| ^^^ Invariant: Encountered a function used before its declaration, which breaks Forget's gating codegen due to hoisting. Rewrite the reference to Foo_withRef to not rely on hoisting to fix this issue (2:2)
3 | return <Bar ref={ref} />;
4 | }
5 |
```
@@ -1,4 +0,0 @@
// @flow @gating
component Foo(ref: React.RefSetter<Controls>) {
return <Bar ref={ref} />;
}
@@ -1,26 +0,0 @@
## Input
```javascript
// @gating
const Foo = React.forwardRef(Foo_withRef);
function Foo_withRef(props, ref) {
return <Bar ref={ref} {...props}></Bar>;
}
```
## Error
```
1 | // @gating
2 | const Foo = React.forwardRef(Foo_withRef);
> 3 | function Foo_withRef(props, ref) {
| ^^^^^^^^^^^ Invariant: Encountered a function used before its declaration, which breaks Forget's gating codegen due to hoisting. Rewrite the reference to Foo_withRef to not rely on hoisting to fix this issue (3:3)
4 | return <Bar ref={ref} {...props}></Bar>;
5 | }
6 |
```
@@ -1,5 +0,0 @@
// @gating
const Foo = React.forwardRef(Foo_withRef);
function Foo_withRef(props, ref) {
return <Bar ref={ref} {...props}></Bar>;
}
@@ -1,24 +0,0 @@
## Input
```javascript
// @gating
import {memo} from 'react';
export default memo(Foo);
function Foo() {}
```
## Error
```
3 |
4 | export default memo(Foo);
> 5 | function Foo() {}
| ^^^ Invariant: Encountered a function used before its declaration, which breaks Forget's gating codegen due to hoisting. Rewrite the reference to Foo to not rely on hoisting to fix this issue (5:5)
6 |
```
@@ -1,5 +0,0 @@
// @gating
import {memo} from 'react';
export default memo(Foo);
function Foo() {}
@@ -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>
@@ -0,0 +1,13 @@
// @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()}],
};
@@ -0,0 +1,64 @@
## Input
```javascript
// @gating
import {memo} from 'react';
import {Stringify} from 'shared-runtime';
export default memo(Foo);
function Foo({prop1, prop2}) {
'use memo';
return <Stringify prop1={prop1} prop2={prop2} />;
}
export const FIXTURE_ENTRYPOINT = {
fn: eval('Foo'),
params: [{prop1: 1, prop2: 2}],
};
```
## Code
```javascript
import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
import { c as _c } from "react/compiler-runtime"; // @gating
import { memo } from "react";
import { Stringify } from "shared-runtime";
export default memo(Foo);
const _isForgetEnabled_Fixtures_result = isForgetEnabled_Fixtures();
function _Foo_optimized(t0) {
"use memo";
const $ = _c(3);
const { prop1, prop2 } = t0;
let t1;
if ($[0] !== prop1 || $[1] !== prop2) {
t1 = <Stringify prop1={prop1} prop2={prop2} />;
$[0] = prop1;
$[1] = prop2;
$[2] = t1;
} else {
t1 = $[2];
}
return t1;
}
function _Foo_unoptimized({ prop1, prop2 }) {
"use memo";
return <Stringify prop1={prop1} prop2={prop2} />;
}
function Foo(arg0) {
if (_isForgetEnabled_Fixtures_result) return _Foo_optimized(arg0);
else return _Foo_unoptimized(arg0);
}
export const FIXTURE_ENTRYPOINT = {
fn: eval("Foo"),
params: [{ prop1: 1, prop2: 2 }],
};
```
### Eval output
(kind: ok) <div>{"prop1":1,"prop2":2}</div>
@@ -0,0 +1,14 @@
// @gating
import {memo} from 'react';
import {Stringify} from 'shared-runtime';
export default memo(Foo);
function Foo({prop1, prop2}) {
'use memo';
return <Stringify prop1={prop1} prop2={prop2} />;
}
export const FIXTURE_ENTRYPOINT = {
fn: eval('Foo'),
params: [{prop1: 1, prop2: 2}],
};
@@ -0,0 +1,59 @@
## Input
```javascript
// @gating
import * as React from 'react';
let Foo;
const MemoFoo = React.memo(Foo);
Foo = () => <div>hello world!</div>;
/**
* Evaluate this fixture module to assert that compiler + original have the same
* runtime error message.
*/
export const FIXTURE_ENTRYPOINT = {
fn: () => {},
params: [],
};
```
## Code
```javascript
import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
import { c as _c } from "react/compiler-runtime"; // @gating
import * as React from "react";
let Foo;
const MemoFoo = React.memo(Foo);
Foo = isForgetEnabled_Fixtures()
? () => {
const $ = _c(1);
let t0;
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
t0 = <div>hello world!</div>;
$[0] = t0;
} else {
t0 = $[0];
}
return t0;
}
: () => <div>hello world!</div>;
/**
* Evaluate this fixture module to assert that compiler + original have the same
* runtime error message.
*/
export const FIXTURE_ENTRYPOINT = {
fn: isForgetEnabled_Fixtures() ? () => {} : () => {},
params: [],
};
```
### Eval output
(kind: ok)
logs: ['memo: The first argument must be a component. Instead received: %s','undefined']
@@ -0,0 +1,15 @@
// @gating
import * as React from 'react';
let Foo;
const MemoFoo = React.memo(Foo);
Foo = () => <div>hello world!</div>;
/**
* Evaluate this fixture module to assert that compiler + original have the same
* runtime error message.
*/
export const FIXTURE_ENTRYPOINT = {
fn: () => {},
params: [],
};
@@ -0,0 +1,86 @@
## Input
```javascript
// @gating
import * as React from 'react';
/**
* Test that the correct `Foo` is printed
*/
let Foo = () => <div>hello world 1!</div>;
const MemoOne = React.memo(Foo);
Foo = () => <div>hello world 2!</div>;
const MemoTwo = React.memo(Foo);
export const FIXTURE_ENTRYPOINT = {
fn: () => {
'use no memo';
return (
<>
<MemoOne />
<MemoTwo />
</>
);
},
params: [],
};
```
## Code
```javascript
import { isForgetEnabled_Fixtures } from "ReactForgetFeatureFlag";
import { c as _c } from "react/compiler-runtime"; // @gating
import * as React from "react";
/**
* Test that the correct `Foo` is printed
*/
let Foo = isForgetEnabled_Fixtures()
? () => {
const $ = _c(1);
let t0;
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
t0 = <div>hello world 1!</div>;
$[0] = t0;
} else {
t0 = $[0];
}
return t0;
}
: () => <div>hello world 1!</div>;
const MemoOne = React.memo(Foo);
Foo = isForgetEnabled_Fixtures()
? () => {
const $ = _c(1);
let t0;
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
t0 = <div>hello world 2!</div>;
$[0] = t0;
} else {
t0 = $[0];
}
return t0;
}
: () => <div>hello world 2!</div>;
const MemoTwo = React.memo(Foo);
export const FIXTURE_ENTRYPOINT = {
fn: () => {
"use no memo";
return (
<>
<MemoOne />
<MemoTwo />
</>
);
},
params: [],
};
```
### Eval output
(kind: ok) <div>hello world 1!</div><div>hello world 2!</div>
@@ -0,0 +1,23 @@
// @gating
import * as React from 'react';
/**
* Test that the correct `Foo` is printed
*/
let Foo = () => <div>hello world 1!</div>;
const MemoOne = React.memo(Foo);
Foo = () => <div>hello world 2!</div>;
const MemoTwo = React.memo(Foo);
export const FIXTURE_ENTRYPOINT = {
fn: () => {
'use no memo';
return (
<>
<MemoOne />
<MemoTwo />
</>
);
},
params: [],
};