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 6c31306274483..65ee683e45e4f 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -98,6 +98,7 @@ import { } from '../Validation'; import {validateLocalsNotReassignedAfterRender} from '../Validation/ValidateLocalsNotReassignedAfterRender'; import {outlineFunctions} from '../Optimization/OutlineFunctions'; +import {lowerContextAccess} from '../Optimization/LowerContextAccess'; export type CompilerPipelineValue = | {kind: 'ast'; name: string; value: CodegenFunction} @@ -199,6 +200,10 @@ function* runWithEnvironment( validateNoCapitalizedCalls(hir); } + if (env.config.enableLowerContextAccess) { + lowerContextAccess(hir); + } + analyseFunctions(hir); yield log({kind: 'hir', name: 'AnalyseFunctions', value: hir}); diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts new file mode 100644 index 0000000000000..11f2a7e63f6f6 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts @@ -0,0 +1,307 @@ +/** + * 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 { + ArrayExpression, + BasicBlock, + CallExpression, + Destructure, + Effect, + Environment, + GeneratedSource, + HIRFunction, + IdentifierId, + Instruction, + LoadLocal, + Place, + PropertyLoad, + ReturnTerminal, + isUseContextHookType, + makeBlockId, + makeIdentifierId, + makeIdentifierName, + makeInstructionId, + makeTemporary, + makeType, + markInstructionIds, + mergeConsecutiveBlocks, + removeUnnecessaryTryCatch, + reversePostorderBlocks, +} from '../HIR'; +import { + removeDeadDoWhileStatements, + removeUnreachableForUpdates, +} from '../HIR/HIRBuilder'; +import {enterSSA} from '../SSA'; +import {inferTypes} from '../TypeInference'; + +export function lowerContextAccess(fn: HIRFunction): void { + const contextAccess: Map = new Map(); + const contextKeys: Map> = new Map(); + + // collect context access and keys + for (const [, block] of fn.body.blocks) { + for (const instr of block.instructions) { + const {value, lvalue} = instr; + + if ( + value.kind === 'CallExpression' && + isUseContextHookType(value.callee.identifier) + ) { + contextAccess.set(lvalue.identifier.id, value); + continue; + } + + if (value.kind !== 'Destructure') { + continue; + } + + const destructureId = value.value.identifier.id; + if (!contextAccess.has(destructureId)) { + continue; + } + + const keys = getContextKeys(value); + if (keys === null) { + continue; + } + + if (contextKeys.has(destructureId)) { + /* + * TODO(gsn): Add support for accessing context over multiple + * statements. + */ + return; + } else { + contextKeys.set(destructureId, keys); + } + } + } + + if (contextAccess.size > 0) { + for (const [, block] of fn.body.blocks) { + const nextInstructions: Array = []; + for (const instr of block.instructions) { + const {lvalue, value} = instr; + if ( + value.kind === 'CallExpression' && + isUseContextHookType(value.callee.identifier) && + contextKeys.has(lvalue.identifier.id) + ) { + const keys = contextKeys.get(lvalue.identifier.id)!; + const selectorFnInstr = emitSelectorFn(fn.env, keys); + nextInstructions.push(selectorFnInstr); + + const selectorFn = selectorFnInstr.lvalue; + value.args.push(selectorFn); + } + + nextInstructions.push(instr); + } + block.instructions = nextInstructions; + } + markInstructionIds(fn.body); + } +} + +function getContextKeys(value: Destructure): Array | null { + const keys = []; + const pattern = value.lvalue.pattern; + + switch (pattern.kind) { + case 'ArrayPattern': { + for (const place of pattern.items) { + if (place.kind !== 'Identifier') { + return null; + } + + if (place.identifier.name === null) { + return null; + } + + keys.push(place.identifier.name.value); + } + return keys; + } + + case 'ObjectPattern': { + for (const place of pattern.properties) { + if ( + place.kind !== 'ObjectProperty' || + place.type !== 'property' || + place.key.kind !== 'identifier' + ) { + return null; + } + keys.push(place.key.name); + } + return keys; + } + } +} + +function emitPropertyLoad( + env: Environment, + obj: Place, + property: string, +): {instructions: Array; element: Place} { + const loadObj: LoadLocal = { + kind: 'LoadLocal', + place: obj, + loc: GeneratedSource, + }; + const object: Place = { + kind: 'Identifier', + identifier: makeTemporary(env.nextIdentifierId, GeneratedSource), + effect: Effect.Unknown, + reactive: false, + loc: GeneratedSource, + }; + const loadLocalInstr: Instruction = { + lvalue: object, + value: loadObj, + id: makeInstructionId(0), + loc: GeneratedSource, + }; + + const loadProp: PropertyLoad = { + kind: 'PropertyLoad', + object, + property, + loc: GeneratedSource, + }; + const element: Place = { + kind: 'Identifier', + identifier: makeTemporary(env.nextIdentifierId, GeneratedSource), + effect: Effect.Unknown, + reactive: false, + loc: GeneratedSource, + }; + const loadPropInstr: Instruction = { + lvalue: element, + value: loadProp, + id: makeInstructionId(0), + loc: GeneratedSource, + }; + return { + instructions: [loadLocalInstr, loadPropInstr], + element: element, + }; +} + +function emitSelectorFn(env: Environment, keys: Array): Instruction { + const obj: Place = { + kind: 'Identifier', + identifier: { + id: makeIdentifierId(env.nextIdentifierId), + name: makeIdentifierName('c'), + mutableRange: {start: makeInstructionId(0), end: makeInstructionId(0)}, + scope: null, + type: makeType(), + loc: GeneratedSource, + }, + effect: Effect.Unknown, + reactive: false, + loc: GeneratedSource, + }; + const instr: Array = []; + const elements = []; + for (const key of keys) { + const {instructions, element: prop} = emitPropertyLoad(env, obj, key); + instr.push(...instructions); + elements.push(prop); + } + + const arrayInstr = emitArrayInstr(elements, env); + instr.push(arrayInstr); + + const block: BasicBlock = { + kind: 'block', + id: makeBlockId(0), + instructions: instr, + terminal: { + id: makeInstructionId(0), + kind: 'return', + loc: GeneratedSource, + value: arrayInstr.lvalue, + }, + preds: new Set(), + phis: new Set(), + }; + + const fn: HIRFunction = { + loc: GeneratedSource, + id: null, + fnType: 'Other', + env, + params: [obj], + returnType: null, + context: [], + effects: null, + body: { + entry: block.id, + blocks: new Map([[block.id, block]]), + }, + generator: false, + async: false, + directives: [], + }; + + reversePostorderBlocks(fn.body); + removeUnreachableForUpdates(fn.body); + removeDeadDoWhileStatements(fn.body); + removeUnnecessaryTryCatch(fn.body); + markInstructionIds(fn.body); + mergeConsecutiveBlocks(fn); + enterSSA(fn); + inferTypes(fn); + + const fnInstr: Instruction = { + id: makeInstructionId(0), + value: { + kind: 'FunctionExpression', + name: null, + loweredFunc: { + func: fn, + dependencies: [], + }, + type: 'ArrowFunctionExpression', + loc: GeneratedSource, + }, + lvalue: { + kind: 'Identifier', + identifier: makeTemporary(env.nextIdentifierId, GeneratedSource), + effect: Effect.Unknown, + reactive: false, + loc: GeneratedSource, + }, + loc: GeneratedSource, + }; + return fnInstr; +} + +function emitArrayInstr(elements: Place[], env: Environment): Instruction { + const array: ArrayExpression = { + kind: 'ArrayExpression', + elements, + loc: GeneratedSource, + }; + const arrayLvalue: Place = { + kind: 'Identifier', + identifier: makeTemporary(env.nextIdentifierId, GeneratedSource), + effect: Effect.Unknown, + reactive: false, + loc: GeneratedSource, + }; + const arrayInstr: Instruction = { + id: makeInstructionId(0), + value: array, + lvalue: arrayLvalue, + loc: GeneratedSource, + }; + return arrayInstr; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.expect.md new file mode 100644 index 0000000000000..b7ee5e1902a39 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.expect.md @@ -0,0 +1,43 @@ + +## Input + +```javascript +// @enableLowerContextAccess +function App() { + const {foo} = useContext(MyContext); + const {bar} = useContext(MyContext); + return ; +} + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +function App() { + const $ = _c(3); + const { foo } = useContext(MyContext, _temp); + const { bar } = useContext(MyContext, _temp2); + let t0; + if ($[0] !== foo || $[1] !== bar) { + t0 = ; + $[0] = foo; + $[1] = bar; + $[2] = t0; + } else { + t0 = $[2]; + } + return t0; +} +function _temp2(c) { + return [c.bar]; +} +function _temp(c) { + return [c.foo]; +} + +``` + +### Eval output +(kind: exception) Fixture not implemented \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.js new file mode 100644 index 0000000000000..8d2894152213d --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.js @@ -0,0 +1,6 @@ +// @enableLowerContextAccess +function App() { + const {foo} = useContext(MyContext); + const {bar} = useContext(MyContext); + return ; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.expect.md new file mode 100644 index 0000000000000..17405b518d365 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.expect.md @@ -0,0 +1,38 @@ + +## Input + +```javascript +// @enableLowerContextAccess +function App() { + const {foo, bar} = useContext(MyContext); + return ; +} + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +function App() { + const $ = _c(3); + const { foo, bar } = useContext(MyContext, _temp); + let t0; + if ($[0] !== foo || $[1] !== bar) { + t0 = ; + $[0] = foo; + $[1] = bar; + $[2] = t0; + } else { + t0 = $[2]; + } + return t0; +} +function _temp(c) { + return [c.foo, c.bar]; +} + +``` + +### Eval output +(kind: exception) Fixture not implemented \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.js new file mode 100644 index 0000000000000..e24d653858752 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.js @@ -0,0 +1,5 @@ +// @enableLowerContextAccess +function App() { + const {foo, bar} = useContext(MyContext); + return ; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/no-lower-context-access.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/no-lower-context-access.expect.md new file mode 100644 index 0000000000000..534fbc28fb71d --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/no-lower-context-access.expect.md @@ -0,0 +1,39 @@ + +## Input + +```javascript +// @enableLowerContextAccess +function App() { + const context = useContext(MyContext); + const foo = context.foo; + const bar = context.bar; + return ; +} + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +function App() { + const $ = _c(3); + const context = useContext(MyContext); + const foo = context.foo; + const bar = context.bar; + let t0; + if ($[0] !== foo || $[1] !== bar) { + t0 = ; + $[0] = foo; + $[1] = bar; + $[2] = t0; + } else { + t0 = $[2]; + } + return t0; +} + +``` + +### Eval output +(kind: exception) Fixture not implemented \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/no-lower-context-access.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/no-lower-context-access.js new file mode 100644 index 0000000000000..9bc2bfadda871 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/no-lower-context-access.js @@ -0,0 +1,7 @@ +// @enableLowerContextAccess +function App() { + const context = useContext(MyContext); + const foo = context.foo; + const bar = context.bar; + return ; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.expect.md new file mode 100644 index 0000000000000..907c1d90a0a31 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.expect.md @@ -0,0 +1,39 @@ + +## Input + +```javascript +// @enableLowerContextAccess +function App() { + const context = useContext(MyContext); + const {foo} = context; + const {bar} = context; + return ; +} + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +function App() { + const $ = _c(3); + const context = useContext(MyContext); + const { foo } = context; + const { bar } = context; + let t0; + if ($[0] !== foo || $[1] !== bar) { + t0 = ; + $[0] = foo; + $[1] = bar; + $[2] = t0; + } else { + t0 = $[2]; + } + return t0; +} + +``` + +### Eval output +(kind: exception) Fixture not implemented \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.js new file mode 100644 index 0000000000000..18115b49b1f3b --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.js @@ -0,0 +1,7 @@ +// @enableLowerContextAccess +function App() { + const context = useContext(MyContext); + const {foo} = context; + const {bar} = context; + return ; +}