[hir] Refactor useMemo inlining
Previously useMemo inlining created a new StoreLocal assignment (not reassignment!) instruction for every return value. This breaks when the return is inside a block (like an if-block) as the scope is tied to the block. For example: ``` let x = useMemo(() => { if (...) { return { ... }; } }) ``` would become: ``` if (...) { const temp = { ... }; } const x = temp; ``` This PR instead changes the inlining to declare a temporary in the function prologue and then reassign values to it when replacing return statements. ``` let x = useMemo(() => { if (...) { return { ... }; } }) ``` becomes ``` let temp; if (...) { temp = { ... }; } const x = temp; ```
Sathya Gunasekaran committed
Apr 17, 2023 at 20:59 UTC
00d4a69459ce57afe10b24331d10e8594e09378d
9 files changed
+87
-46
compiler/forget/src/Inference/InlineUseMemo.ts
+52
-3
@@ -12,11 +12,13 @@ import {
12
Effect,
13
Environment,
14
FunctionExpression,
15
+ GeneratedSource,
16
GotoTerminal,
17
GotoVariant,
18
HIR,
19
HIRFunction,
20
IdentifierId,
21
+ Identifier,
22
InstructionKind,
23
makeInstructionId,
24
makeType,
@@ -198,11 +200,20 @@ export function inlineUseMemo(fn: HIRFunction): void {
200
}
201
}
202
203
+ // We store the result in the useMemo temporary
204
+ const result = instr.lvalue;
205
+
206
+ // Declare the useMemo temporary
207
+ declareTemporary(fn.env, block, result);
208
+
209
+ // Promote the temporary with a name as we require this to persist
210
+ promoteTemporary(result.identifier);
211
+
212
// Rewrite blocks from the lambda to replace any `return` with a
202
- // store the useMemo temporary and `goto` the continuation block
213
+ // store to the result and `goto` the continuation block
214
for (const [id, block] of body.loweredFunc.body.blocks) {
215
block.preds.clear();
205
- rewriteBlock(fn.env, block, continuationBlockId, instr.lvalue);
216
+ rewriteBlock(fn.env, block, continuationBlockId, result);
217
fn.body.blocks.set(id, block);
218
}
219
@@ -342,7 +353,7 @@ function rewriteBlock(
353
},
354
value: {
355
kind: "StoreLocal",
345
- lvalue: { kind: InstructionKind.Const, place: { ...returnValue } },
356
+ lvalue: { kind: InstructionKind.Reassign, place: { ...returnValue } },
357
value: terminal.value,
358
loc: terminal.loc,
359
},
@@ -356,3 +367,41 @@ function rewriteBlock(
367
loc: block.terminal.loc,
368
};
369
}
370
+
371
+function declareTemporary(
372
+ env: Environment,
373
+ block: BasicBlock,
374
+ result: Place
375
+): void {
376
+ block.instructions.push({
377
+ id: makeInstructionId(0),
378
+ loc: GeneratedSource,
379
+ lvalue: {
380
+ effect: Effect.Unknown,
381
+ identifier: {
382
+ id: env.nextIdentifierId,
383
+ mutableRange: {
384
+ start: makeInstructionId(0),
385
+ end: makeInstructionId(0),
386
+ },
387
+ name: null,
388
+ scope: null,
389
+ type: makeType(),
390
+ },
391
+ kind: "Identifier",
392
+ loc: GeneratedSource,
393
+ },
394
+ value: {
395
+ kind: "DeclareLocal",
396
+ lvalue: {
397
+ place: result,
398
+ kind: InstructionKind.Let,
399
+ },
400
+ loc: result.loc,
401
+ },
402
+ });
403
+}
404
+
405
+function promoteTemporary(temp: Identifier): void {
406
+ temp.name = `t${temp.id}`;
407
+}
compiler/forget/src/__tests__/fixtures/compiler/useMemo-if-else-multiple-return.expect.md
+12
-17
@@ -20,7 +20,8 @@ function Component(props) {
20
```javascript
21
// @inlineUseMemo
22
function Component(props) {
23
- const $ = React.unstable_useMemoCache(5);
23
+ const $ = React.unstable_useMemoCache(4);
24
+ let t16 = undefined;
25
if (props.cond) {
26
const c_0 = $[0] !== props.a;
27
let t0;
@@ -31,26 +32,20 @@ function Component(props) {
32
} else {
33
t0 = $[1];
34
}
34
- let t1;
35
- if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
36
- t1 = t0;
37
- $[2] = t1;
38
- } else {
39
- t1 = $[2];
40
- }
35
+ t16 = t0;
36
} else {
42
- const c_3 = $[3] !== props.b;
43
- let t2;
44
- if (c_3) {
45
- t2 = makeObject(props.b);
46
- $[3] = props.b;
47
- $[4] = t2;
37
+ const c_2 = $[2] !== props.b;
38
+ let t1;
39
+ if (c_2) {
40
+ t1 = makeObject(props.b);
41
+ $[2] = props.b;
42
+ $[3] = t1;
43
} else {
49
- t2 = $[4];
44
+ t1 = $[3];
45
}
51
- t1 = t2;
46
+ t16 = t1;
47
}
53
- const x = t1;
48
+ const x = t16;
49
return x;
50
}
51
compiler/forget/src/__tests__/fixtures/compiler/useMemo-independently-memoizeable.expect.md
+2
-2
@@ -52,8 +52,8 @@ function Component(props) {
52
} else {
53
t2 = $[6];
54
}
55
- const t54 = t2;
56
- const [a_0, b_0] = t54;
55
+ const t24 = t2;
56
+ const [a_0, b_0] = t24;
57
const c_7 = $[7] !== a_0;
58
const c_8 = $[8] !== b_0;
59
let t3;
compiler/forget/src/__tests__/fixtures/compiler/useMemo-inlining-block-return.expect.md
+4
-9
@@ -19,7 +19,8 @@ function component(a, b) {
19
```javascript
20
// @inlineUseMemo
21
function component(a, b) {
22
- const $ = React.unstable_useMemoCache(3);
22
+ const $ = React.unstable_useMemoCache(2);
23
+ let t13;
24
if (a) {
25
const c_0 = $[0] !== b;
26
let t0;
@@ -30,15 +31,9 @@ function component(a, b) {
31
} else {
32
t0 = $[1];
33
}
33
- let t1;
34
- if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
35
- t1 = t0;
36
- $[2] = t1;
37
- } else {
38
- t1 = $[2];
39
- }
34
+ t13 = t0;
35
}
41
- const x = t1;
36
+ const x = t13;
37
return x;
38
}
39
compiler/forget/src/__tests__/fixtures/compiler/useMemo-labeled-statement-unconditional-return.expect.md
+2
-2
@@ -19,8 +19,8 @@ function Component(props) {
19
```javascript
20
// @inlineUseMemo
21
function Component(props) {
22
- const t19 = props.value;
23
- const x = t19;
22
+ const t8 = props.value;
23
+ const x = t8;
24
return x;
25
}
26
compiler/forget/src/__tests__/fixtures/compiler/useMemo-logical.expect.md
+2
-2
@@ -15,8 +15,8 @@ function Component(props) {
15
```javascript
16
// @inlineUseMemo
17
function Component(props) {
18
- const t32 = props.a && props.b;
19
- const x = t32;
18
+ const t14 = props.a && props.b;
19
+ const x = t14;
20
return x;
21
}
22
compiler/forget/src/__tests__/fixtures/compiler/useMemo-multiple-if-else.expect.md
+7
-6
@@ -27,24 +27,25 @@ function Component(props) {
27
function Component(props) {
28
const $ = React.unstable_useMemoCache(2);
29
const c_0 = $[0] !== props;
30
- let t0;
30
+ let t25;
31
if (c_0) {
32
const y = [];
33
if (props.cond) {
34
y.push(props.a);
35
}
36
+ t25 = undefined;
37
if (props.cond2) {
37
- t0 = y;
38
+ t25 = y;
39
} else {
40
y.push(props.b);
40
- t0 = y;
41
+ t25 = y;
42
}
43
$[0] = props;
43
- $[1] = t0;
44
+ $[1] = t25;
45
} else {
45
- t0 = $[1];
46
+ t25 = $[1];
47
}
47
- const x = t0;
48
+ const x = t25;
49
return x;
50
}
51
compiler/forget/src/__tests__/fixtures/compiler/useMemo-simple.expect.md
+2
-2
@@ -25,8 +25,8 @@ function component(a) {
25
} else {
26
t0 = $[1];
27
}
28
- const t23 = t0;
29
- const x = t23;
28
+ const t9 = t0;
29
+ const x = t9;
30
const c_2 = $[2] !== x;
31
let t1;
32
if (c_2) {
compiler/forget/src/__tests__/fixtures/compiler/useMemo-switch-no-fallthrough.expect.md
+4
-3
@@ -24,16 +24,17 @@ function Component(props) {
24
```javascript
25
// @inlineUseMemo
26
function Component(props) {
27
+ let t13 = undefined;
28
bb8: switch (props.key) {
29
case "key": {
29
- const t28 = props.value;
30
+ t13 = props.value;
31
break bb8;
32
}
33
default: {
33
- const t28 = props.defaultValue;
34
+ t13 = props.defaultValue;
35
}
36
}
36
- const x = t28;
37
+ const x = t13;
38
return x;
39
}
40