enablePreserveMemo treats memo deps as frozen
See discussion on #2448 for full context. In the new `@enablePreserveExistingMemoizationGuarantees` mode, the goal is to preserve the existing referential equality guarantees from the original code. #2448 lays the groundwork by explicitly marking the _output_ of each useMemo block as memoized, hinting to the compiler that the value cannot subsequently change. This ensures the mutable range doesn't extend _later_, possibly overlapping a hook call and causing memoization to gett pruned. This PR fixes the other direction. There are cases where free variables referenced in the useMemo block could have been inferred as mutated, which could then extend the _start_ of the range earlier past a hook: ```javascript const foo = createObject(); useBar(); const baz = useMemo(() => { const baz = createObject(); maybeMutate(foo, baz); return baz; }, [foo]); ``` Here the compiler would infer that both `baz` and `foo` are mutable at the `maybeMutate()` call, grouping them in the same scope. But that scope would span the `useBar()` call, and be pruned, meaning that `baz` went unmemoized. However, useMemo blocks shouldn't be mutating free variables. Only variables newly created within the useMemo block should be mutable. So this PR extends the feature to treat all free variables referenced in a useMemo block as frozen as of the block itself.
Joe Savona committed
Dec 15, 2023 at 13:47 UTC
723b616c6778fa46fb817e2e4e000fdc7fd1b46d
8 files changed
+281
-8
compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts
+3
-1
@@ -494,6 +494,8 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
494
}
495
case "ObjectMethod":
496
case "FunctionExpression": {
497
+ const kind =
498
+ instrValue.kind === "FunctionExpression" ? "Function" : "ObjectMethod";
499
const name = getFunctionName(instrValue, "");
500
const fn = printFunction(instrValue.loweredFunc.func)
501
.split("\n")
@@ -505,7 +507,7 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
507
const context = instrValue.loweredFunc.func.context
508
.map((dep) => printPlace(dep))
509
.join(",");
508
- value = `Function ${name} @deps[${deps}] @context[${context}]:\n${fn}`;
510
+ value = `${kind} ${name} @deps[${deps}] @context[${context}]:\n${fn}`;
511
break;
512
}
513
case "TaggedTemplateExpression": {
compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts
+40
-6
@@ -8,16 +8,17 @@
8
import { CompilerError } from "..";
9
import {
10
Effect,
11
+ FunctionExpression,
12
HIRFunction,
13
IdentifierId,
14
Instruction,
15
Place,
16
SpreadPattern,
17
makeInstructionId,
17
- markInstructionIds,
18
} from "../HIR";
19
import { createTemporaryPlace } from "../HIR/HIRBuilder";
20
import { HookKind } from "../HIR/ObjectShape";
21
+import { eachInstructionValueOperand } from "../HIR/visitors";
22
23
/*
24
* Removes manual memoization using the `useMemo` and `useCallback` APIs. This pass is designed
@@ -28,6 +29,7 @@ import { HookKind } from "../HIR/ObjectShape";
29
* eg `React.useMemo()`.
30
*/
31
export function dropManualMemoization(func: HIRFunction): void {
32
+ const functions = new Map<IdentifierId, FunctionExpression>();
33
const hooks = new Map<IdentifierId, HookKind>();
34
const react = new Set<IdentifierId>();
35
let hasChanges = false;
@@ -35,10 +37,11 @@ export function dropManualMemoization(func: HIRFunction): void {
37
let nextInstructions: Array<Instruction> | null = null;
38
for (let i = 0; i < block.instructions.length; i++) {
39
const instr = block.instructions[i]!;
38
- if (nextInstructions !== null) {
39
- nextInstructions.push(instr);
40
- }
40
switch (instr.value.kind) {
41
+ case "FunctionExpression": {
42
+ functions.set(instr.lvalue.identifier.id, instr.value);
43
+ break;
44
+ }
45
case "LoadGlobal": {
46
if (
47
instr.value.name === "useMemo" ||
@@ -109,6 +112,7 @@ export function dropManualMemoization(func: HIRFunction): void {
112
args: [],
113
loc: instr.value.loc,
114
};
115
+
116
if (
117
func.env.config.enablePreserveExistingMemoizationGuarantees
118
) {
@@ -137,7 +141,29 @@ export function dropManualMemoization(func: HIRFunction): void {
141
const temp = createTemporaryPlace(func.env);
142
instr.lvalue = { ...temp };
143
nextInstructions =
140
- nextInstructions ?? block.instructions.slice(0, i + 1);
144
+ nextInstructions ?? block.instructions.slice(0, i);
145
+
146
+ const functionExpression = functions.get(fn.identifier.id);
147
+ if (functionExpression !== undefined) {
148
+ for (const operand of eachInstructionValueOperand(
149
+ functionExpression
150
+ )) {
151
+ const operandLValue = createTemporaryPlace(func.env);
152
+ nextInstructions.push({
153
+ id: makeInstructionId(0),
154
+ lvalue: operandLValue,
155
+ value: {
156
+ kind: "Memoize",
157
+ value: { ...operand },
158
+ loc: instr.loc,
159
+ },
160
+ loc: instr.loc,
161
+ });
162
+ }
163
+ }
164
+
165
+ nextInstructions.push(instr);
166
+
167
nextInstructions.push({
168
id: makeInstructionId(0),
169
lvalue,
@@ -148,7 +174,12 @@ export function dropManualMemoization(func: HIRFunction): void {
174
},
175
loc: instr.loc,
176
});
177
+ } else {
178
+ if (nextInstructions !== null) {
179
+ nextInstructions.push(instr);
180
+ }
181
}
182
+ continue;
183
}
184
} else if (hookKind === "useCallback") {
185
const [fn] = instr.value.args as Array<
@@ -229,6 +260,9 @@ export function dropManualMemoization(func: HIRFunction): void {
260
break;
261
}
262
}
263
+ if (nextInstructions !== null) {
264
+ nextInstructions.push(instr);
265
+ }
266
}
267
if (nextInstructions !== null) {
268
block.instructions = nextInstructions;
@@ -236,6 +270,6 @@ export function dropManualMemoization(func: HIRFunction): void {
270
}
271
}
272
if (hasChanges) {
239
- markInstructionIds(func.body);
273
+ // markInstructionIds(func.body);
274
}
275
}
compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts
+8
-1
@@ -313,7 +313,14 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
313
case "StoreContext": {
314
return false;
315
}
316
- case "Memoize":
316
+ case "Memoize": {
317
+ /**
318
+ * This instruction is used by the @enablePreserveExistingMemoizationGuarantees feature
319
+ * to preserve information about memoization semantics in the original code. We can't
320
+ * DCE without losing the memoization guarantees.
321
+ */
322
+ return false;
323
+ }
324
case "RegExpLiteral":
325
case "LoadGlobal":
326
case "ArrayExpression":
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.expect.md
new
+74
@@ -0,0 +1,74 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enablePreserveExistingMemoizationGuarantees:false
6
+import { useMemo } from "react";
7
+import {
8
+ identity,
9
+ makeObject_Primitives,
10
+ mutate,
11
+ useHook,
12
+} from "shared-runtime";
13
+
14
+function Component(props) {
15
+ // With the feature disabled these variables are inferred as being mutated inside the useMemo block
16
+ const free = makeObject_Primitives();
17
+ const free2 = makeObject_Primitives();
18
+ const part = free2.part;
19
+
20
+ // This causes their range to extend to include this hook call, and in turn for the memoization to be pruned
21
+ useHook();
22
+ const object = useMemo(() => {
23
+ const x = makeObject_Primitives();
24
+ x.value = props.value;
25
+ mutate(x, free, part);
26
+ return x;
27
+ }, [props.value]);
28
+ return object;
29
+}
30
+
31
+export const FIXTURE_ENTRYPOINT = {
32
+ fn: Component,
33
+ params: [{ value: 42 }],
34
+};
35
+
36
+```
37
+
38
+## Code
39
+
40
+```javascript
41
+// @enablePreserveExistingMemoizationGuarantees:false
42
+import { useMemo } from "react";
43
+import {
44
+ identity,
45
+ makeObject_Primitives,
46
+ mutate,
47
+ useHook,
48
+} from "shared-runtime";
49
+
50
+function Component(props) {
51
+ const free = makeObject_Primitives();
52
+ const free2 = makeObject_Primitives();
53
+ const part = free2.part;
54
+
55
+ useHook();
56
+ let t39;
57
+
58
+ const x = makeObject_Primitives();
59
+ x.value = props.value;
60
+ mutate(x, free, part);
61
+ t39 = x;
62
+ const object = t39;
63
+ return object;
64
+}
65
+
66
+export const FIXTURE_ENTRYPOINT = {
67
+ fn: Component,
68
+ params: [{ value: 42 }],
69
+};
70
+
71
+```
72
+
73
+### Eval output
74
+(kind: ok) {"a":0,"b":"value1","c":true,"value":42,"wat0":"joe"}
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.js
new
+30
@@ -0,0 +1,30 @@
1
+// @enablePreserveExistingMemoizationGuarantees:false
2
+import { useMemo } from "react";
3
+import {
4
+ identity,
5
+ makeObject_Primitives,
6
+ mutate,
7
+ useHook,
8
+} from "shared-runtime";
9
+
10
+function Component(props) {
11
+ // With the feature disabled these variables are inferred as being mutated inside the useMemo block
12
+ const free = makeObject_Primitives();
13
+ const free2 = makeObject_Primitives();
14
+ const part = free2.part;
15
+
16
+ // This causes their range to extend to include this hook call, and in turn for the memoization to be pruned
17
+ useHook();
18
+ const object = useMemo(() => {
19
+ const x = makeObject_Primitives();
20
+ x.value = props.value;
21
+ mutate(x, free, part);
22
+ return x;
23
+ }, [props.value]);
24
+ return object;
25
+}
26
+
27
+export const FIXTURE_ENTRYPOINT = {
28
+ fn: Component,
29
+ params: [{ value: 42 }],
30
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.expect.md
new
+95
@@ -0,0 +1,95 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enablePreserveExistingMemoizationGuarantees
6
+import { useMemo } from "react";
7
+import {
8
+ identity,
9
+ makeObject_Primitives,
10
+ mutate,
11
+ useHook,
12
+} from "shared-runtime";
13
+
14
+function Component(props) {
15
+ const free = makeObject_Primitives();
16
+ const free2 = makeObject_Primitives();
17
+ const part = free2.part;
18
+ useHook();
19
+ const object = useMemo(() => {
20
+ const x = makeObject_Primitives();
21
+ x.value = props.value;
22
+ mutate(x, free, part);
23
+ return x;
24
+ }, [props.value]);
25
+ return object;
26
+}
27
+
28
+export const FIXTURE_ENTRYPOINT = {
29
+ fn: Component,
30
+ params: [{ value: 42 }],
31
+};
32
+
33
+```
34
+
35
+## Code
36
+
37
+```javascript
38
+// @enablePreserveExistingMemoizationGuarantees
39
+import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
40
+import {
41
+ identity,
42
+ makeObject_Primitives,
43
+ mutate,
44
+ useHook,
45
+} from "shared-runtime";
46
+
47
+function Component(props) {
48
+ const $ = useMemoCache(4);
49
+ let t0;
50
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
51
+ t0 = makeObject_Primitives();
52
+ $[0] = t0;
53
+ } else {
54
+ t0 = $[0];
55
+ }
56
+ const free = t0;
57
+ let t1;
58
+ if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
59
+ t1 = makeObject_Primitives();
60
+ $[1] = t1;
61
+ } else {
62
+ t1 = $[1];
63
+ }
64
+ const free2 = t1;
65
+ const part = free2.part;
66
+ useHook();
67
+
68
+ props.value;
69
+ free;
70
+ part;
71
+ let t44;
72
+ let x;
73
+ if ($[2] !== props.value) {
74
+ x = makeObject_Primitives();
75
+ x.value = props.value;
76
+ mutate(x, free, part);
77
+ $[2] = props.value;
78
+ $[3] = x;
79
+ } else {
80
+ x = $[3];
81
+ }
82
+ t44 = x;
83
+ const object = t44;
84
+ return object;
85
+}
86
+
87
+export const FIXTURE_ENTRYPOINT = {
88
+ fn: Component,
89
+ params: [{ value: 42 }],
90
+};
91
+
92
+```
93
+
94
+### Eval output
95
+(kind: ok) {"a":0,"b":"value1","c":true,"value":42,"wat0":"joe"}
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.js
new
+27
@@ -0,0 +1,27 @@
1
+// @enablePreserveExistingMemoizationGuarantees
2
+import { useMemo } from "react";
3
+import {
4
+ identity,
5
+ makeObject_Primitives,
6
+ mutate,
7
+ useHook,
8
+} from "shared-runtime";
9
+
10
+function Component(props) {
11
+ const free = makeObject_Primitives();
12
+ const free2 = makeObject_Primitives();
13
+ const part = free2.part;
14
+ useHook();
15
+ const object = useMemo(() => {
16
+ const x = makeObject_Primitives();
17
+ x.value = props.value;
18
+ mutate(x, free, part);
19
+ return x;
20
+ }, [props.value]);
21
+ return object;
22
+}
23
+
24
+export const FIXTURE_ENTRYPOINT = {
25
+ fn: Component,
26
+ params: [{ value: 42 }],
27
+};
compiler/packages/sprout/src/shared-runtime.ts
+4
@@ -153,6 +153,10 @@ export function throwInput(x: Object): never {
153
throw x;
154
}
155
156
+export function useHook(): Object {
157
+ return makeObject_Primitives();
158
+}
159
+
160
const noAliasObject = Object.freeze({});
161
export function useNoAlias(...args: Array<any>): object {
162
return noAliasObject;