Preserve memoization guarantees for useCallback
Improves `@enablePreserveExistingMemoizationGuarantees` for the useCallback case. Similar to useMemo, we add an explicit `Memoize` instruction for the callback function itself _and_ for its dependencies. This means we'll assume the callback doesn't mutate any captured variables. TODO: check this with cases involving refs (should be allowed, but also not accidentally freeze the ref) and reassignment of locals (should be disallowed, though that might just be a validation we're missing today)
Joe Savona committed
Dec 15, 2023 at 13:47 UTC
5bfd70ac6fd36983a9e5e532a81b9cf2c6f994a4
5 files changed
+268
-25
compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts
+49
-25
@@ -172,12 +172,8 @@ export function dropManualMemoization(func: HIRFunction): void {
172
},
173
loc: instr.loc,
174
});
175
- } else {
176
- if (nextInstructions !== null) {
177
- nextInstructions.push(instr);
178
- }
175
+ continue;
176
}
180
- continue;
177
}
178
} else if (hookKind === "useCallback") {
179
const [fn] = instr.value.args as Array<
@@ -207,9 +203,22 @@ export function dropManualMemoization(func: HIRFunction): void {
203
* $3 = LoadLocal $2 // reference the function
204
*/
205
if (fn.kind === "Identifier") {
206
+ instr.value = {
207
+ kind: "LoadLocal",
208
+ place: {
209
+ kind: "Identifier",
210
+ identifier: fn.identifier,
211
+ effect: Effect.Unknown,
212
+ reactive: false,
213
+ loc: instr.value.loc,
214
+ },
215
+ loc: instr.value.loc,
216
+ };
217
if (
218
func.env.config.enablePreserveExistingMemoizationGuarantees
219
) {
220
+ nextInstructions =
221
+ nextInstructions ?? block.instructions.slice(0, i);
222
/**
223
* With the flag enabled the output changes to use a Memoize instruction instead
224
* a loadlocal to load the function expression into the original temporary:
@@ -228,29 +237,44 @@ export function dropManualMemoization(func: HIRFunction): void {
237
*
238
* Note the s/LoadLocal/Memoize/
239
*/
231
- instr.value = {
232
- kind: "Memoize",
240
+ const functionExpression = functions.get(fn.identifier.id);
241
+ if (functionExpression !== undefined) {
242
+ for (const operand of eachInstructionValueOperand(
243
+ functionExpression
244
+ )) {
245
+ const temp = createTemporaryPlace(func.env);
246
+ nextInstructions.push({
247
+ id: makeInstructionId(0),
248
+ lvalue: temp,
249
+ value: {
250
+ kind: "Memoize",
251
+ value: { ...operand },
252
+ loc: instr.loc,
253
+ },
254
+ loc: instr.loc,
255
+ });
256
+ }
257
+ }
258
+ nextInstructions.push(instr);
259
+
260
+ const temp = createTemporaryPlace(func.env);
261
+ nextInstructions.push({
262
+ id: makeInstructionId(0),
263
+ lvalue: { ...temp },
264
value: {
234
- kind: "Identifier",
235
- identifier: fn.identifier,
236
- effect: Effect.Unknown,
237
- reactive: false,
238
- loc: instr.value.loc,
239
- },
240
- loc: instr.value.loc,
241
- };
242
- } else {
243
- instr.value = {
244
- kind: "LoadLocal",
245
- place: {
246
- kind: "Identifier",
247
- identifier: fn.identifier,
248
- effect: Effect.Unknown,
249
- reactive: false,
265
+ kind: "Memoize",
266
+ value: {
267
+ kind: "Identifier",
268
+ identifier: fn.identifier,
269
+ effect: Effect.Unknown,
270
+ reactive: false,
271
+ loc: instr.value.loc,
272
+ },
273
loc: instr.value.loc,
274
},
252
- loc: instr.value.loc,
253
- };
275
+ loc: instr.loc,
276
+ });
277
+ continue;
278
}
279
}
280
}
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useCallback-maybe-modify-free-variable-dont-preserve-memoization-guarantee.expect.md
new
+71
@@ -0,0 +1,71 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enablePreserveExistingMemoizationGuarantees:false
6
+import { useCallback } 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 callback = useCallback(() => {
20
+ const x = makeObject_Primitives();
21
+ x.value = props.value;
22
+ mutate(x, free, part);
23
+ }, [props.value]);
24
+
25
+ mutate(free, part);
26
+ return callback;
27
+}
28
+
29
+export const FIXTURE_ENTRYPOINT = {
30
+ fn: Component,
31
+ params: [{ value: 42 }],
32
+};
33
+
34
+```
35
+
36
+## Code
37
+
38
+```javascript
39
+// @enablePreserveExistingMemoizationGuarantees:false
40
+import { useCallback } from "react";
41
+import {
42
+ identity,
43
+ makeObject_Primitives,
44
+ mutate,
45
+ useHook,
46
+} from "shared-runtime";
47
+
48
+function Component(props) {
49
+ const free = makeObject_Primitives();
50
+ const free2 = makeObject_Primitives();
51
+ const part = free2.part;
52
+ useHook();
53
+ const callback = () => {
54
+ const x = makeObject_Primitives();
55
+ x.value = props.value;
56
+ mutate(x, free, part);
57
+ };
58
+
59
+ mutate(free, part);
60
+ return callback;
61
+}
62
+
63
+export const FIXTURE_ENTRYPOINT = {
64
+ fn: Component,
65
+ params: [{ value: 42 }],
66
+};
67
+
68
+```
69
+
70
+### Eval output
71
+(kind: ok) "[[ function params=0 ]]"
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useCallback-maybe-modify-free-variable-dont-preserve-memoization-guarantee.js
new
+28
@@ -0,0 +1,28 @@
1
+// @enablePreserveExistingMemoizationGuarantees:false
2
+import { useCallback } 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 callback = useCallback(() => {
16
+ const x = makeObject_Primitives();
17
+ x.value = props.value;
18
+ mutate(x, free, part);
19
+ }, [props.value]);
20
+
21
+ mutate(free, part);
22
+ return callback;
23
+}
24
+
25
+export const FIXTURE_ENTRYPOINT = {
26
+ fn: Component,
27
+ params: [{ value: 42 }],
28
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useCallback-maybe-modify-free-variable-preserve-memoization-guarantee.expect.md
new
+93
@@ -0,0 +1,93 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enablePreserveExistingMemoizationGuarantees
6
+import { useCallback } 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 callback = useCallback(() => {
20
+ const x = makeObject_Primitives();
21
+ x.value = props.value;
22
+ mutate(x, free, part);
23
+ }, [props.value]);
24
+ mutate(free, part);
25
+ return callback;
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 { useCallback, 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
+ let t2;
68
+ if ($[2] !== props.value) {
69
+ t2 = () => {
70
+ const x = makeObject_Primitives();
71
+ x.value = props.value;
72
+ mutate(x, free, part);
73
+ };
74
+ $[2] = props.value;
75
+ $[3] = t2;
76
+ } else {
77
+ t2 = $[3];
78
+ }
79
+ const callback = t2;
80
+
81
+ mutate(free, part);
82
+ return callback;
83
+}
84
+
85
+export const FIXTURE_ENTRYPOINT = {
86
+ fn: Component,
87
+ params: [{ value: 42 }],
88
+};
89
+
90
+```
91
+
92
+### Eval output
93
+(kind: ok) "[[ function params=0 ]]"
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useCallback-maybe-modify-free-variable-preserve-memoization-guarantee.js
new
+27
@@ -0,0 +1,27 @@
1
+// @enablePreserveExistingMemoizationGuarantees
2
+import { useCallback } 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 callback = useCallback(() => {
16
+ const x = makeObject_Primitives();
17
+ x.value = props.value;
18
+ mutate(x, free, part);
19
+ }, [props.value]);
20
+ mutate(free, part);
21
+ return callback;
22
+}
23
+
24
+export const FIXTURE_ENTRYPOINT = {
25
+ fn: Component,
26
+ params: [{ value: 42 }],
27
+};