[compiler] Infer return types of function expressions
Uses the returnIdentifier added in the previous PR to provide a stable identifier for which we can infer a return type for functions, then wires up the equations in InferTypes to infer the type. ghstack-source-id: 22c0a9ea096daa5f72821fca2a5ff5b199f65c8b Pull Request resolved: https://github.com/facebook/react/pull/30785
Joe Savona committed
Aug 21, 2024 at 21:17 UTC
8410c8b959b8e20adc5577cb7211702cfba0f78f
10 files changed
+74
-50
compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts
+3
-1
@@ -211,7 +211,9 @@ export function lower(
211
null,
212
);
213
214
- const returnIdentifier = builder.makeTemporary(func.node.loc ?? GeneratedSource);
214
+ const returnIdentifier = builder.makeTemporary(
215
+ func.node.loc ?? GeneratedSource,
216
+ );
217
218
return Ok({
219
id,
compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts
+5
-1
@@ -72,6 +72,7 @@ export function printFunction(fn: HIRFunction): string {
72
if (definition.length !== 0) {
73
output.push(definition);
74
}
75
+ output.push(printType(fn.returnIdentifier.type));
76
output.push(printHIR(fn.body));
77
output.push(...fn.directives);
78
return output.join('\n');
@@ -555,7 +556,10 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
556
}
557
})
558
.join(', ') ?? '';
558
- value = `${kind} ${name} @deps[${deps}] @context[${context}] @effects[${effects}]:\n${fn}`;
559
+ const type = printType(
560
+ instrValue.loweredFunc.func.returnIdentifier.type,
561
+ ).trim();
562
+ value = `${kind} ${name} @deps[${deps}] @context[${context}] @effects[${effects}]${type !== '' ? ` return${type}` : ''}:\n${fn}`;
563
break;
564
}
565
case 'TaggedTemplateExpression': {
compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts
+5
-2
@@ -238,7 +238,10 @@ function emitSelectorFn(env: Environment, keys: Array<string>): Instruction {
238
phis: new Set(),
239
};
240
241
- const returnIdentifier = createTemporaryPlace(env, GeneratedSource).identifier;
241
+ const returnIdentifier = createTemporaryPlace(
242
+ env,
243
+ GeneratedSource,
244
+ ).identifier;
245
const fn: HIRFunction = {
246
loc: GeneratedSource,
247
id: null,
@@ -246,7 +249,7 @@ function emitSelectorFn(env: Environment, keys: Array<string>): Instruction {
249
env,
250
params: [obj],
251
returnType: null,
249
- returnIdentifier,
252
+ returnIdentifier,
253
context: [],
254
effects: null,
255
body: {
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts
+13
-7
@@ -481,14 +481,20 @@ function canMergeScopes(
481
}
482
483
function isAlwaysInvalidatingType(type: Type): boolean {
484
- if (type.kind === 'Object') {
485
- switch (type.shapeId) {
486
- case BuiltInArrayId:
487
- case BuiltInObjectId:
488
- case BuiltInFunctionId:
489
- case BuiltInJsxId: {
490
- return true;
484
+ switch (type.kind) {
485
+ case 'Object': {
486
+ switch (type.shapeId) {
487
+ case BuiltInArrayId:
488
+ case BuiltInObjectId:
489
+ case BuiltInFunctionId:
490
+ case BuiltInJsxId: {
491
+ return true;
492
+ }
493
}
494
+ break;
495
+ }
496
+ case 'Function': {
497
+ return true;
498
}
499
}
500
return false;
compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts
+19
-1
@@ -88,6 +88,7 @@ function apply(func: HIRFunction, unifier: Unifier): void {
88
}
89
}
90
}
91
+ func.returnIdentifier.type = unifier.get(func.returnIdentifier.type);
92
}
93
94
type TypeEquation = {
@@ -122,6 +123,7 @@ function* generate(
123
}
124
125
const names = new Map();
126
+ const returnTypes: Array<Type> = [];
127
for (const [_, block] of func.body.blocks) {
128
for (const phi of block.phis) {
129
yield equation(phi.type, {
@@ -133,6 +135,18 @@ function* generate(
135
for (const instr of block.instructions) {
136
yield* generateInstructionTypes(func.env, names, instr);
137
}
138
+ const terminal = block.terminal;
139
+ if (terminal.kind === 'return') {
140
+ returnTypes.push(terminal.value.identifier.type);
141
+ }
142
+ }
143
+ if (returnTypes.length > 1) {
144
+ yield equation(func.returnIdentifier.type, {
145
+ kind: 'Phi',
146
+ operands: returnTypes,
147
+ });
148
+ } else if (returnTypes.length === 1) {
149
+ yield equation(func.returnIdentifier.type, returnTypes[0]!);
150
}
151
}
152
@@ -346,7 +360,11 @@ function* generateInstructionTypes(
360
361
case 'FunctionExpression': {
362
yield* generate(value.loweredFunc.func);
349
- yield equation(left, {kind: 'Object', shapeId: BuiltInFunctionId});
363
+ yield equation(left, {
364
+ kind: 'Function',
365
+ shapeId: BuiltInFunctionId,
366
+ return: value.loweredFunc.func.returnIdentifier.type,
367
+ });
368
break;
369
}
370
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-repro-missed-memoization-from-capture-in-invoked-function-inferred-as-mutation.expect.md
+1
-1
@@ -13,7 +13,7 @@ component Component() {
13
if (data != null) {
14
return true;
15
} else {
16
- return false;
16
+ return {};
17
}
18
};
19
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-repro-missed-memoization-from-capture-in-invoked-function-inferred-as-mutation.js
+1
-1
@@ -9,7 +9,7 @@ component Component() {
9
if (data != null) {
10
return true;
11
} else {
12
- return false;
12
+ return {};
13
}
14
};
15
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-simple-const-declaration.expect.md
+5
-6
@@ -25,18 +25,17 @@ export const FIXTURE_ENTRYPOINT = {
25
import { c as _c } from "react/compiler-runtime";
26
function hoisting() {
27
const $ = _c(1);
28
- let t0;
28
+ let foo;
29
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
30
- const foo = () => bar + baz;
30
+ foo = () => bar + baz;
31
32
const bar = 3;
33
const baz = 2;
34
- t0 = foo();
35
- $[0] = t0;
34
+ $[0] = foo;
35
} else {
37
- t0 = $[0];
36
+ foo = $[0];
37
}
39
- return t0;
38
+ return foo();
39
}
40
41
export const FIXTURE_ENTRYPOINT = {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hoisting-simple-let-declaration.expect.md
+5
-6
@@ -25,18 +25,17 @@ export const FIXTURE_ENTRYPOINT = {
25
import { c as _c } from "react/compiler-runtime";
26
function hoisting() {
27
const $ = _c(1);
28
- let t0;
28
+ let foo;
29
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
30
- const foo = () => bar + baz;
30
+ foo = () => bar + baz;
31
32
let bar = 3;
33
let baz = 2;
34
- t0 = foo();
35
- $[0] = t0;
34
+ $[0] = foo;
35
} else {
37
- t0 = $[0];
36
+ foo = $[0];
37
}
39
- return t0;
38
+ return foo();
39
}
40
41
export const FIXTURE_ENTRYPOINT = {
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/useCallback-nonescaping-invoked-callback-escaping-return.expect.md
+17
-24
@@ -40,7 +40,7 @@ import { c as _c } from "react/compiler-runtime"; // @validatePreserveExistingMe
40
import { useCallback } from "react";
41
42
function Component(t0) {
43
- const $ = _c(11);
43
+ const $ = _c(9);
44
const { entity, children } = t0;
45
let t1;
46
if ($[0] !== entity) {
@@ -51,46 +51,39 @@ function Component(t0) {
51
t1 = $[1];
52
}
53
const showMessage = t1;
54
+
55
+ const shouldShowMessage = showMessage();
56
let t2;
55
- if ($[2] !== showMessage) {
56
- t2 = showMessage();
57
- $[2] = showMessage;
57
+ if ($[2] !== shouldShowMessage) {
58
+ t2 = <div>{shouldShowMessage}</div>;
59
+ $[2] = shouldShowMessage;
60
$[3] = t2;
61
} else {
62
t2 = $[3];
63
}
62
- const shouldShowMessage = t2;
64
let t3;
64
- if ($[4] !== shouldShowMessage) {
65
- t3 = <div>{shouldShowMessage}</div>;
66
- $[4] = shouldShowMessage;
65
+ if ($[4] !== children) {
66
+ t3 = <div>{children}</div>;
67
+ $[4] = children;
68
$[5] = t3;
69
} else {
70
t3 = $[5];
71
}
72
let t4;
72
- if ($[6] !== children) {
73
- t4 = <div>{children}</div>;
74
- $[6] = children;
75
- $[7] = t4;
76
- } else {
77
- t4 = $[7];
78
- }
79
- let t5;
80
- if ($[8] !== t3 || $[9] !== t4) {
81
- t5 = (
73
+ if ($[6] !== t2 || $[7] !== t3) {
74
+ t4 = (
75
<div>
76
+ {t2}
77
{t3}
84
- {t4}
78
</div>
79
);
87
- $[8] = t3;
88
- $[9] = t4;
89
- $[10] = t5;
80
+ $[6] = t2;
81
+ $[7] = t3;
82
+ $[8] = t4;
83
} else {
91
- t5 = $[10];
84
+ t4 = $[8];
85
}
93
- return t5;
86
+ return t4;
87
}
88
89
export const FIXTURE_ENTRYPOINT = {