@samitouri / QOS-React-1 / commits / 8c89fa7643

[compiler][hir-rewrite] Infer non-null props, destructure source

Followup from #30894. This adds a new flagged mode `enablePropagateScopeDepsInHIR: "enabled_with_optimizations"`, under which we infer more hoistable loads: - it's always safe to evaluate loads from `props` (i.e. first parameter of a `component`) - destructuring sources are safe to evaluate loads from (e.g. given `{x} = obj`, we infer that it's safe to evaluate obj.y) - computed load sources are safe to evaluate loads from (e.g. given `arr[0]`, we can infer that it's safe to evaluate arr.length) ghstack-source-id: 32f3bb72e9f85922825579bd785d636f4ccf724d Pull Request resolved: https://github.com/facebook/react/pull/31033

Mofei Zhang committed Sep 30, 2024 at 12:24 UTC 8c89fa76430b3d9fbecd6d535c93171d22f0377f
5 files changed +289 -44
compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts
+95 -44
@@ -8,6 +8,7 @@ import {
8 HIRFunction,
9 Identifier,
10 IdentifierId,
11 + InstructionId,
12 Place,
13 ReactiveScopeDependency,
14 ScopeId,
@@ -66,7 +67,7 @@ export function collectHoistablePropertyLoads(
67 fn: HIRFunction,
68 temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
69 ): ReadonlyMap<ScopeId, BlockInfo> {
69 - const nodes = collectPropertyLoadsInBlocks(fn, temporaries);
70 + const nodes = collectNonNullsInBlocks(fn, temporaries);
71 propagateNonNull(fn, nodes);
72
73 const nodesKeyedByScopeId = new Map<ScopeId, BlockInfo>();
@@ -165,7 +166,7 @@ type PropertyLoadNode =
166 class Tree {
167 roots: Map<Identifier, RootNode> = new Map();
168
168 - #getOrCreateRoot(identifier: Identifier): PropertyLoadNode {
169 + getOrCreateRoot(identifier: Identifier): PropertyLoadNode {
170 /**
171 * Reads from a statically scoped variable are always safe in JS,
172 * with the exception of TDZ (not addressed by this pass).
@@ -207,17 +208,15 @@ class Tree {
208 }
209
210 getPropertyLoadNode(n: ReactiveScopeDependency): PropertyLoadNode {
210 - CompilerError.invariant(n.path.length > 0, {
211 - reason:
212 - '[CollectHoistablePropertyLoads] Expected property node, found root node',
213 - loc: GeneratedSource,
214 - });
211 /**
212 * We add ReactiveScopeDependencies according to instruction ordering,
213 * so all subpaths of a PropertyLoad should already exist
214 * (e.g. a.b is added before a.b.c),
215 */
220 - let currNode = this.#getOrCreateRoot(n.identifier);
216 + let currNode = this.getOrCreateRoot(n.identifier);
217 + if (n.path.length === 0) {
218 + return currNode;
219 + }
220 for (let i = 0; i < n.path.length - 1; i++) {
221 currNode = assertNonNull(currNode.properties.get(n.path[i].property));
222 }
@@ -226,10 +225,44 @@ class Tree {
225 }
226 }
227
229 -function collectPropertyLoadsInBlocks(
228 +function pushPropertyLoadNode(
229 + loadSource: Identifier,
230 + loadSourceNode: PropertyLoadNode,
231 + instrId: InstructionId,
232 + knownImmutableIdentifiers: Set<IdentifierId>,
233 + result: Set<PropertyLoadNode>,
234 +): void {
235 + /**
236 + * Since this runs *after* buildReactiveScopeTerminals, identifier mutable ranges
237 + * are not valid with respect to current instruction id numbering.
238 + * We use attached reactive scope ranges as a proxy for mutable range, but this
239 + * is an overestimate as (1) scope ranges merge and align to form valid program
240 + * blocks and (2) passes like MemoizeFbtAndMacroOperands may assign scopes to
241 + * non-mutable identifiers.
242 + *
243 + * See comment at top of function for why we track known immutable identifiers.
244 + */
245 + const isMutableAtInstr =
246 + loadSource.mutableRange.end > loadSource.mutableRange.start + 1 &&
247 + loadSource.scope != null &&
248 + inRange({id: instrId}, loadSource.scope.range);
249 + if (
250 + !isMutableAtInstr ||
251 + knownImmutableIdentifiers.has(loadSourceNode.fullPath.identifier.id)
252 + ) {
253 + let curr: PropertyLoadNode | null = loadSourceNode;
254 + while (curr != null) {
255 + result.add(curr);
256 + curr = curr.parent;
257 + }
258 + }
259 +}
260 +
261 +function collectNonNullsInBlocks(
262 fn: HIRFunction,
263 temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
264 ): ReadonlyMap<BlockId, BlockInfo> {
265 + const tree = new Tree();
266 /**
267 * Due to current limitations of mutable range inference, there are edge cases in
268 * which we infer known-immutable values (e.g. props or hook params) to have a
@@ -238,53 +271,70 @@ function collectPropertyLoadsInBlocks(
271 * We track known immutable identifiers to reduce regressions (as PropagateScopeDeps
272 * is being rewritten to HIR).
273 */
241 - const knownImmutableIdentifiers = new Set<Identifier>();
274 + const knownImmutableIdentifiers = new Set<IdentifierId>();
275 if (fn.fnType === 'Component' || fn.fnType === 'Hook') {
276 for (const p of fn.params) {
277 if (p.kind === 'Identifier') {
245 - knownImmutableIdentifiers.add(p.identifier);
278 + knownImmutableIdentifiers.add(p.identifier.id);
279 }
280 }
281 }
249 - const tree = new Tree();
282 + /**
283 + * Known non-null objects such as functional component props can be safely
284 + * read from any block.
285 + */
286 + const knownNonNullIdentifiers = new Set<PropertyLoadNode>();
287 + if (
288 + fn.fnType === 'Component' &&
289 + fn.params.length > 0 &&
290 + fn.params[0].kind === 'Identifier'
291 + ) {
292 + const identifier = fn.params[0].identifier;
293 + knownNonNullIdentifiers.add(tree.getOrCreateRoot(identifier));
294 + }
295 const nodes = new Map<BlockId, BlockInfo>();
296 for (const [_, block] of fn.body.blocks) {
252 - const assumedNonNullObjects = new Set<PropertyLoadNode>();
297 + const assumedNonNullObjects = new Set<PropertyLoadNode>(
298 + knownNonNullIdentifiers,
299 + );
300 for (const instr of block.instructions) {
301 if (instr.value.kind === 'PropertyLoad') {
255 - const property = getProperty(
256 - instr.value.object,
257 - instr.value.property,
258 - temporaries,
302 + const source = temporaries.get(instr.value.object.identifier.id) ?? {
303 + identifier: instr.value.object.identifier,
304 + path: [],
305 + };
306 + pushPropertyLoadNode(
307 + instr.value.object.identifier,
308 + tree.getPropertyLoadNode(source),
309 + instr.id,
310 + knownImmutableIdentifiers,
311 + assumedNonNullObjects,
312 );
260 - const propertyNode = tree.getPropertyLoadNode(property);
261 - const object = instr.value.object.identifier;
262 - /**
263 - * Since this runs *after* buildReactiveScopeTerminals, identifier mutable ranges
264 - * are not valid with respect to current instruction id numbering.
265 - * We use attached reactive scope ranges as a proxy for mutable range, but this
266 - * is an overestimate as (1) scope ranges merge and align to form valid program
267 - * blocks and (2) passes like MemoizeFbtAndMacroOperands may assign scopes to
268 - * non-mutable identifiers.
269 - *
270 - * See comment at top of function for why we track known immutable identifiers.
271 - */
272 - const isMutableAtInstr =
273 - object.mutableRange.end > object.mutableRange.start + 1 &&
274 - object.scope != null &&
275 - inRange(instr, object.scope.range);
276 - if (
277 - !isMutableAtInstr ||
278 - knownImmutableIdentifiers.has(propertyNode.fullPath.identifier)
279 - ) {
280 - let curr = propertyNode.parent;
281 - while (curr != null) {
282 - assumedNonNullObjects.add(curr);
283 - curr = curr.parent;
284 - }
313 + } else if (instr.value.kind === 'Destructure') {
314 + const source = instr.value.value.identifier.id;
315 + const sourceNode = temporaries.get(source);
316 + if (sourceNode != null) {
317 + pushPropertyLoadNode(
318 + instr.value.value.identifier,
319 + tree.getPropertyLoadNode(sourceNode),
320 + instr.id,
321 + knownImmutableIdentifiers,
322 + assumedNonNullObjects,
323 + );
324 + }
325 + } else if (instr.value.kind === 'ComputedLoad') {
326 + const source = instr.value.object.identifier.id;
327 + const sourceNode = temporaries.get(source);
328 + if (sourceNode != null) {
329 + pushPropertyLoadNode(
330 + instr.value.object.identifier,
331 + tree.getPropertyLoadNode(sourceNode),
332 + instr.id,
333 + knownImmutableIdentifiers,
334 + assumedNonNullObjects,
335 + );
336 }
337 }
287 - // TODO handle destructuring
338 }
339
340 nodes.set(block.id, {
@@ -449,10 +499,11 @@ function propagateNonNull(
499 );
500
501 for (const [id, node] of nodes) {
452 - node.assumedNonNullObjects = Set_union(
502 + const assumedNonNullObjects = Set_union(
503 assertNonNull(fromEntry.get(id)),
504 assertNonNull(fromExit.get(id)),
505 );
506 + node.assumedNonNullObjects = assumedNonNullObjects;
507 }
508 }
509
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/infer-component-props-non-null.expect.md new
+60
@@ -0,0 +1,60 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enablePropagateDepsInHIR
6 +import {identity, Stringify} from 'shared-runtime';
7 +
8 +function Foo(props) {
9 + /**
10 + * props.value should be inferred as the dependency of this scope
11 + * since we know that props is safe to read from (i.e. non-null)
12 + * as it is arg[0] of a component function
13 + */
14 + const arr = [];
15 + if (cond) {
16 + arr.push(identity(props.value));
17 + }
18 + return <Stringify arr={arr} />;
19 +}
20 +
21 +export const FIXTURE_ENTRYPOINT = {
22 + fn: Foo,
23 + params: [{value: 2}],
24 +};
25 +
26 +```
27 +
28 +## Code
29 +
30 +```javascript
31 +import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
32 +import { identity, Stringify } from "shared-runtime";
33 +
34 +function Foo(props) {
35 + const $ = _c(2);
36 + let t0;
37 + if ($[0] !== props.value) {
38 + const arr = [];
39 + if (cond) {
40 + arr.push(identity(props.value));
41 + }
42 +
43 + t0 = <Stringify arr={arr} />;
44 + $[0] = props.value;
45 + $[1] = t0;
46 + } else {
47 + t0 = $[1];
48 + }
49 + return t0;
50 +}
51 +
52 +export const FIXTURE_ENTRYPOINT = {
53 + fn: Foo,
54 + params: [{ value: 2 }],
55 +};
56 +
57 +```
58 +
59 +### Eval output
60 +(kind: exception) cond is not defined
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/infer-component-props-non-null.tsx new
+20
@@ -0,0 +1,20 @@
1 +// @enablePropagateDepsInHIR
2 +import {identity, Stringify} from 'shared-runtime';
3 +
4 +function Foo(props) {
5 + /**
6 + * props.value should be inferred as the dependency of this scope
7 + * since we know that props is safe to read from (i.e. non-null)
8 + * as it is arg[0] of a component function
9 + */
10 + const arr = [];
11 + if (cond) {
12 + arr.push(identity(props.value));
13 + }
14 + return <Stringify arr={arr} />;
15 +}
16 +
17 +export const FIXTURE_ENTRYPOINT = {
18 + fn: Foo,
19 + params: [{value: 2}],
20 +};
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/infer-non-null-destructure.expect.md new
+91
@@ -0,0 +1,91 @@
1 +
2 +## Input
3 +
4 +```javascript
5 +// @enablePropagateDepsInHIR
6 +import {identity, useIdentity} from 'shared-runtime';
7 +
8 +function useFoo({arg, cond}: {arg: number; cond: boolean}) {
9 + const maybeObj = useIdentity({value: arg});
10 + const {value} = maybeObj;
11 + useIdentity(null);
12 + /**
13 + * maybeObj.value should be inferred as the dependency of this scope
14 + * since we know that maybeObj is safe to read from (i.e. non-null)
15 + * due to the above destructuring instruction
16 + */
17 + const arr = [];
18 + if (cond) {
19 + arr.push(identity(maybeObj.value));
20 + }
21 + return {arr, value};
22 +}
23 +
24 +export const FIXTURE_ENTRYPOINT = {
25 + fn: useFoo,
26 + params: [{arg: 2, cond: false}],
27 +};
28 +
29 +```
30 +
31 +## Code
32 +
33 +```javascript
34 +import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
35 +import { identity, useIdentity } from "shared-runtime";
36 +
37 +function useFoo(t0) {
38 + const $ = _c(10);
39 + const { arg, cond } = t0;
40 + let t1;
41 + if ($[0] !== arg) {
42 + t1 = { value: arg };
43 + $[0] = arg;
44 + $[1] = t1;
45 + } else {
46 + t1 = $[1];
47 + }
48 + const maybeObj = useIdentity(t1);
49 + const { value } = maybeObj;
50 + useIdentity(null);
51 + let arr;
52 + if ($[2] !== cond || $[3] !== maybeObj.value) {
53 + arr = [];
54 + if (cond) {
55 + let t2;
56 + if ($[5] !== maybeObj.value) {
57 + t2 = identity(maybeObj.value);
58 + $[5] = maybeObj.value;
59 + $[6] = t2;
60 + } else {
61 + t2 = $[6];
62 + }
63 + arr.push(t2);
64 + }
65 + $[2] = cond;
66 + $[3] = maybeObj.value;
67 + $[4] = arr;
68 + } else {
69 + arr = $[4];
70 + }
71 + let t2;
72 + if ($[7] !== arr || $[8] !== value) {
73 + t2 = { arr, value };
74 + $[7] = arr;
75 + $[8] = value;
76 + $[9] = t2;
77 + } else {
78 + t2 = $[9];
79 + }
80 + return t2;
81 +}
82 +
83 +export const FIXTURE_ENTRYPOINT = {
84 + fn: useFoo,
85 + params: [{ arg: 2, cond: false }],
86 +};
87 +
88 +```
89 +
90 +### Eval output
91 +(kind: ok) {"arr":[],"value":2}
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/infer-non-null-destructure.ts new
+23
@@ -0,0 +1,23 @@
1 +// @enablePropagateDepsInHIR
2 +import {identity, useIdentity} from 'shared-runtime';
3 +
4 +function useFoo({arg, cond}: {arg: number; cond: boolean}) {
5 + const maybeObj = useIdentity({value: arg});
6 + const {value} = maybeObj;
7 + useIdentity(null);
8 + /**
9 + * maybeObj.value should be inferred as the dependency of this scope
10 + * since we know that maybeObj is safe to read from (i.e. non-null)
11 + * due to the above destructuring instruction
12 + */
13 + const arr = [];
14 + if (cond) {
15 + arr.push(identity(maybeObj.value));
16 + }
17 + return {arr, value};
18 +}
19 +
20 +export const FIXTURE_ENTRYPOINT = {
21 + fn: useFoo,
22 + params: [{arg: 2, cond: false}],
23 +};