[compiler] Add JSX inlining optimization (#30867)

This adds an `InlineJsxTransform` optimization pass, toggled by the
`enableInlineJsxTransform` flag. When enabled, JSX will be transformed
into React Element object literals, preventing runtime overhead during
element creation.

TODO:
- [ ] Add conditionals to make transform PROD-only
- [ ] Make the React element symbol configurable so this works with
runtimes that support `react.element` or `react.transitional.element`
- [ ] Look into additional optimization to pass props spread through
directly if none of the properties are mutated
This commit is contained in:
Jack Pope
2024-09-18 11:51:36 -04:00
committed by GitHub
parent 8dfbd16fce
commit 5dcb009760
8 changed files with 711 additions and 14 deletions
@@ -41,6 +41,7 @@ import {
constantPropagation,
deadCodeElimination,
pruneMaybeThrows,
inlineJsxTransform,
} from '../Optimization';
import {instructionReordering} from '../Optimization/InstructionReordering';
import {
@@ -351,6 +352,15 @@ function* runWithEnvironment(
});
}
if (env.config.enableInlineJsxTransform) {
inlineJsxTransform(hir);
yield log({
kind: 'hir',
name: 'inlineJsxTransform',
value: hir,
});
}
const reactiveFunction = buildReactiveFunction(hir);
yield log({
kind: 'reactive',
@@ -14,6 +14,7 @@ import {
ScopeId,
} from './HIR';
import {
fixScopeAndIdentifierRanges,
markInstructionIds,
markPredecessors,
reversePostorderBlocks,
@@ -176,20 +177,7 @@ export function buildReactiveScopeTerminalsHIR(fn: HIRFunction): void {
* Step 5:
* Fix scope and identifier ranges to account for renumbered instructions
*/
for (const [, block] of fn.body.blocks) {
const terminal = block.terminal;
if (terminal.kind === 'scope' || terminal.kind === 'pruned-scope') {
/*
* Scope ranges should always align to start at the 'scope' terminal
* and end at the first instruction of the fallthrough block
*/
const fallthroughBlock = fn.body.blocks.get(terminal.fallthrough)!;
const firstId =
fallthroughBlock.instructions[0]?.id ?? fallthroughBlock.terminal.id;
terminal.scope.range.start = terminal.id;
terminal.scope.range.end = firstId;
}
}
fixScopeAndIdentifierRanges(fn.body);
}
type TerminalRewriteInfo =
@@ -233,6 +233,13 @@ const EnvironmentConfigSchema = z.object({
*/
enableOptionalDependencies: z.boolean().default(true),
/**
* Enables inlining ReactElement object literals in place of JSX
* An alternative to the standard JSX transform which replaces JSX with React's jsxProd() runtime
* Currently a prod-only optimization, requiring Fast JSX dependencies
*/
enableInlineJsxTransform: z.boolean().default(false),
/*
* Enable validation of hooks to partially check that the component honors the rules of hooks.
* When disabled, the component is assumed to follow the rules (though the Babel plugin looks
@@ -912,3 +912,23 @@ export function clonePlaceToTemporary(env: Environment, place: Place): Place {
temp.reactive = place.reactive;
return temp;
}
/**
* Fix scope and identifier ranges to account for renumbered instructions
*/
export function fixScopeAndIdentifierRanges(func: HIR): void {
for (const [, block] of func.blocks) {
const terminal = block.terminal;
if (terminal.kind === 'scope' || terminal.kind === 'pruned-scope') {
/*
* Scope ranges should always align to start at the 'scope' terminal
* and end at the first instruction of the fallthrough block
*/
const fallthroughBlock = func.blocks.get(terminal.fallthrough)!;
const firstId =
fallthroughBlock.instructions[0]?.id ?? fallthroughBlock.terminal.id;
terminal.scope.range.start = terminal.id;
terminal.scope.range.end = firstId;
}
}
}
@@ -0,0 +1,402 @@
/**
* 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 {
BuiltinTag,
Effect,
HIRFunction,
Instruction,
JsxAttribute,
makeInstructionId,
ObjectProperty,
Place,
SpreadPattern,
} from '../HIR';
import {
createTemporaryPlace,
fixScopeAndIdentifierRanges,
markInstructionIds,
markPredecessors,
reversePostorderBlocks,
} from '../HIR/HIRBuilder';
function createSymbolProperty(
fn: HIRFunction,
instr: Instruction,
nextInstructions: Array<Instruction>,
propertyName: string,
symbolName: string,
): ObjectProperty {
const symbolPlace = createTemporaryPlace(fn.env, instr.value.loc);
const symbolInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...symbolPlace, effect: Effect.Mutate},
value: {
kind: 'LoadGlobal',
binding: {kind: 'Global', name: 'Symbol'},
loc: instr.value.loc,
},
loc: instr.loc,
};
nextInstructions.push(symbolInstruction);
const symbolForPlace = createTemporaryPlace(fn.env, instr.value.loc);
const symbolForInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...symbolForPlace, effect: Effect.Read},
value: {
kind: 'PropertyLoad',
object: {...symbolInstruction.lvalue},
property: 'for',
loc: instr.value.loc,
},
loc: instr.loc,
};
nextInstructions.push(symbolForInstruction);
const symbolValuePlace = createTemporaryPlace(fn.env, instr.value.loc);
const symbolValueInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...symbolValuePlace, effect: Effect.Mutate},
value: {
kind: 'Primitive',
value: symbolName,
loc: instr.value.loc,
},
loc: instr.loc,
};
nextInstructions.push(symbolValueInstruction);
const $$typeofPlace = createTemporaryPlace(fn.env, instr.value.loc);
const $$typeofInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...$$typeofPlace, effect: Effect.Mutate},
value: {
kind: 'MethodCall',
receiver: symbolInstruction.lvalue,
property: symbolForInstruction.lvalue,
args: [symbolValueInstruction.lvalue],
loc: instr.value.loc,
},
loc: instr.loc,
};
const $$typeofProperty: ObjectProperty = {
kind: 'ObjectProperty',
key: {name: propertyName, kind: 'string'},
type: 'property',
place: {...$$typeofPlace, effect: Effect.Capture},
};
nextInstructions.push($$typeofInstruction);
return $$typeofProperty;
}
function createTagProperty(
fn: HIRFunction,
instr: Instruction,
nextInstructions: Array<Instruction>,
componentTag: BuiltinTag | Place,
): ObjectProperty {
let tagProperty: ObjectProperty;
switch (componentTag.kind) {
case 'BuiltinTag': {
const tagPropertyPlace = createTemporaryPlace(fn.env, instr.value.loc);
const tagInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...tagPropertyPlace, effect: Effect.Mutate},
value: {
kind: 'Primitive',
value: componentTag.name,
loc: instr.value.loc,
},
loc: instr.loc,
};
tagProperty = {
kind: 'ObjectProperty',
key: {name: 'type', kind: 'string'},
type: 'property',
place: {...tagPropertyPlace, effect: Effect.Capture},
};
nextInstructions.push(tagInstruction);
break;
}
case 'Identifier': {
tagProperty = {
kind: 'ObjectProperty',
key: {name: 'type', kind: 'string'},
type: 'property',
place: {...componentTag, effect: Effect.Capture},
};
break;
}
}
return tagProperty;
}
function createPropsProperties(
fn: HIRFunction,
instr: Instruction,
nextInstructions: Array<Instruction>,
propAttributes: Array<JsxAttribute>,
children: Array<Place> | null,
): {
refProperty: ObjectProperty;
keyProperty: ObjectProperty;
propsProperty: ObjectProperty;
} {
let refProperty: ObjectProperty | undefined;
let keyProperty: ObjectProperty | undefined;
const props: Array<ObjectProperty | SpreadPattern> = [];
propAttributes.forEach(prop => {
switch (prop.kind) {
case 'JsxAttribute': {
if (prop.name === 'ref') {
refProperty = {
kind: 'ObjectProperty',
key: {name: 'ref', kind: 'string'},
type: 'property',
place: {...prop.place},
};
} else if (prop.name === 'key') {
keyProperty = {
kind: 'ObjectProperty',
key: {name: 'key', kind: 'string'},
type: 'property',
place: {...prop.place},
};
} else {
const attributeProperty: ObjectProperty = {
kind: 'ObjectProperty',
key: {name: prop.name, kind: 'string'},
type: 'property',
place: {...prop.place},
};
props.push(attributeProperty);
}
break;
}
case 'JsxSpreadAttribute': {
// TODO: Optimize spreads to pass object directly if none of its properties are mutated
props.push({
kind: 'Spread',
place: {...prop.argument},
});
break;
}
}
});
const propsPropertyPlace = createTemporaryPlace(fn.env, instr.value.loc);
if (children) {
let childrenPropProperty: ObjectProperty;
if (children.length === 1) {
childrenPropProperty = {
kind: 'ObjectProperty',
key: {name: 'children', kind: 'string'},
type: 'property',
place: {...children[0], effect: Effect.Capture},
};
} else {
const childrenPropPropertyPlace = createTemporaryPlace(
fn.env,
instr.value.loc,
);
const childrenPropInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...childrenPropPropertyPlace, effect: Effect.Mutate},
value: {
kind: 'ArrayExpression',
elements: [...children],
loc: instr.value.loc,
},
loc: instr.loc,
};
nextInstructions.push(childrenPropInstruction);
childrenPropProperty = {
kind: 'ObjectProperty',
key: {name: 'children', kind: 'string'},
type: 'property',
place: {...childrenPropPropertyPlace, effect: Effect.Capture},
};
}
props.push(childrenPropProperty);
}
if (refProperty == null) {
const refPropertyPlace = createTemporaryPlace(fn.env, instr.value.loc);
const refInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...refPropertyPlace, effect: Effect.Mutate},
value: {
kind: 'Primitive',
value: null,
loc: instr.value.loc,
},
loc: instr.loc,
};
refProperty = {
kind: 'ObjectProperty',
key: {name: 'ref', kind: 'string'},
type: 'property',
place: {...refPropertyPlace, effect: Effect.Capture},
};
nextInstructions.push(refInstruction);
}
if (keyProperty == null) {
const keyPropertyPlace = createTemporaryPlace(fn.env, instr.value.loc);
const keyInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...keyPropertyPlace, effect: Effect.Mutate},
value: {
kind: 'Primitive',
value: null,
loc: instr.value.loc,
},
loc: instr.loc,
};
keyProperty = {
kind: 'ObjectProperty',
key: {name: 'key', kind: 'string'},
type: 'property',
place: {...keyPropertyPlace, effect: Effect.Capture},
};
nextInstructions.push(keyInstruction);
}
const propsInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...propsPropertyPlace, effect: Effect.Mutate},
value: {
kind: 'ObjectExpression',
properties: props,
loc: instr.value.loc,
},
loc: instr.loc,
};
const propsProperty: ObjectProperty = {
kind: 'ObjectProperty',
key: {name: 'props', kind: 'string'},
type: 'property',
place: {...propsPropertyPlace, effect: Effect.Capture},
};
nextInstructions.push(propsInstruction);
return {refProperty, keyProperty, propsProperty};
}
// TODO: Make PROD only with conditional statements
export function inlineJsxTransform(fn: HIRFunction): void {
for (const [, block] of fn.body.blocks) {
let nextInstructions: Array<Instruction> | null = null;
for (let i = 0; i < block.instructions.length; i++) {
const instr = block.instructions[i]!;
switch (instr.value.kind) {
case 'JsxExpression': {
nextInstructions ??= block.instructions.slice(0, i);
const {refProperty, keyProperty, propsProperty} =
createPropsProperties(
fn,
instr,
nextInstructions,
instr.value.props,
instr.value.children,
);
const reactElementInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...instr.lvalue, effect: Effect.Store},
value: {
kind: 'ObjectExpression',
properties: [
createSymbolProperty(
fn,
instr,
nextInstructions,
'$$typeof',
/**
* TODO: Add this to config so we can switch between
* react.element / react.transitional.element
*/
'react.transitional.element',
),
createTagProperty(fn, instr, nextInstructions, instr.value.tag),
refProperty,
keyProperty,
propsProperty,
],
loc: instr.value.loc,
},
loc: instr.loc,
};
nextInstructions.push(reactElementInstruction);
break;
}
case 'JsxFragment': {
nextInstructions ??= block.instructions.slice(0, i);
const {refProperty, keyProperty, propsProperty} =
createPropsProperties(
fn,
instr,
nextInstructions,
[],
instr.value.children,
);
const reactElementInstruction: Instruction = {
id: makeInstructionId(0),
lvalue: {...instr.lvalue, effect: Effect.Store},
value: {
kind: 'ObjectExpression',
properties: [
createSymbolProperty(
fn,
instr,
nextInstructions,
'$$typeof',
/**
* TODO: Add this to config so we can switch between
* react.element / react.transitional.element
*/
'react.transitional.element',
),
createSymbolProperty(
fn,
instr,
nextInstructions,
'type',
'react.fragment',
),
refProperty,
keyProperty,
propsProperty,
],
loc: instr.value.loc,
},
loc: instr.loc,
};
nextInstructions.push(reactElementInstruction);
break;
}
default: {
if (nextInstructions !== null) {
nextInstructions.push(instr);
}
}
}
}
if (nextInstructions !== null) {
block.instructions = nextInstructions;
}
}
// Fixup the HIR to restore RPO, ensure correct predecessors, and renumber instructions.
reversePostorderBlocks(fn.body);
markPredecessors(fn.body);
markInstructionIds(fn.body);
// The renumbering instructions invalidates scope and identifier ranges
fixScopeAndIdentifierRanges(fn.body);
}
@@ -8,3 +8,4 @@
export {constantPropagation} from './ConstantPropagation';
export {deadCodeElimination} from './DeadCodeElimination';
export {pruneMaybeThrows} from './PruneMaybeThrows';
export {inlineJsxTransform} from './InlineJsxTransform';
@@ -0,0 +1,226 @@
## Input
```javascript
// @enableInlineJsxTransform
function Parent({children, a: _a, b: _b, c: _c, ref}) {
return <div ref={ref}>{children}</div>;
}
function Child({children}) {
return <>{children}</>;
}
function GrandChild({className}) {
return (
<span className={className}>
<React.Fragment key="fragmentKey">Hello world</React.Fragment>
</span>
);
}
function ParentAndRefAndKey(props) {
const testRef = useRef();
return <Parent a="a" b={{b: 'b'}} c={C} key="testKey" ref={testRef} />;
}
function ParentAndChildren(props) {
return (
<Parent>
<Child key="a" />
<Child key="b">
<GrandChild className={props.foo} />
</Child>
</Parent>
);
}
const propsToSpread = {a: 'a', b: 'b', c: 'c'};
function PropsSpread() {
return <Test {...propsToSpread} />;
}
export const FIXTURE_ENTRYPOINT = {
fn: ParentAndChildren,
params: [{foo: 'abc'}],
};
```
## Code
```javascript
import { c as _c2 } from "react/compiler-runtime"; // @enableInlineJsxTransform
function Parent(t0) {
const $ = _c2(2);
const { children, ref } = t0;
let t1;
if ($[0] !== children) {
t1 = {
$$typeof: Symbol.for("react.transitional.element"),
type: "div",
ref: ref,
key: null,
props: { children: children },
};
$[0] = children;
$[1] = t1;
} else {
t1 = $[1];
}
return t1;
}
function Child(t0) {
const $ = _c2(2);
const { children } = t0;
let t1;
if ($[0] !== children) {
t1 = {
$$typeof: Symbol.for("react.transitional.element"),
type: Symbol.for("react.fragment"),
ref: null,
key: null,
props: { children: children },
};
$[0] = children;
$[1] = t1;
} else {
t1 = $[1];
}
return t1;
}
function GrandChild(t0) {
const $ = _c2(3);
const { className } = t0;
let t1;
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
t1 = {
$$typeof: Symbol.for("react.transitional.element"),
type: React.Fragment,
ref: null,
key: "fragmentKey",
props: { children: "Hello world" },
};
$[0] = t1;
} else {
t1 = $[0];
}
let t2;
if ($[1] !== className) {
t2 = {
$$typeof: Symbol.for("react.transitional.element"),
type: "span",
ref: null,
key: null,
props: { className: className, children: t1 },
};
$[1] = className;
$[2] = t2;
} else {
t2 = $[2];
}
return t2;
}
function ParentAndRefAndKey(props) {
const $ = _c2(1);
const testRef = useRef();
let t0;
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
t0 = {
$$typeof: Symbol.for("react.transitional.element"),
type: Parent,
ref: testRef,
key: "testKey",
props: { a: "a", b: { b: "b" }, c: C },
};
$[0] = t0;
} else {
t0 = $[0];
}
return t0;
}
function ParentAndChildren(props) {
const $ = _c2(3);
let t0;
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
t0 = {
$$typeof: Symbol.for("react.transitional.element"),
type: Child,
ref: null,
key: "a",
props: {},
};
$[0] = t0;
} else {
t0 = $[0];
}
let t1;
if ($[1] !== props.foo) {
t1 = {
$$typeof: Symbol.for("react.transitional.element"),
type: Parent,
ref: null,
key: null,
props: {
children: [
t0,
{
$$typeof: Symbol.for("react.transitional.element"),
type: Child,
ref: null,
key: "b",
props: {
children: {
$$typeof: Symbol.for("react.transitional.element"),
type: GrandChild,
ref: null,
key: null,
props: { className: props.foo },
},
},
},
],
},
};
$[1] = props.foo;
$[2] = t1;
} else {
t1 = $[2];
}
return t1;
}
const propsToSpread = { a: "a", b: "b", c: "c" };
function PropsSpread() {
const $ = _c2(1);
let t0;
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
t0 = {
$$typeof: Symbol.for("react.transitional.element"),
type: Test,
ref: null,
key: null,
props: { ...propsToSpread },
};
$[0] = t0;
} else {
t0 = $[0];
}
return t0;
}
export const FIXTURE_ENTRYPOINT = {
fn: ParentAndChildren,
params: [{ foo: "abc" }],
};
```
### Eval output
(kind: ok) <div><span class="abc">Hello world</span></div>
@@ -0,0 +1,43 @@
// @enableInlineJsxTransform
function Parent({children, a: _a, b: _b, c: _c, ref}) {
return <div ref={ref}>{children}</div>;
}
function Child({children}) {
return <>{children}</>;
}
function GrandChild({className}) {
return (
<span className={className}>
<React.Fragment key="fragmentKey">Hello world</React.Fragment>
</span>
);
}
function ParentAndRefAndKey(props) {
const testRef = useRef();
return <Parent a="a" b={{b: 'b'}} c={C} key="testKey" ref={testRef} />;
}
function ParentAndChildren(props) {
return (
<Parent>
<Child key="a" />
<Child key="b">
<GrandChild className={props.foo} />
</Child>
</Parent>
);
}
const propsToSpread = {a: 'a', b: 'b', c: 'c'};
function PropsSpread() {
return <Test {...propsToSpread} />;
}
export const FIXTURE_ENTRYPOINT = {
fn: ParentAndChildren,
params: [{foo: 'abc'}],
};