From 64589aede69c51f7d3d5484432f1e8666512cad9 Mon Sep 17 00:00:00 2001 From: Sathya Gunsasekaran Date: Tue, 6 Aug 2024 16:30:23 +0100 Subject: [PATCH] [compiler] Rewrite useContext callee If a value is specified for the LowerContextAccess environment config, we rewrite the callee from 'useContext' to the specificed value. This will allow us run an experiment internally. ghstack-source-id: 60120858972dacc54bfc55daaadd0924a10f2cba Pull Request resolved: https://github.com/facebook/react/pull/30612 --- .../src/Entrypoint/Pipeline.ts | 4 +- .../src/Entrypoint/Program.ts | 5 +++ .../src/HIR/Environment.ts | 11 +++-- .../src/Optimization/LowerContextAccess.ts | 45 +++++++++++++++++-- .../lower-context-acess-multiple.expect.md | 9 ++-- .../compiler/lower-context-acess-multiple.js | 2 +- .../lower-context-selector-simple.expect.md | 7 +-- .../compiler/lower-context-selector-simple.js | 2 +- ...ntext-access-array-destructuring.expect.md | 5 ++- ...ower-context-access-array-destructuring.js | 2 +- ...text-access-destructure-multiple.expect.md | 5 ++- ...wer-context-access-destructure-multiple.js | 2 +- ...r-context-access-mixed-array-obj.expect.md | 5 ++- ...do.lower-context-access-mixed-array-obj.js | 2 +- ...text-access-nested-destructuring.expect.md | 5 ++- ...wer-context-access-nested-destructuring.js | 2 +- ...wer-context-access-property-load.expect.md | 5 ++- ...todo.lower-context-access-property-load.js | 2 +- compiler/packages/snap/src/compiler.ts | 15 +++++++ 19 files changed, 101 insertions(+), 34 deletions(-) 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 d515de2900eb5..3e557ebe0f9f3 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -200,8 +200,8 @@ function* runWithEnvironment( validateNoCapitalizedCalls(hir); } - if (env.config.enableLowerContextAccess) { - lowerContextAccess(hir); + if (env.config.lowerContextAccess) { + lowerContextAccess(hir, env.config.lowerContextAccess); } analyseFunctions(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 049fe58e1b419..f976b9031f6da 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts @@ -511,6 +511,11 @@ export function compileProgram( externalFunctions.push(gating); } + const lowerContextAccess = pass.opts.environment?.lowerContextAccess; + if (lowerContextAccess) { + externalFunctions.push(tryParseExternalFunction(lowerContextAccess)); + } + const enableEmitInstrumentForget = pass.opts.environment?.enableEmitInstrumentForget; if (enableEmitInstrumentForget != null) { 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 dad27965afd19..4bbd8cb17f89a 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -436,8 +436,11 @@ const EnvironmentConfigSchema = z.object({ enableTreatRefLikeIdentifiersAsRefs: z.boolean().nullable().default(false), /* - * If enabled, this lowers any calls to `useContext` hook to use a selector - * function. + * If specified a value, the compiler lowers any calls to `useContext` to use + * this value as the callee. + * + * A selector function is compiled and passed as an argument along with the + * context to this function call. * * The compiler automatically figures out the keys by looking for the immediate * destructuring of the return value from the useContext call. In the future, @@ -449,10 +452,10 @@ const EnvironmentConfigSchema = z.object({ * const {foo, bar} = useContext(MyContext); * * // output - * const {foo, bar} = useContext(MyContext, (c) => [c.foo, c.bar]); + * const {foo, bar} = useCompiledContext(MyContext, (c) => [c.foo, c.bar]); * ``` */ - enableLowerContextAccess: z.boolean().nullable().default(false), + lowerContextAccess: ExternalFunctionSchema.nullish(), }); export type EnvironmentConfig = z.infer; diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts index 6a8b23ffa5c31..efa0c0cdaa231 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts @@ -10,12 +10,13 @@ import { BasicBlock, CallExpression, Destructure, - Effect, Environment, + ExternalFunction, GeneratedSource, HIRFunction, IdentifierId, Instruction, + LoadGlobal, LoadLocal, Place, PropertyLoad, @@ -30,7 +31,10 @@ import {createTemporaryPlace} from '../HIR/HIRBuilder'; import {enterSSA} from '../SSA'; import {inferTypes} from '../TypeInference'; -export function lowerContextAccess(fn: HIRFunction): void { +export function lowerContextAccess( + fn: HIRFunction, + loweredContextCallee: ExternalFunction, +): void { const contextAccess: Map = new Map(); const contextKeys: Map> = new Map(); @@ -85,13 +89,23 @@ export function lowerContextAccess(fn: HIRFunction): void { isUseContextHookType(value.callee.identifier) && contextKeys.has(lvalue.identifier.id) ) { - const keys = contextKeys.get(lvalue.identifier.id)!; - const selectorFnInstr = emitSelectorFn(fn.env, keys); + const loweredContextCalleeInstr = emitLoadLoweredContextCallee( + fn.env, + loweredContextCallee, + ); + if (nextInstructions === null) { nextInstructions = block.instructions.slice(0, i); } + nextInstructions.push(loweredContextCalleeInstr); + + const keys = contextKeys.get(lvalue.identifier.id)!; + const selectorFnInstr = emitSelectorFn(fn.env, keys); nextInstructions.push(selectorFnInstr); + const lowerContextCallId = loweredContextCalleeInstr.lvalue; + value.callee = lowerContextCallId; + const selectorFn = selectorFnInstr.lvalue; value.args.push(selectorFn); } @@ -105,9 +119,32 @@ export function lowerContextAccess(fn: HIRFunction): void { } } markInstructionIds(fn.body); + inferTypes(fn); } } +function emitLoadLoweredContextCallee( + env: Environment, + loweredContextCallee: ExternalFunction, +): Instruction { + const loadGlobal: LoadGlobal = { + kind: 'LoadGlobal', + binding: { + kind: 'ImportNamespace', + module: loweredContextCallee.source, + name: loweredContextCallee.importSpecifierName, + }, + loc: GeneratedSource, + }; + + return { + id: makeInstructionId(0), + loc: GeneratedSource, + lvalue: createTemporaryPlace(env, GeneratedSource), + value: loadGlobal, + }; +} + function getContextKeys(value: Destructure): Array | null { const keys = []; const pattern = value.lvalue.pattern; 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 index 0f9124536c55f..7b81ed33bb0ac 100644 --- 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 @@ -2,7 +2,7 @@ ## Input ```javascript -// @enableLowerContextAccess +// @lowerContextAccess function App() { const {foo} = useContext(MyContext); const {bar} = useContext(MyContext); @@ -14,11 +14,12 @@ function App() { ## Code ```javascript -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +import { useContext_withSelector } from "react-compiler-runtime"; +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess function App() { const $ = _c(3); - const { foo } = useContext(MyContext, _temp); - const { bar } = useContext(MyContext, _temp2); + const { foo } = useContext_withSelector(MyContext, _temp); + const { bar } = useContext_withSelector(MyContext, _temp2); let t0; if ($[0] !== foo || $[1] !== bar) { t0 = ; 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 index 8d2894152213d..314ff41c280e5 100644 --- 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 @@ -1,4 +1,4 @@ -// @enableLowerContextAccess +// @lowerContextAccess function App() { const {foo} = useContext(MyContext); const {bar} = useContext(MyContext); 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 index 8109cb5b81f02..0b12c2e250875 100644 --- 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 @@ -2,7 +2,7 @@ ## Input ```javascript -// @enableLowerContextAccess +// @lowerContextAccess function App() { const {foo, bar} = useContext(MyContext); return ; @@ -13,10 +13,11 @@ function App() { ## Code ```javascript -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +import { useContext_withSelector } from "react-compiler-runtime"; +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess function App() { const $ = _c(3); - const { foo, bar } = useContext(MyContext, _temp); + const { foo, bar } = useContext_withSelector(MyContext, _temp); let t0; if ($[0] !== foo || $[1] !== bar) { t0 = ; 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 index e24d653858752..0d53548e67c56 100644 --- 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 @@ -1,4 +1,4 @@ -// @enableLowerContextAccess +// @lowerContextAccess function App() { const {foo, bar} = useContext(MyContext); return ; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.expect.md index 0ebd7989a4e7a..7feba7d408cf4 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.expect.md @@ -2,7 +2,7 @@ ## Input ```javascript -// @enableLowerContextAccess +// @lowerContextAccess function App() { const [foo, bar] = useContext(MyContext); return ; @@ -13,7 +13,8 @@ function App() { ## Code ```javascript -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +import { useContext_withSelector } from "react-compiler-runtime"; +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess function App() { const $ = _c(3); const [foo, bar] = useContext(MyContext); diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.js index b1f60bf8e7d97..387c30cdc1e7e 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.js +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.js @@ -1,4 +1,4 @@ -// @enableLowerContextAccess +// @lowerContextAccess function App() { const [foo, bar] = useContext(MyContext); 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 index 21322ed17864c..ac38876bce3f6 100644 --- 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 @@ -2,7 +2,7 @@ ## Input ```javascript -// @enableLowerContextAccess +// @lowerContextAccess function App() { const context = useContext(MyContext); const {foo} = context; @@ -15,7 +15,8 @@ function App() { ## Code ```javascript -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +import { useContext_withSelector } from "react-compiler-runtime"; +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess function App() { const $ = _c(3); const context = useContext(MyContext); 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 index 18115b49b1f3b..e7b106e9db864 100644 --- 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 @@ -1,4 +1,4 @@ -// @enableLowerContextAccess +// @lowerContextAccess function App() { const context = useContext(MyContext); const {foo} = context; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.expect.md index 5217fbd6168e3..d808a5e3f04cd 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.expect.md @@ -2,7 +2,7 @@ ## Input ```javascript -// @enableLowerContextAccess +// @lowerContextAccess function App() { const context = useContext(MyContext); const [foo] = context; @@ -15,7 +15,8 @@ function App() { ## Code ```javascript -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +import { useContext_withSelector } from "react-compiler-runtime"; +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess function App() { const $ = _c(3); const context = useContext(MyContext); diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.js index 2eefdcd76f9fd..fa511732e3fb2 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.js +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.js @@ -1,4 +1,4 @@ -// @enableLowerContextAccess +// @lowerContextAccess function App() { const context = useContext(MyContext); const [foo] = context; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.expect.md index 8f4ec58c6de89..2e5b517ef23b1 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.expect.md @@ -2,7 +2,7 @@ ## Input ```javascript -// @enableLowerContextAccess +// @lowerContextAccess function App() { const { joe: {foo}, @@ -16,7 +16,8 @@ function App() { ## Code ```javascript -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +import { useContext_withSelector } from "react-compiler-runtime"; +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess function App() { const $ = _c(3); const { joe: t0, bar } = useContext(MyContext); diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.js index ceecaaaaffb43..89a7ccd7db47d 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.js +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.js @@ -1,4 +1,4 @@ -// @enableLowerContextAccess +// @lowerContextAccess function App() { const { joe: {foo}, diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.expect.md index 67350a6ee0340..482ed458237e5 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.expect.md @@ -2,7 +2,7 @@ ## Input ```javascript -// @enableLowerContextAccess +// @lowerContextAccess function App() { const context = useContext(MyContext); const foo = context.foo; @@ -15,7 +15,8 @@ function App() { ## Code ```javascript -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess +import { useContext_withSelector } from "react-compiler-runtime"; +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess function App() { const $ = _c(3); const context = useContext(MyContext); diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.js index 9bc2bfadda871..a055114c1ec7e 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.js +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.js @@ -1,4 +1,4 @@ -// @enableLowerContextAccess +// @lowerContextAccess function App() { const context = useContext(MyContext); const foo = context.foo; diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts index 8031038f3c992..5c2b452093de4 100644 --- a/compiler/packages/snap/src/compiler.ts +++ b/compiler/packages/snap/src/compiler.ts @@ -154,6 +154,20 @@ function makePluginOptions( .map(s => s.trim()) .filter(s => s.length > 0); } + if (firstLine.includes('@enableEmitFreeze')) { + enableEmitFreeze = { + source: 'react-compiler-runtime', + importSpecifierName: 'makeReadOnly', + }; + } + + let lowerContextAccess = null; + if (firstLine.includes('@lowerContextAccess')) { + lowerContextAccess = { + source: 'react-compiler-runtime', + importSpecifierName: 'useContext_withSelector', + }; + } let logs: Array<{filename: string | null; event: LoggerEvent}> = []; let logger: Logger | null = null; @@ -207,6 +221,7 @@ function makePluginOptions( hookPattern, validatePreserveExistingMemoizationGuarantees, enableChangeDetectionForDebugging, + lowerContextAccess, }, compilationMode, logger,