@samitouri / QOS-React / commits / 83cc13f746

[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: 00e161b988c8f8a1cf96efff8095f050cb534cc1 Pull Request resolved: https://github.com/facebook/react/pull/30612

Sathya Gunsasekaran committed Aug 7, 2024 at 15:55 UTC 83cc13f74676920bf6b189af9665bd343e056d58
19 files changed +95 -33
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+2 -2
@@ -205,8 +205,8 @@ function* runWithEnvironment(
205 validateNoCapitalizedCalls(hir);
206 }
207
208 - if (env.config.enableLowerContextAccess) {
209 - lowerContextAccess(hir);
208 + if (env.config.lowerContextAccess) {
209 + lowerContextAccess(hir, env.config.lowerContextAccess);
210 }
211
212 analyseFunctions(hir);
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Program.ts
+5
@@ -511,6 +511,11 @@ export function compileProgram(
511 externalFunctions.push(gating);
512 }
513
514 + const lowerContextAccess = pass.opts.environment?.lowerContextAccess;
515 + if (lowerContextAccess) {
516 + externalFunctions.push(tryParseExternalFunction(lowerContextAccess));
517 + }
518 +
519 const enableEmitInstrumentForget =
520 pass.opts.environment?.enableEmitInstrumentForget;
521 if (enableEmitInstrumentForget != null) {
compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+7 -4
@@ -436,8 +436,11 @@ const EnvironmentConfigSchema = z.object({
436 enableTreatRefLikeIdentifiersAsRefs: z.boolean().nullable().default(false),
437
438 /*
439 - * If enabled, this lowers any calls to `useContext` hook to use a selector
440 - * function.
439 + * If specified a value, the compiler lowers any calls to `useContext` to use
440 + * this value as the callee.
441 + *
442 + * A selector function is compiled and passed as an argument along with the
443 + * context to this function call.
444 *
445 * The compiler automatically figures out the keys by looking for the immediate
446 * destructuring of the return value from the useContext call. In the future,
@@ -449,10 +452,10 @@ const EnvironmentConfigSchema = z.object({
452 * const {foo, bar} = useContext(MyContext);
453 *
454 * // output
452 - * const {foo, bar} = useContext(MyContext, (c) => [c.foo, c.bar]);
455 + * const {foo, bar} = useCompiledContext(MyContext, (c) => [c.foo, c.bar]);
456 * ```
457 */
455 - enableLowerContextAccess: z.boolean().nullable().default(false),
458 + lowerContextAccess: ExternalFunctionSchema.nullish(),
459 });
460
461 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;
compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts
+41 -3
@@ -11,10 +11,12 @@ import {
11 CallExpression,
12 Destructure,
13 Environment,
14 + ExternalFunction,
15 GeneratedSource,
16 HIRFunction,
17 IdentifierId,
18 Instruction,
19 + LoadGlobal,
20 LoadLocal,
21 Place,
22 PropertyLoad,
@@ -29,7 +31,10 @@ import {createTemporaryPlace} from '../HIR/HIRBuilder';
31 import {enterSSA} from '../SSA';
32 import {inferTypes} from '../TypeInference';
33
32 -export function lowerContextAccess(fn: HIRFunction): void {
34 +export function lowerContextAccess(
35 + fn: HIRFunction,
36 + loweredContextCallee: ExternalFunction,
37 +): void {
38 const contextAccess: Map<IdentifierId, CallExpression> = new Map();
39 const contextKeys: Map<IdentifierId, Array<string>> = new Map();
40
@@ -84,13 +89,23 @@ export function lowerContextAccess(fn: HIRFunction): void {
89 isUseContextHookType(value.callee.identifier) &&
90 contextKeys.has(lvalue.identifier.id)
91 ) {
87 - const keys = contextKeys.get(lvalue.identifier.id)!;
88 - const selectorFnInstr = emitSelectorFn(fn.env, keys);
92 + const loweredContextCalleeInstr = emitLoadLoweredContextCallee(
93 + fn.env,
94 + loweredContextCallee,
95 + );
96 +
97 if (nextInstructions === null) {
98 nextInstructions = block.instructions.slice(0, i);
99 }
100 + nextInstructions.push(loweredContextCalleeInstr);
101 +
102 + const keys = contextKeys.get(lvalue.identifier.id)!;
103 + const selectorFnInstr = emitSelectorFn(fn.env, keys);
104 nextInstructions.push(selectorFnInstr);
105
106 + const lowerContextCallId = loweredContextCalleeInstr.lvalue;
107 + value.callee = lowerContextCallId;
108 +
109 const selectorFn = selectorFnInstr.lvalue;
110 value.args.push(selectorFn);
111 }
@@ -104,9 +119,32 @@ export function lowerContextAccess(fn: HIRFunction): void {
119 }
120 }
121 markInstructionIds(fn.body);
122 + inferTypes(fn);
123 }
124 }
125
126 +function emitLoadLoweredContextCallee(
127 + env: Environment,
128 + loweredContextCallee: ExternalFunction,
129 +): Instruction {
130 + const loadGlobal: LoadGlobal = {
131 + kind: 'LoadGlobal',
132 + binding: {
133 + kind: 'ImportNamespace',
134 + module: loweredContextCallee.source,
135 + name: loweredContextCallee.importSpecifierName,
136 + },
137 + loc: GeneratedSource,
138 + };
139 +
140 + return {
141 + id: makeInstructionId(0),
142 + loc: GeneratedSource,
143 + lvalue: createTemporaryPlace(env, GeneratedSource),
144 + value: loadGlobal,
145 + };
146 +}
147 +
148 function getContextKeys(value: Destructure): Array<string> | null {
149 const keys = [];
150 const pattern = value.lvalue.pattern;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.expect.md
+5 -4
@@ -2,7 +2,7 @@
2 ## Input
3
4 ```javascript
5 -// @enableLowerContextAccess
5 +// @lowerContextAccess
6 function App() {
7 const {foo} = useContext(MyContext);
8 const {bar} = useContext(MyContext);
@@ -14,11 +14,12 @@ function App() {
14 ## Code
15
16 ```javascript
17 -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
17 +import { useContext_withSelector } from "react-compiler-runtime";
18 +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
19 function App() {
20 const $ = _c(3);
20 - const { foo } = useContext(MyContext, _temp);
21 - const { bar } = useContext(MyContext, _temp2);
21 + const { foo } = useContext_withSelector(MyContext, _temp);
22 + const { bar } = useContext_withSelector(MyContext, _temp2);
23 let t0;
24 if ($[0] !== foo || $[1] !== bar) {
25 t0 = <Bar foo={foo} bar={bar} />;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.js
+1 -1
@@ -1,4 +1,4 @@
1 -// @enableLowerContextAccess
1 +// @lowerContextAccess
2 function App() {
3 const {foo} = useContext(MyContext);
4 const {bar} = useContext(MyContext);
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.expect.md
+4 -3
@@ -2,7 +2,7 @@
2 ## Input
3
4 ```javascript
5 -// @enableLowerContextAccess
5 +// @lowerContextAccess
6 function App() {
7 const {foo, bar} = useContext(MyContext);
8 return <Bar foo={foo} bar={bar} />;
@@ -13,10 +13,11 @@ function App() {
13 ## Code
14
15 ```javascript
16 -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
16 +import { useContext_withSelector } from "react-compiler-runtime";
17 +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
18 function App() {
19 const $ = _c(3);
19 - const { foo, bar } = useContext(MyContext, _temp);
20 + const { foo, bar } = useContext_withSelector(MyContext, _temp);
21 let t0;
22 if ($[0] !== foo || $[1] !== bar) {
23 t0 = <Bar foo={foo} bar={bar} />;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.js
+1 -1
@@ -1,4 +1,4 @@
1 -// @enableLowerContextAccess
1 +// @lowerContextAccess
2 function App() {
3 const {foo, bar} = useContext(MyContext);
4 return <Bar foo={foo} bar={bar} />;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.expect.md
+3 -2
@@ -2,7 +2,7 @@
2 ## Input
3
4 ```javascript
5 -// @enableLowerContextAccess
5 +// @lowerContextAccess
6 function App() {
7 const [foo, bar] = useContext(MyContext);
8 return <Bar foo={foo} bar={bar} />;
@@ -13,7 +13,8 @@ function App() {
13 ## Code
14
15 ```javascript
16 -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
16 +import { useContext_withSelector } from "react-compiler-runtime";
17 +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
18 function App() {
19 const $ = _c(3);
20 const [foo, bar] = useContext(MyContext);
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.js
+1 -1
@@ -1,4 +1,4 @@
1 -// @enableLowerContextAccess
1 +// @lowerContextAccess
2 function App() {
3 const [foo, bar] = useContext(MyContext);
4 return <Bar foo={foo} bar={bar} />;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.expect.md
+3 -2
@@ -2,7 +2,7 @@
2 ## Input
3
4 ```javascript
5 -// @enableLowerContextAccess
5 +// @lowerContextAccess
6 function App() {
7 const context = useContext(MyContext);
8 const {foo} = context;
@@ -15,7 +15,8 @@ function App() {
15 ## Code
16
17 ```javascript
18 -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
18 +import { useContext_withSelector } from "react-compiler-runtime";
19 +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
20 function App() {
21 const $ = _c(3);
22 const context = useContext(MyContext);
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.js
+1 -1
@@ -1,4 +1,4 @@
1 -// @enableLowerContextAccess
1 +// @lowerContextAccess
2 function App() {
3 const context = useContext(MyContext);
4 const {foo} = context;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.expect.md
+3 -2
@@ -2,7 +2,7 @@
2 ## Input
3
4 ```javascript
5 -// @enableLowerContextAccess
5 +// @lowerContextAccess
6 function App() {
7 const context = useContext(MyContext);
8 const [foo] = context;
@@ -15,7 +15,8 @@ function App() {
15 ## Code
16
17 ```javascript
18 -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
18 +import { useContext_withSelector } from "react-compiler-runtime";
19 +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
20 function App() {
21 const $ = _c(3);
22 const context = useContext(MyContext);
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.js
+1 -1
@@ -1,4 +1,4 @@
1 -// @enableLowerContextAccess
1 +// @lowerContextAccess
2 function App() {
3 const context = useContext(MyContext);
4 const [foo] = context;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.expect.md
+3 -2
@@ -2,7 +2,7 @@
2 ## Input
3
4 ```javascript
5 -// @enableLowerContextAccess
5 +// @lowerContextAccess
6 function App() {
7 const {
8 joe: {foo},
@@ -16,7 +16,8 @@ function App() {
16 ## Code
17
18 ```javascript
19 -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
19 +import { useContext_withSelector } from "react-compiler-runtime";
20 +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
21 function App() {
22 const $ = _c(3);
23 const { joe: t0, bar } = useContext(MyContext);
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.js
+1 -1
@@ -1,4 +1,4 @@
1 -// @enableLowerContextAccess
1 +// @lowerContextAccess
2 function App() {
3 const {
4 joe: {foo},
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.expect.md
+3 -2
@@ -2,7 +2,7 @@
2 ## Input
3
4 ```javascript
5 -// @enableLowerContextAccess
5 +// @lowerContextAccess
6 function App() {
7 const context = useContext(MyContext);
8 const foo = context.foo;
@@ -15,7 +15,8 @@ function App() {
15 ## Code
16
17 ```javascript
18 -import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
18 +import { useContext_withSelector } from "react-compiler-runtime";
19 +import { c as _c } from "react/compiler-runtime"; // @lowerContextAccess
20 function App() {
21 const $ = _c(3);
22 const context = useContext(MyContext);
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.js
+1 -1
@@ -1,4 +1,4 @@
1 -// @enableLowerContextAccess
1 +// @lowerContextAccess
2 function App() {
3 const context = useContext(MyContext);
4 const foo = context.foo;
compiler/packages/snap/src/compiler.ts
+9
@@ -155,6 +155,14 @@ function makePluginOptions(
155 .filter(s => s.length > 0);
156 }
157
158 + let lowerContextAccess = null;
159 + if (firstLine.includes('@lowerContextAccess')) {
160 + lowerContextAccess = {
161 + source: 'react-compiler-runtime',
162 + importSpecifierName: 'useContext_withSelector',
163 + };
164 + }
165 +
166 let logs: Array<{filename: string | null; event: LoggerEvent}> = [];
167 let logger: Logger | null = null;
168 if (firstLine.includes('@logger')) {
@@ -207,6 +215,7 @@ function makePluginOptions(
215 hookPattern,
216 validatePreserveExistingMemoizationGuarantees,
217 enableChangeDetectionForDebugging,
218 + lowerContextAccess,
219 },
220 compilationMode,
221 logger,