From 8b69a6f74b1f4e123927a05f240410e957541ffb Mon Sep 17 00:00:00 2001 From: Jack Pope Date: Tue, 3 Sep 2024 12:08:23 -0400 Subject: [PATCH] Add inline jsx transform --- .../src/Entrypoint/Pipeline.ts | 10 + .../src/HIR/BuildReactiveScopeTerminalsHIR.ts | 16 +- .../src/HIR/Environment.ts | 7 + .../src/HIR/HIRBuilder.ts | 21 + .../src/Optimization/InlineJsxTransform.ts | 388 ++++++++++++++++++ .../src/Optimization/index.ts | 1 + .../compiler/inline-jsx-transform.expect.md | 227 ++++++++++ .../fixtures/compiler/inline-jsx-transform.js | 43 ++ 8 files changed, 699 insertions(+), 14 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts index aef18c90c2e5e..e29dbe2d3ac5e 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -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', diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts index 0999a3492b4a4..7c1fb54ea8058 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildReactiveScopeTerminalsHIR.ts @@ -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 = diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts index 75f3086011fd0..66270345fdf35 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -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 diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts index c694cf310fb39..14d5abcc3f486 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts @@ -17,6 +17,7 @@ import { GeneratedSource, GotoVariant, HIR, + HIRFunction, Identifier, IdentifierId, Instruction, @@ -912,3 +913,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; + } + } +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts new file mode 100644 index 0000000000000..84b3b3055327b --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/Optimization/InlineJsxTransform.ts @@ -0,0 +1,388 @@ +/** + * 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, + 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, + 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, + propAttributes: Array, + children: Array | null, +): { + refProperty: ObjectProperty; + keyProperty: ObjectProperty; + propsProperty: ObjectProperty; +} { + let refProperty: ObjectProperty | undefined; + let keyProperty: ObjectProperty | undefined; + const props: Array = []; + 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) { + 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); + const childrenPropProperty: ObjectProperty = { + 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 | 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); +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/index.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/index.ts index 722b05a809960..bb060b8dc285c 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/index.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Optimization/index.ts @@ -8,3 +8,4 @@ export {constantPropagation} from './ConstantPropagation'; export {deadCodeElimination} from './DeadCodeElimination'; export {pruneMaybeThrows} from './PruneMaybeThrows'; +export {inlineJsxTransform} from './InlineJsxTransform'; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md new file mode 100644 index 0000000000000..93eeda8ee598f --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.expect.md @@ -0,0 +1,227 @@ + +## Input + +```javascript +// @enableInlineJsxTransform + +function Parent({children, a: _a, b: _b, c: _c, ref}) { + return
{children}
+} + +function Child({children}) { + return <>{children} +} + +function GrandChild({ + className +}) { + return Hello world; +} + +function ParentAndRefAndKey(props) { + const testRef = useRef(); + return +} + +function ParentAndChildren(props) { + return ( + + + + + + + ) +} + +const propsToSpread = {a: 'a', b: 'b', c: 'c'}; +function PropsSpread() { + return +} + +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(4); + 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] === Symbol.for("react.memo_cache_sentinel")) { + t1 = { + $$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: "gc-1" }, + }, + ], + }, + }; + $[1] = t1; + } else { + t1 = $[1]; + } + let t2; + if ($[2] !== props.foo) { + t2 = { + $$typeof: Symbol.for("react.transitional.element"), + type: Parent, + ref: null, + key: null, + props: { foo: props.foo, children: [t0, t1] }, + }; + $[2] = props.foo; + $[3] = t2; + } else { + t2 = $[3]; + } + return t2; +} + +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)
Hello world
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js new file mode 100644 index 0000000000000..95a1b5cc7b8c5 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/inline-jsx-transform.js @@ -0,0 +1,43 @@ +// @enableInlineJsxTransform + +function Parent({children, a: _a, b: _b, c: _c, ref}) { + return
{children}
; +} + +function Child({children}) { + return <>{children}; +} + +function GrandChild({className}) { + return ( + + Hello world + + ); +} + +function ParentAndRefAndKey(props) { + const testRef = useRef(); + return ; +} + +function ParentAndChildren(props) { + return ( + + + + + + + ); +} + +const propsToSpread = {a: 'a', b: 'b', c: 'c'}; +function PropsSpread() { + return ; +} + +export const FIXTURE_ENTRYPOINT = { + fn: ParentAndChildren, + params: [{foo: 'abc'}], +};