[compiler] Infer phi types, extend mutable ranges to account for Store effects
Redo of an earlier (pre-OSS) PR to infer types of phi nodes. There are a few pieces to this: 1. Update InferTypes to infer the type of `phi.id.type`, not the unused `phi.type`. 2. Update the algorithm to verify that all the phi types are actually equal, not just have the same kind. 3. Handle circular types by removing the cycle. However, that reveals another issue: InferMutableRanges currently infers the results of `Store` effects _after_ its fixpoint loop. That was fine when a Store could never occur on a phi (since they wouldn't have a type to get a function signature from). Now though, we can have Store effects occur on phis, and we need to ensure that this correctly updates the mutable range of the phi operands - recursively. See new test that fails without the fixpoint loop. ghstack-source-id: 2e1b02844d3a814dce094b7e3812df799e54343f Pull Request resolved: https://github.com/facebook/react/pull/30796
Joe Savona committed
Aug 22, 2024 at 17:26 UTC
4f546740788ae95ced0338a9cbafc50124b47194
6 files changed
+313
-29
compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRanges.ts
+32
-2
@@ -50,8 +50,38 @@ export function inferMutableRanges(ir: HIRFunction): void {
50
// Re-infer mutable ranges for all values
51
inferMutableLifetimes(ir, true);
52
53
- // Re-infer mutable ranges for aliases
54
- inferMutableRangesForAlias(ir, aliases);
53
+ /**
54
+ * The second inferMutableLifetimes() call updates mutable ranges
55
+ * of values to account for Store effects. Now we need to update
56
+ * all aliases of such values to extend their ranges as well. Note
57
+ * that the store only mutates the the directly aliased value and
58
+ * not any of its inner captured references. For example:
59
+ *
60
+ * ```
61
+ * let y;
62
+ * if (cond) {
63
+ * y = [];
64
+ * } else {
65
+ * y = [{}];
66
+ * }
67
+ * y.push(z);
68
+ * ```
69
+ *
70
+ * The Store effect from the `y.push` modifies the values that `y`
71
+ * directly aliases - the two arrays from the if/else branches -
72
+ * but does not modify values that `y` "contains" such as the
73
+ * object literal or `z`.
74
+ */
75
+ prevAliases = aliases.canonicalize();
76
+ while (true) {
77
+ inferMutableRangesForAlias(ir, aliases);
78
+ inferAliasForPhis(ir, aliases);
79
+ const nextAliases = aliases.canonicalize();
80
+ if (areEqualMaps(prevAliases, nextAliases)) {
81
+ break;
82
+ }
83
+ prevAliases = nextAliases;
84
+ }
85
}
86
87
function areEqualMaps<T>(a: Map<T, T>, b: Map<T, T>): boolean {
compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts
+112
-7
@@ -483,30 +483,135 @@ class Unifier {
483
}
484
485
if (type.kind === 'Phi') {
486
- const operands = new Set(type.operands.map(i => this.get(i).kind));
487
-
488
- CompilerError.invariant(operands.size > 0, {
486
+ CompilerError.invariant(type.operands.length > 0, {
487
reason: 'there should be at least one operand',
488
description: null,
489
loc: null,
490
suggestions: null,
491
});
494
- const kind = operands.values().next().value;
492
496
- // there's only one unique type and it's not a type var
497
- if (operands.size === 1 && kind !== 'Type') {
498
- this.unify(v, type.operands[0]);
493
+ let candidateType: Type | null = null;
494
+ for (const operand of type.operands) {
495
+ const resolved = this.get(operand);
496
+ if (candidateType === null) {
497
+ candidateType = resolved;
498
+ } else if (!typeEquals(resolved, candidateType)) {
499
+ candidateType = null;
500
+ break;
501
+ } // else same type, continue
502
+ }
503
+
504
+ if (candidateType !== null) {
505
+ this.unify(v, candidateType);
506
return;
507
}
508
}
509
510
if (this.occursCheck(v, type)) {
511
+ const resolvedType = this.tryResolveType(v, type);
512
+ if (resolvedType !== null) {
513
+ this.substitutions.set(v.id, resolvedType);
514
+ return;
515
+ }
516
throw new Error('cycle detected');
517
}
518
519
this.substitutions.set(v.id, type);
520
}
521
522
+ tryResolveType(v: TypeVar, type: Type): Type | null {
523
+ switch (type.kind) {
524
+ case 'Phi': {
525
+ /**
526
+ * Resolve the type of the phi by recursively removing `v` as an operand.
527
+ * For example we can end up with types like this:
528
+ *
529
+ * v = Phi [
530
+ * T1
531
+ * T2
532
+ * Phi [
533
+ * T3
534
+ * Phi [
535
+ * T4
536
+ * v <-- cycle!
537
+ * ]
538
+ * ]
539
+ * ]
540
+ *
541
+ * By recursively removing `v`, we end up with:
542
+ *
543
+ * v = Phi [
544
+ * T1
545
+ * T2
546
+ * Phi [
547
+ * T3
548
+ * Phi [
549
+ * T4
550
+ * ]
551
+ * ]
552
+ * ]
553
+ *
554
+ * Which avoids the cycle
555
+ */
556
+ const operands = [];
557
+ for (const operand of type.operands) {
558
+ if (operand.kind === 'Type' && operand.id === v.id) {
559
+ continue;
560
+ }
561
+ const resolved = this.tryResolveType(v, operand);
562
+ if (resolved === null) {
563
+ return null;
564
+ }
565
+ operands.push(resolved);
566
+ }
567
+ return {kind: 'Phi', operands};
568
+ }
569
+ case 'Type': {
570
+ const substitution = this.get(type);
571
+ if (substitution !== type) {
572
+ const resolved = this.tryResolveType(v, substitution);
573
+ if (resolved !== null) {
574
+ this.substitutions.set(type.id, resolved);
575
+ }
576
+ return resolved;
577
+ }
578
+ return type;
579
+ }
580
+ case 'Property': {
581
+ const objectType = this.tryResolveType(v, this.get(type.objectType));
582
+ if (objectType === null) {
583
+ return null;
584
+ }
585
+ return {
586
+ kind: 'Property',
587
+ objectName: type.objectName,
588
+ objectType,
589
+ propertyName: type.propertyName,
590
+ };
591
+ }
592
+ case 'Function': {
593
+ const returnType = this.tryResolveType(v, this.get(type.return));
594
+ if (returnType === null) {
595
+ return null;
596
+ }
597
+ return {
598
+ kind: 'Function',
599
+ return: returnType,
600
+ shapeId: type.shapeId,
601
+ };
602
+ }
603
+ case 'ObjectMethod':
604
+ case 'Object':
605
+ case 'Primitive':
606
+ case 'Poly': {
607
+ return type;
608
+ }
609
+ default: {
610
+ assertExhaustive(type, `Unexpected type kind '${(type as any).kind}'`);
611
+ }
612
+ }
613
+ }
614
+
615
occursCheck(v: TypeVar, type: Type): boolean {
616
if (typeEquals(v, type)) return true;
617
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/phi-type-inference-array-push-consecutive-phis.expect.md
new
+110
@@ -0,0 +1,110 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+import {makeArray} from 'shared-runtime';
6
+
7
+function Component(props) {
8
+ const x = {};
9
+ let y;
10
+ if (props.cond) {
11
+ if (props.cond2) {
12
+ y = [props.value];
13
+ } else {
14
+ y = [props.value2];
15
+ }
16
+ } else {
17
+ y = [];
18
+ }
19
+ // This should be inferred as `<store> y` s.t. `x` can still
20
+ // be independently memoized. *But* this also must properly
21
+ // extend the mutable range of the array literals in the
22
+ // if/else branches
23
+ y.push(x);
24
+
25
+ return [x, y];
26
+}
27
+
28
+export const FIXTURE_ENTRYPOINT = {
29
+ fn: Component,
30
+ params: [{cond: true, cond2: true, value: 42}],
31
+ sequentialRenders: [
32
+ {cond: true, cond2: true, value: 3.14},
33
+ {cond: true, cond2: true, value: 42},
34
+ {cond: true, cond2: true, value: 3.14},
35
+ {cond: true, cond2: false, value2: 3.14},
36
+ {cond: true, cond2: false, value2: 42},
37
+ {cond: true, cond2: false, value2: 3.14},
38
+ {cond: false},
39
+ {cond: false},
40
+ ],
41
+};
42
+
43
+```
44
+
45
+## Code
46
+
47
+```javascript
48
+import { c as _c } from "react/compiler-runtime";
49
+import { makeArray } from "shared-runtime";
50
+
51
+function Component(props) {
52
+ const $ = _c(3);
53
+ let t0;
54
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
55
+ t0 = {};
56
+ $[0] = t0;
57
+ } else {
58
+ t0 = $[0];
59
+ }
60
+ const x = t0;
61
+ let t1;
62
+ if ($[1] !== props) {
63
+ let y;
64
+ if (props.cond) {
65
+ if (props.cond2) {
66
+ y = [props.value];
67
+ } else {
68
+ y = [props.value2];
69
+ }
70
+ } else {
71
+ y = [];
72
+ }
73
+
74
+ y.push(x);
75
+
76
+ t1 = [x, y];
77
+ $[1] = props;
78
+ $[2] = t1;
79
+ } else {
80
+ t1 = $[2];
81
+ }
82
+ return t1;
83
+}
84
+
85
+export const FIXTURE_ENTRYPOINT = {
86
+ fn: Component,
87
+ params: [{ cond: true, cond2: true, value: 42 }],
88
+ sequentialRenders: [
89
+ { cond: true, cond2: true, value: 3.14 },
90
+ { cond: true, cond2: true, value: 42 },
91
+ { cond: true, cond2: true, value: 3.14 },
92
+ { cond: true, cond2: false, value2: 3.14 },
93
+ { cond: true, cond2: false, value2: 42 },
94
+ { cond: true, cond2: false, value2: 3.14 },
95
+ { cond: false },
96
+ { cond: false },
97
+ ],
98
+};
99
+
100
+```
101
+
102
+### Eval output
103
+(kind: ok) [{},[3.14,"[[ cyclic ref *1 ]]"]]
104
+[{},[42,"[[ cyclic ref *1 ]]"]]
105
+[{},[3.14,"[[ cyclic ref *1 ]]"]]
106
+[{},[3.14,"[[ cyclic ref *1 ]]"]]
107
+[{},[42,"[[ cyclic ref *1 ]]"]]
108
+[{},[3.14,"[[ cyclic ref *1 ]]"]]
109
+[{},["[[ cyclic ref *1 ]]"]]
110
+[{},["[[ cyclic ref *1 ]]"]]
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/phi-type-inference-array-push-consecutive-phis.js
new
+37
@@ -0,0 +1,37 @@
1
+import {makeArray} from 'shared-runtime';
2
+
3
+function Component(props) {
4
+ const x = {};
5
+ let y;
6
+ if (props.cond) {
7
+ if (props.cond2) {
8
+ y = [props.value];
9
+ } else {
10
+ y = [props.value2];
11
+ }
12
+ } else {
13
+ y = [];
14
+ }
15
+ // This should be inferred as `<store> y` s.t. `x` can still
16
+ // be independently memoized. *But* this also must properly
17
+ // extend the mutable range of the array literals in the
18
+ // if/else branches
19
+ y.push(x);
20
+
21
+ return [x, y];
22
+}
23
+
24
+export const FIXTURE_ENTRYPOINT = {
25
+ fn: Component,
26
+ params: [{cond: true, cond2: true, value: 42}],
27
+ sequentialRenders: [
28
+ {cond: true, cond2: true, value: 3.14},
29
+ {cond: true, cond2: true, value: 42},
30
+ {cond: true, cond2: true, value: 3.14},
31
+ {cond: true, cond2: false, value2: 3.14},
32
+ {cond: true, cond2: false, value2: 42},
33
+ {cond: true, cond2: false, value2: 3.14},
34
+ {cond: false},
35
+ {cond: false},
36
+ ],
37
+};
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/phi-type-inference-array-push.expect.md
+15
-8
@@ -36,10 +36,17 @@ export const FIXTURE_ENTRYPOINT = {
36
```javascript
37
import { c as _c } from "react/compiler-runtime";
38
function Component(props) {
39
- const $ = _c(2);
39
+ const $ = _c(3);
40
let t0;
41
- if ($[0] !== props) {
42
- const x = {};
41
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
42
+ t0 = {};
43
+ $[0] = t0;
44
+ } else {
45
+ t0 = $[0];
46
+ }
47
+ const x = t0;
48
+ let t1;
49
+ if ($[1] !== props) {
50
let y;
51
if (props.cond) {
52
y = [props.value];
@@ -49,13 +56,13 @@ function Component(props) {
56
57
y.push(x);
58
52
- t0 = [x, y];
53
- $[0] = props;
54
- $[1] = t0;
59
+ t1 = [x, y];
60
+ $[1] = props;
61
+ $[2] = t1;
62
} else {
56
- t0 = $[1];
63
+ t1 = $[2];
64
}
58
- return t0;
65
+ return t1;
66
}
67
68
export const FIXTURE_ENTRYPOINT = {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/phi-type-inference-property-store.expect.md
+7
-12
@@ -32,7 +32,7 @@ export const FIXTURE_ENTRYPOINT = {
32
```javascript
33
import { c as _c } from "react/compiler-runtime"; // @debug
34
function Component(props) {
35
- const $ = _c(5);
35
+ const $ = _c(3);
36
let t0;
37
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
38
t0 = {};
@@ -41,8 +41,9 @@ function Component(props) {
41
t0 = $[0];
42
}
43
const x = t0;
44
- let y;
44
+ let t1;
45
if ($[1] !== props) {
46
+ let y;
47
if (props.cond) {
48
y = {};
49
} else {
@@ -50,18 +51,12 @@ function Component(props) {
51
}
52
53
y.x = x;
53
- $[1] = props;
54
- $[2] = y;
55
- } else {
56
- y = $[2];
57
- }
58
- let t1;
59
- if ($[3] !== y) {
54
+
55
t1 = [x, y];
61
- $[3] = y;
62
- $[4] = t1;
56
+ $[1] = props;
57
+ $[2] = t1;
58
} else {
64
- t1 = $[4];
59
+ t1 = $[2];
60
}
61
return t1;
62
}