@samitouri / QOS-React / commits / e948a5ac68

[compiler] Add lowerContextAccess pass

*This is only for internal profiling, not intended to ship.* This pass is intended to be used with https://github.com/facebook/react/pull/30407. This pass synthesizes selector functions by collecting immediately destructured context acesses. We bailout for other types of context access. This pass lowers context access to use a selector function by passing the synthesized selector function as the second argument. ghstack-source-id: 92d0f6ff2fe95cda93f66786f56e97ba9ace95fa Pull Request resolved: https://github.com/facebook/react/pull/30548

Sathya Gunsasekaran committed Aug 7, 2024 at 15:55 UTC e948a5ac6876a016ff57f7177f93452ea2eeb574
17 files changed +574
compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+5
@@ -103,6 +103,7 @@ import {
103 import {validateLocalsNotReassignedAfterRender} from '../Validation/ValidateLocalsNotReassignedAfterRender';
104 import {outlineFunctions} from '../Optimization/OutlineFunctions';
105 import {propagatePhiTypes} from '../TypeInference/PropagatePhiTypes';
106 +import {lowerContextAccess} from '../Optimization/LowerContextAccess';
107
108 export type CompilerPipelineValue =
109 | {kind: 'ast'; name: string; value: CodegenFunction}
@@ -204,6 +205,10 @@ function* runWithEnvironment(
205 validateNoCapitalizedCalls(hir);
206 }
207
208 + if (env.config.enableLowerContextAccess) {
209 + lowerContextAccess(hir);
210 + }
211 +
212 analyseFunctions(hir);
213 yield log({kind: 'hir', name: 'AnalyseFunctions', value: hir});
214
compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts new
+257
@@ -0,0 +1,257 @@
1 +/**
2 + * Copyright (c) Meta Platforms, Inc. and affiliates.
3 + *
4 + * This source code is licensed under the MIT license found in the
5 + * LICENSE file in the root directory of this source tree.
6 + */
7 +
8 +import {
9 + ArrayExpression,
10 + BasicBlock,
11 + CallExpression,
12 + Destructure,
13 + Environment,
14 + GeneratedSource,
15 + HIRFunction,
16 + IdentifierId,
17 + Instruction,
18 + LoadLocal,
19 + Place,
20 + PropertyLoad,
21 + isUseContextHookType,
22 + makeBlockId,
23 + makeInstructionId,
24 + markInstructionIds,
25 + promoteTemporary,
26 + reversePostorderBlocks,
27 +} from '../HIR';
28 +import {createTemporaryPlace} from '../HIR/HIRBuilder';
29 +import {enterSSA} from '../SSA';
30 +import {inferTypes} from '../TypeInference';
31 +
32 +export function lowerContextAccess(fn: HIRFunction): void {
33 + const contextAccess: Map<IdentifierId, CallExpression> = new Map();
34 + const contextKeys: Map<IdentifierId, Array<string>> = new Map();
35 +
36 + // collect context access and keys
37 + for (const [, block] of fn.body.blocks) {
38 + for (const instr of block.instructions) {
39 + const {value, lvalue} = instr;
40 +
41 + if (
42 + value.kind === 'CallExpression' &&
43 + isUseContextHookType(value.callee.identifier)
44 + ) {
45 + contextAccess.set(lvalue.identifier.id, value);
46 + continue;
47 + }
48 +
49 + if (value.kind !== 'Destructure') {
50 + continue;
51 + }
52 +
53 + const destructureId = value.value.identifier.id;
54 + if (!contextAccess.has(destructureId)) {
55 + continue;
56 + }
57 +
58 + const keys = getContextKeys(value);
59 + if (keys === null) {
60 + return;
61 + }
62 +
63 + if (contextKeys.has(destructureId)) {
64 + /*
65 + * TODO(gsn): Add support for accessing context over multiple
66 + * statements.
67 + */
68 + return;
69 + } else {
70 + contextKeys.set(destructureId, keys);
71 + }
72 + }
73 + }
74 +
75 + if (contextAccess.size > 0) {
76 + for (const [, block] of fn.body.blocks) {
77 + let nextInstructions: Array<Instruction> | null = null;
78 +
79 + for (let i = 0; i < block.instructions.length; i++) {
80 + const instr = block.instructions[i];
81 + const {lvalue, value} = instr;
82 + if (
83 + value.kind === 'CallExpression' &&
84 + isUseContextHookType(value.callee.identifier) &&
85 + contextKeys.has(lvalue.identifier.id)
86 + ) {
87 + const keys = contextKeys.get(lvalue.identifier.id)!;
88 + const selectorFnInstr = emitSelectorFn(fn.env, keys);
89 + if (nextInstructions === null) {
90 + nextInstructions = block.instructions.slice(0, i);
91 + }
92 + nextInstructions.push(selectorFnInstr);
93 +
94 + const selectorFn = selectorFnInstr.lvalue;
95 + value.args.push(selectorFn);
96 + }
97 +
98 + if (nextInstructions) {
99 + nextInstructions.push(instr);
100 + }
101 + }
102 + if (nextInstructions) {
103 + block.instructions = nextInstructions;
104 + }
105 + }
106 + markInstructionIds(fn.body);
107 + }
108 +}
109 +
110 +function getContextKeys(value: Destructure): Array<string> | null {
111 + const keys = [];
112 + const pattern = value.lvalue.pattern;
113 +
114 + switch (pattern.kind) {
115 + case 'ArrayPattern': {
116 + return null;
117 + }
118 +
119 + case 'ObjectPattern': {
120 + for (const place of pattern.properties) {
121 + if (
122 + place.kind !== 'ObjectProperty' ||
123 + place.type !== 'property' ||
124 + place.key.kind !== 'identifier' ||
125 + place.place.identifier.name === null ||
126 + place.place.identifier.name.kind !== 'named'
127 + ) {
128 + return null;
129 + }
130 + keys.push(place.key.name);
131 + }
132 + return keys;
133 + }
134 + }
135 +}
136 +
137 +function emitPropertyLoad(
138 + env: Environment,
139 + obj: Place,
140 + property: string,
141 +): {instructions: Array<Instruction>; element: Place} {
142 + const loadObj: LoadLocal = {
143 + kind: 'LoadLocal',
144 + place: obj,
145 + loc: GeneratedSource,
146 + };
147 + const object: Place = createTemporaryPlace(env, GeneratedSource);
148 + const loadLocalInstr: Instruction = {
149 + lvalue: object,
150 + value: loadObj,
151 + id: makeInstructionId(0),
152 + loc: GeneratedSource,
153 + };
154 +
155 + const loadProp: PropertyLoad = {
156 + kind: 'PropertyLoad',
157 + object,
158 + property,
159 + loc: GeneratedSource,
160 + };
161 + const element: Place = createTemporaryPlace(env, GeneratedSource);
162 + const loadPropInstr: Instruction = {
163 + lvalue: element,
164 + value: loadProp,
165 + id: makeInstructionId(0),
166 + loc: GeneratedSource,
167 + };
168 + return {
169 + instructions: [loadLocalInstr, loadPropInstr],
170 + element: element,
171 + };
172 +}
173 +
174 +function emitSelectorFn(env: Environment, keys: Array<string>): Instruction {
175 + const obj: Place = createTemporaryPlace(env, GeneratedSource);
176 + promoteTemporary(obj.identifier);
177 + const instr: Array<Instruction> = [];
178 + const elements = [];
179 + for (const key of keys) {
180 + const {instructions, element: prop} = emitPropertyLoad(env, obj, key);
181 + instr.push(...instructions);
182 + elements.push(prop);
183 + }
184 +
185 + const arrayInstr = emitArrayInstr(elements, env);
186 + instr.push(arrayInstr);
187 +
188 + const block: BasicBlock = {
189 + kind: 'block',
190 + id: makeBlockId(0),
191 + instructions: instr,
192 + terminal: {
193 + id: makeInstructionId(0),
194 + kind: 'return',
195 + loc: GeneratedSource,
196 + value: arrayInstr.lvalue,
197 + },
198 + preds: new Set(),
199 + phis: new Set(),
200 + };
201 +
202 + const fn: HIRFunction = {
203 + loc: GeneratedSource,
204 + id: null,
205 + fnType: 'Other',
206 + env,
207 + params: [obj],
208 + returnType: null,
209 + context: [],
210 + effects: null,
211 + body: {
212 + entry: block.id,
213 + blocks: new Map([[block.id, block]]),
214 + },
215 + generator: false,
216 + async: false,
217 + directives: [],
218 + };
219 +
220 + reversePostorderBlocks(fn.body);
221 + markInstructionIds(fn.body);
222 + enterSSA(fn);
223 + inferTypes(fn);
224 +
225 + const fnInstr: Instruction = {
226 + id: makeInstructionId(0),
227 + value: {
228 + kind: 'FunctionExpression',
229 + name: null,
230 + loweredFunc: {
231 + func: fn,
232 + dependencies: [],
233 + },
234 + type: 'ArrowFunctionExpression',
235 + loc: GeneratedSource,
236 + },
237 + lvalue: createTemporaryPlace(env, GeneratedSource),
238 + loc: GeneratedSource,
239 + };
240 + return fnInstr;
241 +}
242 +
243 +function emitArrayInstr(elements: Array<Place>, env: Environment): Instruction {
244 + const array: ArrayExpression = {
245 + kind: 'ArrayExpression',
246 + elements,
247 + loc: GeneratedSource,
248 + };
249 + const arrayLvalue: Place = createTemporaryPlace(env, GeneratedSource);
250 + const arrayInstr: Instruction = {
251 + id: makeInstructionId(0),
252 + value: array,
253 + lvalue: arrayLvalue,
254 + loc: GeneratedSource,
255 + };
256 + return arrayInstr;
257 +}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.expect.md new
+41
@@ -0,0 +1,41 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableLowerContextAccess
6 +function App() {
7 + const {foo} = useContext(MyContext);
8 + const {bar} = useContext(MyContext);
9 + return <Bar foo={foo} bar={bar} />;
10 +}
11 +
12 +```
13 +
14 +## Code
15 +
16 +```javascript
17 +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
18 +function App() {
19 + const $ = _c(3);
20 + const { foo } = useContext(MyContext, _temp);
21 + const { bar } = useContext(MyContext, _temp2);
22 + let t0;
23 + if ($[0] !== foo || $[1] !== bar) {
24 + t0 = <Bar foo={foo} bar={bar} />;
25 + $[0] = foo;
26 + $[1] = bar;
27 + $[2] = t0;
28 + } else {
29 + t0 = $[2];
30 + }
31 + return t0;
32 +}
33 +function _temp2(t0) {
34 + return [t0.bar];
35 +}
36 +function _temp(t0) {
37 + return [t0.foo];
38 +}
39 +
40 +```
41 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-acess-multiple.js new
+6
@@ -0,0 +1,6 @@
1 +// @enableLowerContextAccess
2 +function App() {
3 + const {foo} = useContext(MyContext);
4 + const {bar} = useContext(MyContext);
5 + return <Bar foo={foo} bar={bar} />;
6 +}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.expect.md new
+36
@@ -0,0 +1,36 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableLowerContextAccess
6 +function App() {
7 + const {foo, bar} = useContext(MyContext);
8 + return <Bar foo={foo} bar={bar} />;
9 +}
10 +
11 +```
12 +
13 +## Code
14 +
15 +```javascript
16 +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
17 +function App() {
18 + const $ = _c(3);
19 + const { foo, bar } = useContext(MyContext, _temp);
20 + let t0;
21 + if ($[0] !== foo || $[1] !== bar) {
22 + t0 = <Bar foo={foo} bar={bar} />;
23 + $[0] = foo;
24 + $[1] = bar;
25 + $[2] = t0;
26 + } else {
27 + t0 = $[2];
28 + }
29 + return t0;
30 +}
31 +function _temp(t0) {
32 + return [t0.foo, t0.bar];
33 +}
34 +
35 +```
36 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/lower-context-selector-simple.js new
+5
@@ -0,0 +1,5 @@
1 +// @enableLowerContextAccess
2 +function App() {
3 + const {foo, bar} = useContext(MyContext);
4 + return <Bar foo={foo} bar={bar} />;
5 +}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.expect.md new
+33
@@ -0,0 +1,33 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableLowerContextAccess
6 +function App() {
7 + const [foo, bar] = useContext(MyContext);
8 + return <Bar foo={foo} bar={bar} />;
9 +}
10 +
11 +```
12 +
13 +## Code
14 +
15 +```javascript
16 +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
17 +function App() {
18 + const $ = _c(3);
19 + const [foo, bar] = useContext(MyContext);
20 + let t0;
21 + if ($[0] !== foo || $[1] !== bar) {
22 + t0 = <Bar foo={foo} bar={bar} />;
23 + $[0] = foo;
24 + $[1] = bar;
25 + $[2] = t0;
26 + } else {
27 + t0 = $[2];
28 + }
29 + return t0;
30 +}
31 +
32 +```
33 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-array-destructuring.js new
+5
@@ -0,0 +1,5 @@
1 +// @enableLowerContextAccess
2 +function App() {
3 + const [foo, bar] = useContext(MyContext);
4 + return <Bar foo={foo} bar={bar} />;
5 +}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.expect.md new
+37
@@ -0,0 +1,37 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableLowerContextAccess
6 +function App() {
7 + const context = useContext(MyContext);
8 + const {foo} = context;
9 + const {bar} = context;
10 + return <Bar foo={foo} bar={bar} />;
11 +}
12 +
13 +```
14 +
15 +## Code
16 +
17 +```javascript
18 +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
19 +function App() {
20 + const $ = _c(3);
21 + const context = useContext(MyContext);
22 + const { foo } = context;
23 + const { bar } = context;
24 + let t0;
25 + if ($[0] !== foo || $[1] !== bar) {
26 + t0 = <Bar foo={foo} bar={bar} />;
27 + $[0] = foo;
28 + $[1] = bar;
29 + $[2] = t0;
30 + } else {
31 + t0 = $[2];
32 + }
33 + return t0;
34 +}
35 +
36 +```
37 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-destructure-multiple.js new
+7
@@ -0,0 +1,7 @@
1 +// @enableLowerContextAccess
2 +function App() {
3 + const context = useContext(MyContext);
4 + const {foo} = context;
5 + const {bar} = context;
6 + return <Bar foo={foo} bar={bar} />;
7 +}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.expect.md new
+37
@@ -0,0 +1,37 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableLowerContextAccess
6 +function App() {
7 + const context = useContext(MyContext);
8 + const [foo] = context;
9 + const {bar} = context;
10 + return <Bar foo={foo} bar={bar} />;
11 +}
12 +
13 +```
14 +
15 +## Code
16 +
17 +```javascript
18 +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
19 +function App() {
20 + const $ = _c(3);
21 + const context = useContext(MyContext);
22 + const [foo] = context;
23 + const { bar } = context;
24 + let t0;
25 + if ($[0] !== foo || $[1] !== bar) {
26 + t0 = <Bar foo={foo} bar={bar} />;
27 + $[0] = foo;
28 + $[1] = bar;
29 + $[2] = t0;
30 + } else {
31 + t0 = $[2];
32 + }
33 + return t0;
34 +}
35 +
36 +```
37 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-mixed-array-obj.js new
+7
@@ -0,0 +1,7 @@
1 +// @enableLowerContextAccess
2 +function App() {
3 + const context = useContext(MyContext);
4 + const [foo] = context;
5 + const {bar} = context;
6 + return <Bar foo={foo} bar={bar} />;
7 +}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.expect.md new
+37
@@ -0,0 +1,37 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableLowerContextAccess
6 +function App() {
7 + const {
8 + joe: {foo},
9 + bar,
10 + } = useContext(MyContext);
11 + return <Bar foo={foo} bar={bar} />;
12 +}
13 +
14 +```
15 +
16 +## Code
17 +
18 +```javascript
19 +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
20 +function App() {
21 + const $ = _c(3);
22 + const { joe: t0, bar } = useContext(MyContext);
23 + const { foo } = t0;
24 + let t1;
25 + if ($[0] !== foo || $[1] !== bar) {
26 + t1 = <Bar foo={foo} bar={bar} />;
27 + $[0] = foo;
28 + $[1] = bar;
29 + $[2] = t1;
30 + } else {
31 + t1 = $[2];
32 + }
33 + return t1;
34 +}
35 +
36 +```
37 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-nested-destructuring.js new
+8
@@ -0,0 +1,8 @@
1 +// @enableLowerContextAccess
2 +function App() {
3 + const {
4 + joe: {foo},
5 + bar,
6 + } = useContext(MyContext);
7 + return <Bar foo={foo} bar={bar} />;
8 +}
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.expect.md new
+37
@@ -0,0 +1,37 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enableLowerContextAccess
6 +function App() {
7 + const context = useContext(MyContext);
8 + const foo = context.foo;
9 + const bar = context.bar;
10 + return <Bar foo={foo} bar={bar} />;
11 +}
12 +
13 +```
14 +
15 +## Code
16 +
17 +```javascript
18 +import { c as _c } from "react/compiler-runtime"; // @enableLowerContextAccess
19 +function App() {
20 + const $ = _c(3);
21 + const context = useContext(MyContext);
22 + const foo = context.foo;
23 + const bar = context.bar;
24 + let t0;
25 + if ($[0] !== foo || $[1] !== bar) {
26 + t0 = <Bar foo={foo} bar={bar} />;
27 + $[0] = foo;
28 + $[1] = bar;
29 + $[2] = t0;
30 + } else {
31 + t0 = $[2];
32 + }
33 + return t0;
34 +}
35 +
36 +```
37 +
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/todo.lower-context-access-property-load.js new
+7
@@ -0,0 +1,7 @@
1 +// @enableLowerContextAccess
2 +function App() {
3 + const context = useContext(MyContext);
4 + const foo = context.foo;
5 + const bar = context.bar;
6 + return <Bar foo={foo} bar={bar} />;
7 +}
compiler/packages/snap/src/SproutTodoFilter.ts
+9
@@ -502,6 +502,15 @@ const skipFilter = new Set([
502
503 // needs to be executed as a module
504 'meta-property',
505 +
506 + // needs context lowering support in React
507 + 'todo.lower-context-access-property-load',
508 + 'todo.lower-context-access-nested-destructuring',
509 + 'todo.lower-context-access-mixed-array-obj',
510 + 'todo.lower-context-access-destructure-multiple',
511 + 'todo.lower-context-access-array-destructuring',
512 + 'lower-context-selector-simple',
513 + 'lower-context-acess-multiple',
514 ]);
515
516 export default skipFilter;