[RFC] Refine memoization for Array#map with non-mutating callbacks
Improves memoization for cases such as #2409: ```javascript const x = []; useEffect(...); return <div>{x.map(item => <span>{item}</span>)}</div>; ``` We previously thought that the `x.map(...)` call mutated `x` since its kind was Mutable. However, in this case we can determine that the map call cannot mutate `x` (or anything else): the lambda does not mutate any free variables and does not mutate its arguments. This PR adds a new flag to function signatures, used for method calls only, that checks for such cases. The idea is that if the receiver is the only thing that is mutable — including that there are no args which are function expressions which mutate their parameters — then we can infer the effect as a read. See tests which confirm that function expressions which capture or mutate their params bypass the optimization.
Joe Savona committed
Nov 29, 2023 at 10:46 UTC
f7d16db54480558fe792d04d44ae3b8d18603894
10 files changed
+456
-13
compiler/packages/babel-plugin-react-forget/src/HIR/HIR.ts
+1
@@ -971,6 +971,7 @@ export enum Effect {
971
* But we do not error if the value is known to be immutable.
972
*/
973
ConditionallyMutate = "mutate?",
974
+
975
/*
976
* This reference *does* write to (mutate) the value. It is an error (invalid input)
977
* if an immutable value flows into a location with this effect.
compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts
+79
@@ -147,6 +147,20 @@ export type FunctionSignature = {
147
* may choose not to memoize arguments if they do not otherwise escape.
148
*/
149
noAlias?: boolean;
150
+
151
+ /**
152
+ * Supported only for methods (no-op when used on functions in CallExpression.callee position).
153
+ *
154
+ * Indicates that the method can only modify its receiver if any of the arguments
155
+ * are mutable or are function expressions which mutate their arguments. This is designed
156
+ * for methods such as Array.prototype.map(), which only mutate the receiver array if they are
157
+ * passed a callback which has mutable side-effects (including mutating its inputs).
158
+ *
159
+ * MethodCalls to such functions will use a different behavior depending on their arguments:
160
+ * - If arguments are all non-mutable, the arguments get the Read effect and the receiver is Capture.
161
+ * - Else uses the effects specified by this signature.
162
+ */
163
+ mutableOnlyIfOperandsAreMutable?: boolean;
164
};
165
166
/*
@@ -231,6 +245,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
245
calleeEffect: Effect.ConditionallyMutate,
246
returnValueKind: ValueKind.Mutable,
247
noAlias: true,
248
+ mutableOnlyIfOperandsAreMutable: true,
249
}),
250
],
251
[
@@ -247,6 +262,70 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
262
calleeEffect: Effect.ConditionallyMutate,
263
returnValueKind: ValueKind.Mutable,
264
noAlias: true,
265
+ mutableOnlyIfOperandsAreMutable: true,
266
+ }),
267
+ ],
268
+ [
269
+ "every",
270
+ addFunction(BUILTIN_SHAPES, [], {
271
+ positionalParams: [],
272
+ restParam: Effect.ConditionallyMutate,
273
+ returnType: { kind: "Primitive" },
274
+ /*
275
+ * callee is ConditionallyMutate because items of the array
276
+ * flow into the lambda and may be mutated there, even though
277
+ * the array object itself is not modified
278
+ */
279
+ calleeEffect: Effect.ConditionallyMutate,
280
+ returnValueKind: ValueKind.Immutable,
281
+ noAlias: true,
282
+ mutableOnlyIfOperandsAreMutable: true,
283
+ }),
284
+ ],
285
+ [
286
+ "some",
287
+ addFunction(BUILTIN_SHAPES, [], {
288
+ positionalParams: [],
289
+ restParam: Effect.ConditionallyMutate,
290
+ returnType: { kind: "Primitive" },
291
+ /*
292
+ * callee is ConditionallyMutate because items of the array
293
+ * flow into the lambda and may be mutated there, even though
294
+ * the array object itself is not modified
295
+ */
296
+ calleeEffect: Effect.ConditionallyMutate,
297
+ returnValueKind: ValueKind.Immutable,
298
+ noAlias: true,
299
+ mutableOnlyIfOperandsAreMutable: true,
300
+ }),
301
+ ],
302
+ [
303
+ "find",
304
+ addFunction(BUILTIN_SHAPES, [], {
305
+ positionalParams: [],
306
+ restParam: Effect.ConditionallyMutate,
307
+ returnType: { kind: "Poly" },
308
+ calleeEffect: Effect.ConditionallyMutate,
309
+ returnValueKind: ValueKind.Mutable,
310
+ noAlias: true,
311
+ mutableOnlyIfOperandsAreMutable: true,
312
+ }),
313
+ ],
314
+ [
315
+ "findIndex",
316
+ addFunction(BUILTIN_SHAPES, [], {
317
+ positionalParams: [],
318
+ restParam: Effect.ConditionallyMutate,
319
+ returnType: { kind: "Primitive" },
320
+ /*
321
+ * callee is ConditionallyMutate because items of the array
322
+ * flow into the lambda and may be mutated there, even though
323
+ * the array object itself is not modified
324
+ */
325
+ calleeEffect: Effect.ConditionallyMutate,
326
+ returnValueKind: ValueKind.Immutable,
327
+ noAlias: true,
328
+ mutableOnlyIfOperandsAreMutable: true,
329
}),
330
],
331
[
compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
+76
@@ -228,6 +228,17 @@ class InferenceState {
228
this.#values.set(value, kind);
229
}
230
231
+ values(place: Place): Array<InstructionValue> {
232
+ const values = this.#variables.get(place.identifier.id);
233
+ CompilerError.invariant(values != null, {
234
+ reason: `[hoisting] Expected value kind to be initialized`,
235
+ description: `${printPlace(place)}`,
236
+ loc: place.loc,
237
+ suggestions: null,
238
+ });
239
+ return Array.from(values);
240
+ }
241
+
242
// Lookup the kind of the given @param value.
243
kind(place: Place): ValueKind {
244
const values = this.#variables.get(place.identifier.id);
@@ -833,6 +844,26 @@ function inferBlock(
844
instrValue.property.identifier.type
845
);
846
847
+ if (
848
+ signature !== null &&
849
+ signature.mutableOnlyIfOperandsAreMutable &&
850
+ areArgumentsImmutableAndNonMutating(state, instrValue.args)
851
+ ) {
852
+ /*
853
+ * None of the args are mutable or mutate their params, we can downgrade to
854
+ * treating as all reads
855
+ */
856
+ for (const arg of instrValue.args) {
857
+ const place = arg.kind === "Identifier" ? arg : arg.place;
858
+ state.reference(place, Effect.Read);
859
+ }
860
+ state.reference(instrValue.receiver, Effect.Read);
861
+ state.initialize(instrValue, signature.returnValueKind);
862
+ state.define(instr.lvalue, instrValue);
863
+ instr.lvalue.effect = Effect.ConditionallyMutate;
864
+ continue;
865
+ }
866
+
867
const effects =
868
signature !== null ? getFunctionEffects(instrValue, signature) : null;
869
const returnValueKind =
@@ -1185,3 +1216,48 @@ function getFunctionEffects(
1216
}
1217
return results;
1218
}
1219
+
1220
+/**
1221
+ * Returns true if all of the arguments are both non-mutable (immutable or frozen)
1222
+ * _and_ are not functions which might mutate their arguments. Note that function
1223
+ * expressions count as frozen so long as they do not mutate free variables: this
1224
+ * function checks that such functions also don't mutate their inputs.
1225
+ */
1226
+function areArgumentsImmutableAndNonMutating(
1227
+ state: InferenceState,
1228
+ args: MethodCall["args"]
1229
+): boolean {
1230
+ for (const arg of args) {
1231
+ const place = arg.kind === "Identifier" ? arg : arg.place;
1232
+ const kind = state.kind(place);
1233
+ switch (kind) {
1234
+ case ValueKind.Immutable:
1235
+ case ValueKind.Frozen: {
1236
+ /*
1237
+ * Only immutable values, or frozen lambdas are allowed.
1238
+ * A lambda may appear frozen even if it may mutate its inputs,
1239
+ * so we have a second check even for frozen value types
1240
+ */
1241
+ break;
1242
+ }
1243
+ default: {
1244
+ return false;
1245
+ }
1246
+ }
1247
+ const values = state.values(place);
1248
+ for (const value of values) {
1249
+ if (
1250
+ value.kind === "FunctionExpression" &&
1251
+ value.loweredFunc.func.params.some((param) => {
1252
+ const place = param.kind === "Identifier" ? param : param.place;
1253
+ const range = place.identifier.mutableRange;
1254
+ return range.end > range.start + 1;
1255
+ })
1256
+ ) {
1257
+ // This is a function which may mutate its inputs
1258
+ return false;
1259
+ }
1260
+ }
1261
+ }
1262
+ return true;
1263
+}
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-captures-receiver-noAlias.expect.md
+21
-5
@@ -23,18 +23,34 @@ export const FIXTURE_ENTRYPOINT = {
23
```javascript
24
import { unstable_useMemoCache as useMemoCache } from "react";
25
function Component(props) {
26
- const $ = useMemoCache(2);
26
+ const $ = useMemoCache(6);
27
let t0;
28
if ($[0] !== props.a) {
29
- const item = { a: props.a };
30
- const items = [item];
31
- t0 = items.map((item_0) => item_0);
29
+ t0 = { a: props.a };
30
$[0] = props.a;
31
$[1] = t0;
32
} else {
33
t0 = $[1];
34
}
37
- const mapped = t0;
35
+ const item = t0;
36
+ let t1;
37
+ if ($[2] !== item) {
38
+ t1 = [item];
39
+ $[2] = item;
40
+ $[3] = t1;
41
+ } else {
42
+ t1 = $[3];
43
+ }
44
+ const items = t1;
45
+ let t2;
46
+ if ($[4] !== items) {
47
+ t2 = items.map((item_0) => item_0);
48
+ $[4] = items;
49
+ $[5] = t2;
50
+ } else {
51
+ t2 = $[5];
52
+ }
53
+ const mapped = t2;
54
return mapped;
55
}
56
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.expect.md
renamed
+23
-8
@@ -3,25 +3,29 @@
3
4
```javascript
5
import { useEffect, useState } from "react";
6
+import { mutate } from "shared-runtime";
7
8
function Component(props) {
8
- const x = [props.value];
9
+ const x = [{ ...props.value }];
10
useEffect(() => {}, []);
11
const onClick = () => {
12
console.log(x.length);
13
};
14
+ let y;
15
return (
16
<div onClick={onClick}>
17
{x.map((item) => {
16
- return <span key={item}>{item}</span>;
18
+ y = item;
19
+ return <span key={item.id}>{item.text}</span>;
20
})}
21
+ {mutate(y)}
22
</div>
23
);
24
}
25
26
export const FIXTURE_ENTRYPOINT = {
27
fn: Component,
24
- params: [{ value: 42 }],
28
+ params: [{ value: { id: 0, text: "Hello!" } }],
29
isComponent: true,
30
};
31
@@ -35,10 +39,11 @@ import {
39
useState,
40
unstable_useMemoCache as useMemoCache,
41
} from "react";
42
+import { mutate } from "shared-runtime";
43
44
function Component(props) {
45
const $ = useMemoCache(5);
41
- const x = [props.value];
46
+ const x = [{ ...props.value }];
47
let t0;
48
let t1;
49
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -55,10 +60,20 @@ function Component(props) {
60
console.log(x.length);
61
};
62
58
- const t2 = x.map((item) => <span key={item}>{item}</span>);
63
+ let y;
64
+
65
+ const t2 = x.map((item) => {
66
+ y = item;
67
+ return <span key={item.id}>{item.text}</span>;
68
+ });
69
let t3;
70
if ($[2] !== onClick || $[3] !== t2) {
61
- t3 = <div onClick={onClick}>{t2}</div>;
71
+ t3 = (
72
+ <div onClick={onClick}>
73
+ {t2}
74
+ {mutate(y)}
75
+ </div>
76
+ );
77
$[2] = onClick;
78
$[3] = t2;
79
$[4] = t3;
@@ -70,11 +85,11 @@ function Component(props) {
85
86
export const FIXTURE_ENTRYPOINT = {
87
fn: Component,
73
- params: [{ value: 42 }],
88
+ params: [{ value: { id: 0, text: "Hello!" } }],
89
isComponent: true,
90
};
91
92
```
93
94
### Eval output
80
-(kind: ok) <div><span>42</span></div>
\ No newline at end of file
95
+(kind: ok) <div><span>Hello!</span></div>
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-capturing-map-after-hook.js
new
+26
@@ -0,0 +1,26 @@
1
+import { useEffect, useState } from "react";
2
+import { mutate } from "shared-runtime";
3
+
4
+function Component(props) {
5
+ const x = [{ ...props.value }];
6
+ useEffect(() => {}, []);
7
+ const onClick = () => {
8
+ console.log(x.length);
9
+ };
10
+ let y;
11
+ return (
12
+ <div onClick={onClick}>
13
+ {x.map((item) => {
14
+ y = item;
15
+ return <span key={item.id}>{item.text}</span>;
16
+ })}
17
+ {mutate(y)}
18
+ </div>
19
+ );
20
+}
21
+
22
+export const FIXTURE_ENTRYPOINT = {
23
+ fn: Component,
24
+ params: [{ value: { id: 0, text: "Hello!" } }],
25
+ isComponent: true,
26
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.expect.md
new
+102
@@ -0,0 +1,102 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+import { useEffect, useState } from "react";
6
+import { mutate } from "shared-runtime";
7
+
8
+function Component(props) {
9
+ const x = [{ ...props.value }];
10
+ useEffect(() => {}, []);
11
+ const onClick = () => {
12
+ console.log(x.length);
13
+ };
14
+ let y;
15
+ return (
16
+ <div onClick={onClick}>
17
+ {x.map((item) => {
18
+ item.flag = true;
19
+ return <span key={item.id}>{item.text}</span>;
20
+ })}
21
+ {mutate(y)}
22
+ </div>
23
+ );
24
+}
25
+
26
+export const FIXTURE_ENTRYPOINT = {
27
+ fn: Component,
28
+ params: [{ value: { id: 0, text: "Hello", flag: false } }],
29
+ isComponent: true,
30
+};
31
+
32
+```
33
+
34
+## Code
35
+
36
+```javascript
37
+import {
38
+ useEffect,
39
+ useState,
40
+ unstable_useMemoCache as useMemoCache,
41
+} from "react";
42
+import { mutate } from "shared-runtime";
43
+
44
+function Component(props) {
45
+ const $ = useMemoCache(6);
46
+ const x = [{ ...props.value }];
47
+ let t0;
48
+ let t1;
49
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
50
+ t0 = () => {};
51
+ t1 = [];
52
+ $[0] = t0;
53
+ $[1] = t1;
54
+ } else {
55
+ t0 = $[0];
56
+ t1 = $[1];
57
+ }
58
+ useEffect(t0, t1);
59
+ const onClick = () => {
60
+ console.log(x.length);
61
+ };
62
+
63
+ let y;
64
+
65
+ const t3 = x.map((item) => {
66
+ item.flag = true;
67
+ return <span key={item.id}>{item.text}</span>;
68
+ });
69
+ let t2;
70
+ if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
71
+ t2 = mutate(y);
72
+ $[2] = t2;
73
+ } else {
74
+ t2 = $[2];
75
+ }
76
+ let t4;
77
+ if ($[3] !== onClick || $[4] !== t3) {
78
+ t4 = (
79
+ <div onClick={onClick}>
80
+ {t3}
81
+ {t2}
82
+ </div>
83
+ );
84
+ $[3] = onClick;
85
+ $[4] = t3;
86
+ $[5] = t4;
87
+ } else {
88
+ t4 = $[5];
89
+ }
90
+ return t4;
91
+}
92
+
93
+export const FIXTURE_ENTRYPOINT = {
94
+ fn: Component,
95
+ params: [{ value: { id: 0, text: "Hello", flag: false } }],
96
+ isComponent: true,
97
+};
98
+
99
+```
100
+
101
+### Eval output
102
+(kind: ok) <div><span>Hello</span></div>
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-dont-memoize-array-with-mutable-map-after-hook.js
new
+26
@@ -0,0 +1,26 @@
1
+import { useEffect, useState } from "react";
2
+import { mutate } from "shared-runtime";
3
+
4
+function Component(props) {
5
+ const x = [{ ...props.value }];
6
+ useEffect(() => {}, []);
7
+ const onClick = () => {
8
+ console.log(x.length);
9
+ };
10
+ let y;
11
+ return (
12
+ <div onClick={onClick}>
13
+ {x.map((item) => {
14
+ item.flag = true;
15
+ return <span key={item.id}>{item.text}</span>;
16
+ })}
17
+ {mutate(y)}
18
+ </div>
19
+ );
20
+}
21
+
22
+export const FIXTURE_ENTRYPOINT = {
23
+ fn: Component,
24
+ params: [{ value: { id: 0, text: "Hello", flag: false } }],
25
+ isComponent: true,
26
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/repro-memoize-array-with-immutable-map-after-hook.expect.md
new
+102
@@ -0,0 +1,102 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+import { useEffect, useState } from "react";
6
+
7
+function Component(props) {
8
+ const x = [props.value];
9
+ useEffect(() => {}, []);
10
+ const onClick = () => {
11
+ console.log(x.length);
12
+ };
13
+ return (
14
+ <div onClick={onClick}>
15
+ {x.map((item) => {
16
+ return <span key={item}>{item}</span>;
17
+ })}
18
+ </div>
19
+ );
20
+}
21
+
22
+export const FIXTURE_ENTRYPOINT = {
23
+ fn: Component,
24
+ params: [{ value: 42 }],
25
+ isComponent: true,
26
+};
27
+
28
+```
29
+
30
+## Code
31
+
32
+```javascript
33
+import {
34
+ useEffect,
35
+ useState,
36
+ unstable_useMemoCache as useMemoCache,
37
+} from "react";
38
+
39
+function Component(props) {
40
+ const $ = useMemoCache(11);
41
+ let t0;
42
+ if ($[0] !== props.value) {
43
+ t0 = [props.value];
44
+ $[0] = props.value;
45
+ $[1] = t0;
46
+ } else {
47
+ t0 = $[1];
48
+ }
49
+ const x = t0;
50
+ let t1;
51
+ let t2;
52
+ if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
53
+ t1 = () => {};
54
+ t2 = [];
55
+ $[2] = t1;
56
+ $[3] = t2;
57
+ } else {
58
+ t1 = $[2];
59
+ t2 = $[3];
60
+ }
61
+ useEffect(t1, t2);
62
+ let t3;
63
+ if ($[4] !== x.length) {
64
+ t3 = () => {
65
+ console.log(x.length);
66
+ };
67
+ $[4] = x.length;
68
+ $[5] = t3;
69
+ } else {
70
+ t3 = $[5];
71
+ }
72
+ const onClick = t3;
73
+ let t4;
74
+ if ($[6] !== x) {
75
+ t4 = x.map((item) => <span key={item}>{item}</span>);
76
+ $[6] = x;
77
+ $[7] = t4;
78
+ } else {
79
+ t4 = $[7];
80
+ }
81
+ let t5;
82
+ if ($[8] !== onClick || $[9] !== t4) {
83
+ t5 = <div onClick={onClick}>{t4}</div>;
84
+ $[8] = onClick;
85
+ $[9] = t4;
86
+ $[10] = t5;
87
+ } else {
88
+ t5 = $[10];
89
+ }
90
+ return t5;
91
+}
92
+
93
+export const FIXTURE_ENTRYPOINT = {
94
+ fn: Component,
95
+ params: [{ value: 42 }],
96
+ isComponent: true,
97
+};
98
+
99
+```
100
+
101
+### Eval output
102
+(kind: ok) <div><span>42</span></div>
\ No newline at end of file