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..182f32e61f2b6 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -102,6 +102,7 @@ import {lowerContextAccess} from '../Optimization/LowerContextAccess'; import {validateNoSetStateInPassiveEffects} from '../Validation/ValidateNoSetStateInPassiveEffects'; import {validateNoJSXInTryStatement} from '../Validation/ValidateNoJSXInTryStatement'; import {propagateScopeDependenciesHIR} from '../HIR/PropagateScopeDependenciesHIR'; +import {outlineJSX} from '../Optimization/OutlineJsx'; export type CompilerPipelineValue = | {kind: 'ast'; name: string; value: CodegenFunction} @@ -277,6 +278,10 @@ function* runWithEnvironment( value: hir, }); + if (env.config.enableJsxOutlining) { + outlineJSX(hir); + } + if (env.config.enableFunctionOutlining) { outlineFunctions(hir, fbtOperands); yield log({kind: 'hir', name: 'OutlineFunctions', value: hir}); diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts index c2c7d8d640846..6c37ec94c3de0 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts @@ -199,7 +199,7 @@ function insertNewOutlinedFunctionNode( program: NodePath, originalFn: BabelFn, compiledFn: CodegenFunction, -): NodePath { +): BabelFn { switch (originalFn.type) { case 'FunctionDeclaration': { return originalFn.insertAfter( @@ -492,18 +492,11 @@ export function compileProgram( fn.skip(); ALREADY_COMPILED.add(fn.node); if (outlined.type !== null) { - CompilerError.throwTodo({ - reason: `Implement support for outlining React functions (components/hooks)`, - loc: outlined.fn.loc, + queue.push({ + kind: 'outlined', + fn, + fnType: outlined.type, }); - /* - * Above should be as simple as the following, but needs testing: - * queue.push({ - * kind: "outlined", - * fn, - * fnType: outlined.type, - * }); - */ } } compiledFns.push({ 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..012e5b1149e8e 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -338,6 +338,11 @@ const EnvironmentConfigSchema = z.object({ */ enableFunctionOutlining: z.boolean().default(true), + /** + * TODO(gsn): Fill this out + */ + enableJsxOutlining: z.boolean().default(false), + /* * Enables instrumentation codegen. This emits a dev-mode only call to an * instrumentation function, for components and hooks that Forget compiles. diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts index 930dd79f2fd59..e663000390251 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts @@ -921,15 +921,7 @@ export type InstructionValue = type: Type; loc: SourceLocation; } - | { - kind: 'JsxExpression'; - tag: Place | BuiltinTag; - props: Array; - children: Array | null; // null === no children - loc: SourceLocation; - openingLoc: SourceLocation; - closingLoc: SourceLocation; - } + | JsxExpression | { kind: 'ObjectExpression'; properties: Array; @@ -1075,6 +1067,16 @@ export type InstructionValue = loc: SourceLocation; }; +export type JsxExpression = { + kind: 'JsxExpression'; + tag: Place | BuiltinTag; + props: Array; + children: Array | null; // null === no children + loc: SourceLocation; + openingLoc: SourceLocation; + closingLoc: SourceLocation; +}; + export type JsxAttribute = | {kind: 'JsxSpreadAttribute'; argument: Place} | {kind: 'JsxAttribute'; name: string; place: Place}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts new file mode 100644 index 0000000000000..84a4ccf938024 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts @@ -0,0 +1,359 @@ +/** + * 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 invariant from 'invariant'; +import {Environment} from '../HIR'; +import { + BasicBlock, + GeneratedSource, + HIRFunction, + IdentifierId, + Instruction, + InstructionKind, + JsxAttribute, + JsxExpression, + LoadGlobal, + makeBlockId, + makeIdentifierName, + makeInstructionId, + makeType, + ObjectProperty, + Place, + promoteTemporary, + promoteTemporaryJsxTag, +} from '../HIR/HIR'; +import {createTemporaryPlace} from '../HIR/HIRBuilder'; +import {deadCodeElimination} from './DeadCodeElimination'; + +export function outlineJSX(fn: HIRFunction): void { + const outlinedFns: Array = []; + outlineJsxImpl(fn, outlinedFns); + + for (const outlinedFn of outlinedFns) { + fn.env.outlineFunction(outlinedFn, 'Component'); + } +} + +type JsxInstruction = Instruction & {value: JsxExpression}; +type LoadGlobalInstruction = Instruction & {value: LoadGlobal}; +type LoadGlobalMap = Map; + +function outlineJsxImpl( + fn: HIRFunction, + outlinedFns: Array, +): void { + const globals: LoadGlobalMap = new Map(); + let shouldRunDeadCodeElimination = false; + + for (const [, block] of fn.body.blocks) { + const newInstrs = new Map(); + let jsx: Array = []; + + for (const instr of block.instructions) { + const {value, lvalue} = instr; + if (value.kind === 'LoadGlobal') { + globals.set(lvalue.identifier.id, instr as LoadGlobalInstruction); + continue; + } + + if (instr.value.kind === 'JsxExpression') { + jsx.push(instr as JsxInstruction); + continue; + } + + if (jsx.length !== 0) { + const result = process(fn, jsx, globals); + if (result) { + outlinedFns.push(result.fn); + newInstrs.set(jsx[0].id, { + end: instr.id, + instrs: result.instrs, + }); + } + + jsx = []; + } + + if (value.kind === 'FunctionExpression') { + outlineJsxImpl(value.loweredFunc.func, outlinedFns); + } + } + + if (jsx.length !== 0) { + const result = process(fn, jsx, globals); + if (result) { + outlinedFns.push(result.fn); + newInstrs.set(jsx[0].id, { + end: block.instructions.length, + instrs: result.instrs, + }); + } + } + + if (newInstrs.size > 0) { + shouldRunDeadCodeElimination = true; + const newInstr = []; + for (let i = 0; i < block.instructions.length; i++) { + if (newInstrs.has(i)) { + const {end, instrs} = newInstrs.get(i); + newInstr.push(...instrs); + i = end - 1; + } else { + newInstr.push(block.instructions[i]); + } + } + block.instructions = newInstr; + } + } + + if (shouldRunDeadCodeElimination) { + deadCodeElimination(fn); + } +} + +type OutlinedResult = { + instrs: Array; + fn: HIRFunction; +}; + +function process( + fn: HIRFunction, + jsx: Array, + globals: LoadGlobalMap, +): OutlinedResult | null { + /** + * In the future, add a check for backedge to outline jsx inside loops in a + * top level component. For now, only outline jsx in callbacks. + */ + if (fn.fnType === 'Component') { + return null; + } + + // Only outline nested jsx. + if (jsx.length < 2) { + return null; + } + + const props = collectProps(jsx); + if (!props) return null; + + const outlinedTag = fn.env.generateGloballyUniqueIdentifierName(null).value; + const newInstrs = emitOutlinedJsx(fn.env, jsx, props, outlinedTag); + if (!newInstrs) return null; + + const outlinedFn = emitOutlinedFn(fn.env, jsx, props, globals); + if (!outlinedFn) return null; + outlinedFn.id = outlinedTag; + + return {instrs: newInstrs, fn: outlinedFn}; +} + +function collectProps( + instructions: Array, +): Array | null { + const attributes: Array = []; + for (const instr of instructions) { + const {value} = instr; + + for (const at of value.props) { + if (at.kind === 'JsxSpreadAttribute') { + return null; + } + + if (at.kind === 'JsxAttribute') { + attributes.push(at); + } + } + } + return attributes; +} + +function emitOutlinedJsx( + env: Environment, + instructions: Array, + props: Array, + outlinedTag: string, +): Array { + const loadJsx: Instruction = { + id: makeInstructionId(0), + loc: GeneratedSource, + lvalue: createTemporaryPlace(env, GeneratedSource), + value: { + kind: 'LoadGlobal', + binding: { + kind: 'ModuleLocal', + name: outlinedTag, + }, + loc: GeneratedSource, + }, + }; + promoteTemporaryJsxTag(loadJsx.lvalue.identifier); + const jsxExpr: Instruction = { + id: makeInstructionId(0), + loc: GeneratedSource, + lvalue: instructions.at(-1)!.lvalue, + value: { + kind: 'JsxExpression', + tag: loadJsx.lvalue, + props, + children: null, + loc: GeneratedSource, + openingLoc: GeneratedSource, + closingLoc: GeneratedSource, + }, + }; + + return [loadJsx, jsxExpr]; +} + +function emitOutlinedFn( + env: Environment, + jsx: Array, + oldProps: Array, + globals: LoadGlobalMap, +): HIRFunction | null { + const instructions: Array = []; + const oldToNewProps = createOldToNewPropsMapping(env, oldProps); + + const propsObj: Place = createTemporaryPlace(env, GeneratedSource); + promoteTemporary(propsObj.identifier); + + const destructurePropsInstr = emitDestructureProps(env, propsObj, [ + ...oldToNewProps.values(), + ]); + instructions.push(destructurePropsInstr); + + updateProps(jsx, oldToNewProps); + const loadGlobalInstrs = emitLoadGlobals(jsx, globals); + if (!loadGlobalInstrs) { + return null; + } + instructions.push(...loadGlobalInstrs); + instructions.push(...jsx); + + const block: BasicBlock = { + kind: 'block', + id: makeBlockId(0), + instructions, + terminal: { + id: makeInstructionId(0), + kind: 'return', + loc: GeneratedSource, + value: instructions.at(-1)!.lvalue, + }, + preds: new Set(), + phis: new Set(), + }; + + const fn: HIRFunction = { + loc: GeneratedSource, + id: null, + fnType: 'Other', + env, + params: [propsObj], + returnTypeAnnotation: null, + returnType: makeType(), + context: [], + effects: null, + body: { + entry: block.id, + blocks: new Map([[block.id, block]]), + }, + generator: false, + async: false, + directives: [], + }; + return fn; +} + +function emitLoadGlobals( + jsx: Array, + globals: LoadGlobalMap, +): Array | null { + const instructions: Array = []; + for (const {value} of jsx) { + // Add load globals instructions for jsx tags + if (value.tag.kind === 'Identifier') { + const loadGlobalInstr = globals.get(value.tag.identifier.id); + if (!loadGlobalInstr) { + return null; + } + instructions.push(loadGlobalInstr); + } + } + + return instructions; +} + +function updateProps( + jsx: Array, + oldToNewProps: Map, +): void { + for (const {value} of jsx) { + // Update old props references to use the newly destructured props param + for (const prop of value.props) { + invariant( + prop.kind === 'JsxAttribute', + `Expected only attributes but found ${prop.kind}`, + ); + const newProp = oldToNewProps.get(prop.place.identifier.id); + invariant(newProp !== undefined, ''); + prop.place.identifier = newProp.place.identifier; + } + } +} +function createOldToNewPropsMapping( + env: Environment, + oldProps: Array, +): Map { + const oldToNewProps = new Map(); + + for (const oldProp of oldProps) { + invariant( + oldProp.kind === 'JsxAttribute', + `Expected only attributes but found ${oldProp.kind}`, + ); + const newProp: ObjectProperty = { + kind: 'ObjectProperty', + key: { + kind: 'string', + name: oldProp.name, + }, + type: 'property', + place: createTemporaryPlace(env, GeneratedSource), + }; + newProp.place.identifier.name = makeIdentifierName(oldProp.name); + oldToNewProps.set(oldProp.place.identifier.id, newProp); + } + + return oldToNewProps; +} + +function emitDestructureProps( + env: Environment, + propsObj: Place, + properties: Array, +): Instruction { + const destructurePropsInstr: Instruction = { + id: makeInstructionId(0), + lvalue: createTemporaryPlace(env, GeneratedSource), + loc: GeneratedSource, + value: { + kind: 'Destructure', + lvalue: { + pattern: { + kind: 'ObjectPattern', + properties, + }, + kind: InstructionKind.Let, + }, + loc: GeneratedSource, + value: propsObj, + }, + }; + return destructurePropsInstr; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-simple.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-simple.expect.md new file mode 100644 index 0000000000000..a1e688451b2a2 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-simple.expect.md @@ -0,0 +1,90 @@ + +## Input + +```javascript +// @enableJsxOutlining +function Component(arr) { + const x = useX(); + return arr.map(i => { + return ( + + + + ); + }); +} + +function useX() { + return 'x'; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [['foo', 'bar']], +}; + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @enableJsxOutlining +function Component(arr) { + const $ = _c(5); + const x = useX(); + let t0; + if ($[0] !== x || $[1] !== arr) { + let t1; + if ($[3] !== x) { + t1 = (i) => { + const T0 = _temp; + return ; + }; + $[3] = x; + $[4] = t1; + } else { + t1 = $[4]; + } + t0 = arr.map(t1); + $[0] = x; + $[1] = arr; + $[2] = t0; + } else { + t0 = $[2]; + } + return t0; +} +function _temp(t0) { + const $ = _c(5); + const { i: i$0, x: x$0 } = t0; + let t1; + if ($[0] !== i$0) { + t1 = ; + $[0] = i$0; + $[1] = t1; + } else { + t1 = $[1]; + } + let t2; + if ($[2] !== x$0 || $[3] !== t1) { + t2 = {t1}; + $[2] = x$0; + $[3] = t1; + $[4] = t2; + } else { + t2 = $[4]; + } + return t2; +} + +function useX() { + return "x"; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [["foo", "bar"]], +}; + +``` + \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-simple.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-simple.js new file mode 100644 index 0000000000000..445b2274eb8fc --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/jsx-outlining-simple.js @@ -0,0 +1,20 @@ +// @enableJsxOutlining +function Component(arr) { + const x = useX(); + return arr.map(i => { + return ( + + + + ); + }); +} + +function useX() { + return 'x'; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [['foo', 'bar']], +}; diff --git a/compiler/packages/snap/src/SproutTodoFilter.ts b/compiler/packages/snap/src/SproutTodoFilter.ts index 7d7476dc74739..98f574ddba4b5 100644 --- a/compiler/packages/snap/src/SproutTodoFilter.ts +++ b/compiler/packages/snap/src/SproutTodoFilter.ts @@ -504,6 +504,9 @@ const skipFilter = new Set([ 'todo.lower-context-access-array-destructuring', 'lower-context-selector-simple', 'lower-context-acess-multiple', + + // TODO + 'jsx-outlining-simple', ]); export default skipFilter;