Reorder InlineUseMemo after type inference
The goal of this stack is to generalize `InlineUseMemo` into a pass that inlines all immediately invoked function expressions (IIFEs). Rather than specialize just useMemo calls, we'll rely on DropManualMemoization running first and turning useMemo calls into IIFEs. Then the generalized inlining pass can handle those IIFEs as well as others present in the source. For now, moving the order of the pass makes the output closer to what it will eventually be after this stack is complete.
Joe Savona committed
Oct 9, 2023 at 15:25 UTC
444a1ad9c111b5148c343d2a79ff80f1076eaa29
20 files changed
+174
-166
compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+3
-3
@@ -113,9 +113,6 @@ function* runWithEnvironment(
113
pruneMaybeThrows(hir);
114
yield log({ kind: "hir", name: "PruneMaybeThrows", value: hir });
115
116
- inlineUseMemo(hir);
117
- yield log({ kind: "hir", name: "RewriteUseMemo", value: hir });
118
-
116
mergeConsecutiveBlocks(hir);
117
yield log({ kind: "hir", name: "MergeConsecutiveBlocks", value: hir });
118
@@ -146,6 +143,9 @@ function* runWithEnvironment(
143
});
144
}
145
146
+ inlineUseMemo(hir);
147
+ yield log({ kind: "hir", name: "InlineUseMemo", value: hir });
148
+
149
dropManualMemoization(hir);
150
yield log({ kind: "hir", name: "DropManualMemoization", value: hir });
151
compiler/packages/babel-plugin-react-forget/src/Inference/InlineUseMemo.ts
+79
-83
@@ -20,6 +20,7 @@ import {
20
InstructionKind,
21
LabelTerminal,
22
Place,
23
+ getHookKind,
24
makeInstructionId,
25
makeType,
26
reversePostorderBlocks,
@@ -61,8 +62,6 @@ import { retainWhere } from "../Utils/utils";
62
export function inlineUseMemo(fn: HIRFunction): void {
63
// Track all function expressions in case they appear as the argument to a useMemo
64
const functions = new Map<IdentifierId, FunctionExpression>();
64
- // Track all references to `useMemo`
65
- const useMemoGlobals = new Set<IdentifierId>();
65
// Identifiers (lvalues) for known useMemo functions, so that we can prune them
66
// at the end of the pass
67
const useMemoFunctions = new Set<IdentifierId>();
@@ -77,103 +76,100 @@ export function inlineUseMemo(fn: HIRFunction): void {
76
for (let ii = 0; ii < block.instructions.length; ii++) {
77
const instr = block.instructions[ii]!;
78
switch (instr.value.kind) {
80
- case "LoadGlobal": {
81
- if (instr.value.name === "useMemo") {
82
- useMemoGlobals.add(instr.lvalue.identifier.id);
83
- }
84
- break;
85
- }
79
case "FunctionExpression": {
80
functions.set(instr.lvalue.identifier.id, instr.value);
81
break;
82
}
83
+ case "MethodCall":
84
case "CallExpression": {
91
- if (useMemoGlobals.has(instr.value.callee.identifier.id)) {
92
- const [lambda] = instr.value.args;
93
- if (lambda.kind === "Spread") {
94
- continue;
95
- }
96
- const body = functions.get(lambda.identifier.id);
97
- if (body === undefined) {
98
- // Allow passing a named function to useMemo, eg `useMemo(someImportedFunction, [])`
99
- continue;
100
- }
101
-
102
- if (body.loweredFunc.func.params.length > 0) {
103
- CompilerError.invalidReact({
104
- reason: "useMemo callbacks may not accept any arguments",
105
- description: null,
106
- loc: body.loc,
107
- suggestions: null,
108
- });
109
- }
85
+ const hookKind =
86
+ instr.value.kind === "CallExpression"
87
+ ? getHookKind(fn.env, instr.value.callee.identifier)
88
+ : getHookKind(fn.env, instr.value.property.identifier);
89
+ if (hookKind !== "useMemo") {
90
+ continue;
91
+ }
92
+ const [lambda] = instr.value.args;
93
+ if (lambda.kind === "Spread") {
94
+ continue;
95
+ }
96
+ const body = functions.get(lambda.identifier.id);
97
+ if (body === undefined) {
98
+ // Allow passing a named function to useMemo, eg `useMemo(someImportedFunction, [])`
99
+ continue;
100
+ }
101
111
- if (
112
- body.loweredFunc.func.async ||
113
- body.loweredFunc.func.generator
114
- ) {
115
- CompilerError.invalidReact({
116
- reason:
117
- "useMemo callbacks may not be async or generator functions",
118
- description: null,
119
- loc: body.loc,
120
- suggestions: null,
121
- });
122
- }
102
+ if (body.loweredFunc.func.params.length > 0) {
103
+ CompilerError.invalidReact({
104
+ reason: "useMemo callbacks may not accept any arguments",
105
+ description: null,
106
+ loc: body.loc,
107
+ suggestions: null,
108
+ });
109
+ }
110
124
- // We know this function is used for useMemo and can prune it later
125
- useMemoFunctions.add(lambda.identifier.id);
111
+ if (body.loweredFunc.func.async || body.loweredFunc.func.generator) {
112
+ CompilerError.invalidReact({
113
+ reason:
114
+ "useMemo callbacks may not be async or generator functions",
115
+ description: null,
116
+ loc: body.loc,
117
+ suggestions: null,
118
+ });
119
+ }
120
127
- // Create a new block which will contain code following the useMemo call
128
- const continuationBlockId = fn.env.nextBlockId;
129
- const continuationBlock: BasicBlock = {
130
- id: continuationBlockId,
131
- instructions: block.instructions.slice(ii + 1),
132
- kind: block.kind,
133
- phis: new Set(),
134
- preds: new Set(),
135
- terminal: block.terminal,
136
- };
137
- fn.body.blocks.set(continuationBlockId, continuationBlock);
121
+ // We know this function is used for useMemo and can prune it later
122
+ useMemoFunctions.add(lambda.identifier.id);
123
139
- // Trim the original block to contain instructions up to (but not including)
140
- // the useMemo
141
- block.instructions.length = ii;
124
+ // Create a new block which will contain code following the useMemo call
125
+ const continuationBlockId = fn.env.nextBlockId;
126
+ const continuationBlock: BasicBlock = {
127
+ id: continuationBlockId,
128
+ instructions: block.instructions.slice(ii + 1),
129
+ kind: block.kind,
130
+ phis: new Set(),
131
+ preds: new Set(),
132
+ terminal: block.terminal,
133
+ };
134
+ fn.body.blocks.set(continuationBlockId, continuationBlock);
135
143
- // To account for complex control flow within the lambda, we treat the lambda
144
- // as if it were a single labeled statement, and replace all returns with gotos
145
- // to the label fallthrough.
146
- const newTerminal: LabelTerminal = {
147
- block: body.loweredFunc.func.body.entry,
148
- id: makeInstructionId(0),
149
- kind: "label",
150
- fallthrough: continuationBlockId,
151
- loc: block.terminal.loc,
152
- };
153
- block.terminal = newTerminal;
136
+ // Trim the original block to contain instructions up to (but not including)
137
+ // the useMemo
138
+ block.instructions.length = ii;
139
155
- // We store the result in the useMemo temporary
156
- const result = instr.lvalue;
140
+ // To account for complex control flow within the lambda, we treat the lambda
141
+ // as if it were a single labeled statement, and replace all returns with gotos
142
+ // to the label fallthrough.
143
+ const newTerminal: LabelTerminal = {
144
+ block: body.loweredFunc.func.body.entry,
145
+ id: makeInstructionId(0),
146
+ kind: "label",
147
+ fallthrough: continuationBlockId,
148
+ loc: block.terminal.loc,
149
+ };
150
+ block.terminal = newTerminal;
151
158
- // Declare the useMemo temporary
159
- declareTemporary(fn.env, block, result);
152
+ // We store the result in the useMemo temporary
153
+ const result = instr.lvalue;
154
161
- // Promote the temporary with a name as we require this to persist
162
- promoteTemporary(result.identifier);
155
+ // Declare the useMemo temporary
156
+ declareTemporary(fn.env, block, result);
157
164
- // Rewrite blocks from the lambda to replace any `return` with a
165
- // store to the result and `goto` the continuation block
166
- for (const [id, block] of body.loweredFunc.func.body.blocks) {
167
- block.preds.clear();
168
- rewriteBlock(fn.env, block, continuationBlockId, result);
169
- fn.body.blocks.set(id, block);
170
- }
158
+ // Promote the temporary with a name as we require this to persist
159
+ promoteTemporary(result.identifier);
160
172
- // Ensure we visit the continuation block, since there may have been
173
- // sequential useMemos that need to be visited.
174
- queue.push(continuationBlock);
175
- continue queue;
161
+ // Rewrite blocks from the lambda to replace any `return` with a
162
+ // store to the result and `goto` the continuation block
163
+ for (const [id, block] of body.loweredFunc.func.body.blocks) {
164
+ block.preds.clear();
165
+ rewriteBlock(fn.env, block, continuationBlockId, result);
166
+ fn.body.blocks.set(id, block);
167
}
168
+
169
+ // Ensure we visit the continuation block, since there may have been
170
+ // sequential useMemos that need to be visited.
171
+ queue.push(continuationBlock);
172
+ continue queue;
173
}
174
}
175
}
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/babel-existing-react-import.expect.md
+6
-4
@@ -32,6 +32,7 @@ import {
32
function Component(props) {
33
const $ = useMemoCache(4);
34
const [x] = useState(0);
35
+ let t35;
36
const c_0 = $[0] !== x;
37
let t0;
38
if (c_0) {
@@ -41,8 +42,8 @@ function Component(props) {
42
} else {
43
t0 = $[1];
44
}
44
- const t15 = t0;
45
- const expensiveNumber = t15;
45
+ t35 = t0;
46
+ const expensiveNumber = t35;
47
const c_2 = $[2] !== expensiveNumber;
48
let t1;
49
if (c_2) {
@@ -58,6 +59,7 @@ function Component(props) {
59
function Component2(props) {
60
const $ = useMemoCache(4);
61
const [x] = useState(0);
62
+ let t35;
63
const c_0 = $[0] !== x;
64
let t0;
65
if (c_0) {
@@ -67,8 +69,8 @@ function Component2(props) {
69
} else {
70
t0 = $[1];
71
}
70
- const t15 = t0;
71
- const expensiveNumber = t15;
72
+ t35 = t0;
73
+ const expensiveNumber = t35;
74
const c_2 = $[2] !== expensiveNumber;
75
let t1;
76
if (c_2) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/babel-existing-react-kitchensink-import.expect.md
+6
-4
@@ -34,6 +34,7 @@ import {
34
function Component(props) {
35
const $ = useMemoCache(4);
36
const [x] = useState(0);
37
+ let t35;
38
const c_0 = $[0] !== x;
39
let t0;
40
if (c_0) {
@@ -43,8 +44,8 @@ function Component(props) {
44
} else {
45
t0 = $[1];
46
}
46
- const t15 = t0;
47
- const expensiveNumber = t15;
47
+ t35 = t0;
48
+ const expensiveNumber = t35;
49
const c_2 = $[2] !== expensiveNumber;
50
let t1;
51
if (c_2) {
@@ -60,6 +61,7 @@ function Component(props) {
61
function Component2(props) {
62
const $ = useMemoCache(4);
63
const [x] = useState(0);
64
+ let t35;
65
const c_0 = $[0] !== x;
66
let t0;
67
if (c_0) {
@@ -69,8 +71,8 @@ function Component2(props) {
71
} else {
72
t0 = $[1];
73
}
72
- const t15 = t0;
73
- const expensiveNumber = t15;
74
+ t35 = t0;
75
+ const expensiveNumber = t35;
76
const c_2 = $[2] !== expensiveNumber;
77
let t1;
78
if (c_2) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/babel-existing-react-namespace-import.expect.md
+3
-1
@@ -29,7 +29,9 @@ import { calculateExpensiveNumber } from "shared-runtime";
29
function Component(props) {
30
const $ = useMemoCache(2);
31
const [x] = React.useState(0);
32
- const expensiveNumber = (() => calculateExpensiveNumber(x))();
32
+ let t39;
33
+ t39 = calculateExpensiveNumber(x);
34
+ const expensiveNumber = t39;
35
const c_0 = $[0] !== expensiveNumber;
36
let t0;
37
if (c_0) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/bug.useMemo-deps-array-not-cleared.expect.md
+5
-4
@@ -24,11 +24,12 @@ export const FIXTURE_ENTRYPOINT = {
24
25
```javascript
26
import { unstable_useMemoCache as useMemoCache } from "react";
27
-function App(t25) {
27
+function App(t23) {
28
const $ = useMemoCache(2);
29
- const { text, hasDeps } = t25;
29
+ const { text, hasDeps } = t23;
30
31
hasDeps ? null : [text];
32
+ let t44;
33
const c_0 = $[0] !== text;
34
let t0;
35
if (c_0) {
@@ -38,8 +39,8 @@ function App(t25) {
39
} else {
40
t0 = $[1];
41
}
41
- const t18 = t0;
42
- const resolvedText = t18;
42
+ t44 = t0;
43
+ const resolvedText = t44;
44
return resolvedText;
45
}
46
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/drop-methodcall-usememo.expect.md
+9
-19
@@ -27,30 +27,20 @@ import { unstable_useMemoCache as useMemoCache } from "react";
27
import * as React from "react";
28
29
function Component(props) {
30
- const $ = useMemoCache(4);
30
+ const $ = useMemoCache(2);
31
+ let t42;
32
const c_0 = $[0] !== props.value;
32
- let t0;
33
+ let x;
34
if (c_0) {
34
- t0 = () => {
35
- const x = [];
36
- x.push(props.value);
37
- return x;
38
- };
35
+ x = [];
36
+ x.push(props.value);
37
$[0] = props.value;
40
- $[1] = t0;
41
- } else {
42
- t0 = $[1];
43
- }
44
- const c_2 = $[2] !== t0;
45
- let t1;
46
- if (c_2) {
47
- t1 = t0();
48
- $[2] = t0;
49
- $[3] = t1;
38
+ $[1] = x;
39
} else {
51
- t1 = $[3];
40
+ x = $[1];
41
}
53
- const x_0 = t1;
42
+ t42 = x;
43
+ const x_0 = t42;
44
return x_0;
45
}
46
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/remove-memoization-kitchen-sink.expect.md
+3
-2
@@ -28,10 +28,11 @@ function Component(props) {
28
};
29
30
const object = { x, onChange };
31
+ let t86;
32
33
const { x: x_0, onChange: onChange_0 } = object;
33
- const t43 = <input value={x_0} onChange={onChange_0} />;
34
- return t43;
34
+ t86 = <input value={x_0} onChange={onChange_0} />;
35
+ return t86;
36
}
37
38
```
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.expect.md
+6
-6
@@ -20,8 +20,8 @@ function Component(props) {
20
import { unstable_useMemoCache as useMemoCache } from "react";
21
function Component(props) {
22
const $ = useMemoCache(4);
23
- let t20 = undefined;
24
- bb7: {
23
+ let t44;
24
+ bb8: {
25
if (props.cond) {
26
const c_0 = $[0] !== props.a;
27
let t0;
@@ -32,8 +32,8 @@ function Component(props) {
32
} else {
33
t0 = $[1];
34
}
35
- t20 = t0;
36
- break bb7;
35
+ t44 = t0;
36
+ break bb8;
37
}
38
const c_2 = $[2] !== props.b;
39
let t1;
@@ -44,9 +44,9 @@ function Component(props) {
44
} else {
45
t1 = $[3];
46
}
47
- t20 = t1;
47
+ t44 = t1;
48
}
49
- const x = t20;
49
+ const x = t44;
50
return x;
51
}
52
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-independently-memoizeable.expect.md
+3
-2
@@ -20,6 +20,7 @@ function Component(props) {
20
import { unstable_useMemoCache as useMemoCache } from "react";
21
function Component(props) {
22
const $ = useMemoCache(10);
23
+ let t59;
24
const c_0 = $[0] !== props.a;
25
let t0;
26
if (c_0) {
@@ -51,8 +52,8 @@ function Component(props) {
52
} else {
53
t2 = $[6];
54
}
54
- const t26 = t2;
55
- const [a_0, b_0] = t26;
55
+ t59 = t2;
56
+ const [a_0, b_0] = t59;
57
const c_7 = $[7] !== a_0;
58
const c_8 = $[8] !== b_0;
59
let t3;
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-inlining-block-return.expect.md
+6
-6
@@ -25,8 +25,8 @@ export const FIXTURE_ENTRYPOINT = {
25
import { unstable_useMemoCache as useMemoCache } from "react";
26
function component(a, b) {
27
const $ = useMemoCache(2);
28
- let t13 = undefined;
29
- bb6: {
28
+ let t31;
29
+ bb7: {
30
if (a) {
31
const c_0 = $[0] !== b;
32
let t0;
@@ -37,12 +37,12 @@ function component(a, b) {
37
} else {
38
t0 = $[1];
39
}
40
- t13 = t0;
41
- break bb6;
40
+ t31 = t0;
41
+ break bb7;
42
}
43
- t13 = undefined;
43
+ t31 = undefined;
44
}
45
- const x = t13;
45
+ const x = t31;
46
return x;
47
}
48
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-inverted-if.expect.md
+6
-6
@@ -27,20 +27,20 @@ export const FIXTURE_ENTRYPOINT = {
27
28
```javascript
29
function Component(props) {
30
- let t16 = undefined;
31
- bb10: {
30
+ let t36;
31
+ bb11: {
32
bb5: {
33
if (props.cond) {
34
break bb5;
35
}
36
37
- t16 = props.a;
38
- break bb10;
37
+ t36 = props.a;
38
+ break bb11;
39
}
40
41
- t16 = props.b;
41
+ t36 = props.b;
42
}
43
- const x = t16;
43
+ const x = t36;
44
return x;
45
}
46
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.expect.md
+4
-2
@@ -23,8 +23,10 @@ export const FIXTURE_ENTRYPOINT = {
23
24
```javascript
25
function Component(props) {
26
- const t8 = props.value;
27
- const x = t8;
26
+ let t20;
27
+
28
+ t20 = props.value;
29
+ const x = t20;
30
return x;
31
}
32
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-logical.expect.md
+3
-2
@@ -19,8 +19,9 @@ export const FIXTURE_ENTRYPOINT = {
19
20
```javascript
21
function Component(props) {
22
- const t16 = props.a && props.b;
23
- const x = t16;
22
+ let t38;
23
+ t38 = props.a && props.b;
24
+ const x = t38;
25
return x;
26
}
27
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md
+8
-8
@@ -33,8 +33,8 @@ import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
33
34
function Component(props) {
35
const $ = useMemoCache(3);
36
- let t31 = undefined;
37
- bb9: {
36
+ let t68;
37
+ bb10: {
38
const c_0 = $[0] !== props;
39
let y;
40
if (c_0) {
@@ -43,21 +43,21 @@ function Component(props) {
43
y.push(props.a);
44
}
45
if (props.cond2) {
46
- t31 = y;
47
- break bb9;
46
+ t68 = y;
47
+ break bb10;
48
}
49
50
y.push(props.b);
51
$[0] = props;
52
$[1] = y;
53
- $[2] = t31;
53
+ $[2] = t68;
54
} else {
55
y = $[1];
56
- t31 = $[2];
56
+ t68 = $[2];
57
}
58
- t31 = y;
58
+ t68 = y;
59
}
60
- const x = t31;
60
+ const x = t68;
61
return x;
62
}
63
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-nested-ifs.expect.md
+4
@@ -24,10 +24,14 @@ export const FIXTURE_ENTRYPOINT = {
24
25
```javascript
26
function Component(props) {
27
+ let t31;
28
if (props.cond) {
29
if (props.cond) {
30
}
31
}
32
+ t31 = undefined;
33
+ const x = t31;
34
+ return x;
35
}
36
37
export const FIXTURE_ENTRYPOINT = {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-return-empty.expect.md
+5
@@ -15,7 +15,12 @@ function component(a) {
15
16
```javascript
17
function component(a) {
18
+ let t23;
19
+
20
mutate(a);
21
+ t23 = undefined;
22
+ const x = t23;
23
+ return x;
24
}
25
26
```
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md
+3
-2
@@ -15,6 +15,7 @@ function component(a) {
15
import { unstable_useMemoCache as useMemoCache } from "react";
16
function component(a) {
17
const $ = useMemoCache(4);
18
+ let t24;
19
const c_0 = $[0] !== a;
20
let t0;
21
if (c_0) {
@@ -24,8 +25,8 @@ function component(a) {
25
} else {
26
t0 = $[1];
27
}
27
- const t9 = t0;
28
- const x = t9;
28
+ t24 = t0;
29
+ const x = t24;
30
const c_2 = $[2] !== x;
31
let t1;
32
if (c_2) {
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.expect.md
+6
-6
@@ -28,17 +28,17 @@ export const FIXTURE_ENTRYPOINT = {
28
29
```javascript
30
function Component(props) {
31
- let t17 = undefined;
32
- bb8: switch (props.key) {
31
+ let t38;
32
+ bb9: switch (props.key) {
33
case "key": {
34
- t17 = props.value;
35
- break bb8;
34
+ t38 = props.value;
35
+ break bb9;
36
}
37
default: {
38
- t17 = props.defaultValue;
38
+ t38 = props.defaultValue;
39
}
40
}
41
- const x = t17;
41
+ const x = t38;
42
return x;
43
}
44
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-switch-return.expect.md
+6
-6
@@ -34,13 +34,13 @@ export const FIXTURE_ENTRYPOINT = {
34
35
```javascript
36
function Component(props) {
37
- let t21 = undefined;
38
- bb10: {
37
+ let t49;
38
+ bb11: {
39
let y = undefined;
40
bb2: switch (props.switch) {
41
case "foo": {
42
- t21 = "foo";
43
- break bb10;
42
+ t49 = "foo";
43
+ break bb11;
44
}
45
case "bar": {
46
y = "bar";
@@ -51,9 +51,9 @@ function Component(props) {
51
}
52
}
53
54
- t21 = y;
54
+ t49 = y;
55
}
56
- const x = t21;
56
+ const x = t49;
57
return x;
58
}
59