Option to not memoize non-aliased function parameters
This has been nagging at me for a _long_ time: we unnecessarily memoize function callbacks passed to things like Array.prototype.map, even though we know these functions can't escape. This PR fixes this as follows: * Adds a `noAlias?: boolean` flag to builtin function signatures, defaulting to false if not specified. * Adds a feature flag, `enableNoAliasOptimizations`, to gate optimizations based on the value of that new flag. * When the feature is enabled, `PruneNonEscapingScopes` now looks up the signature of method calls, and avoids memoizing the arguments if the signatures specifies `noAlias: true`. * Annotates Array.prototype.map and Array.prototype.filter as `noAlias`. This does not mean we'll never memoize arguments to Array.prototype.map, it just means that the argument itself won't be considered as escaping. If the function still escapes by some other means it will get memoized: ``` function Component(props) { const f = () => {}; // memoized! const x = [].map(f); // not from here.. return [x, f]; // but bc it escapes here } ``` Note: this delivers some of the wins from #1640. That PR tried to do a bunch of things, part of which I already landed w the introduction of ConditionallyMutate, which allowed us to type Array.prototype.map. This PR further gives us the ability to understand functions that don't alias their params at all. The remaining bit from #1640 is the idea of understanding that key hooks such as `useFragment()` return transitively readonly, transitively array/object/primitive values, and any `.map()` or `.filter()` calls must be on arrays, allowing us to optimize them. Without that extra step, we'll still have to memoize a lot of `array.map()` lambdas just because we aren't sure that the receiver is an Array. But this PR helps with some cases, and lays the groundwork for the rest of that PR.
Joe Savona committed
Sep 15, 2023 at 16:18 UTC
1879fbe644105de155eda887bfae733a5f5a80ca
10 files changed
+294
-11
compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+12
@@ -109,6 +109,15 @@ export type EnvironmentConfig = Partial<{
109
*/
110
enableFunctionCallSignatureOptimizations: boolean;
111
112
+ /**
113
+ * Enable optimizations based on the `noAlias` flag of method signatures. When enabled,
114
+ * function signatures can declare that they do not alias their arguments, allowing
115
+ * Forget to (in some cases) avoid memoizing arguments if they do not otherwise escape.
116
+ *
117
+ * Defaults to false
118
+ */
119
+ enableNoAliasOptimizations: boolean;
120
+
121
/**
122
* When enabled, the compiler assumes that hooks follow the Rules of React:
123
* - Hooks may memoize computation based on any of their parameters, thus
@@ -182,6 +191,7 @@ export class Environment {
191
enableFunctionCallSignatureOptimizations: boolean;
192
enableAssumeHooksFollowRulesOfReact: boolean;
193
enableTreatHooksAsFunctions: boolean;
194
+ enableNoAliasOptimizations: boolean;
195
disableAllMemoization: boolean;
196
enableEmitFreeze: ExternalFunction | null;
197
assertValidMutableRanges: boolean;
@@ -225,6 +235,8 @@ export class Environment {
235
this.validateFrozenLambdas = config?.validateFrozenLambdas ?? false;
236
this.enableFunctionCallSignatureOptimizations =
237
config?.enableFunctionCallSignatureOptimizations ?? false;
238
+ this.enableNoAliasOptimizations =
239
+ config?.enableNoAliasOptimizations ?? false;
240
this.enableAssumeHooksFollowRulesOfReact =
241
config?.enableAssumeHooksFollowRulesOfReact ?? false;
242
this.enableTreatHooksAsFunctions =
compiler/packages/babel-plugin-react-forget/src/HIR/ObjectShape.ts
+8
@@ -138,6 +138,12 @@ export type FunctionSignature = {
138
returnValueKind: ValueKind;
139
calleeEffect: Effect;
140
hookKind: HookKind | null;
141
+ /**
142
+ * Whether any of the parameters may be aliased by each other or the return
143
+ * value. Defaults to false (parameters may alias). When true, the compiler
144
+ * may choose not to memoize arguments if they do not otherwise escape.
145
+ */
146
+ noAlias?: boolean;
147
};
148
149
/**
@@ -217,6 +223,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
223
// the array object itself is not modified
224
calleeEffect: Effect.ConditionallyMutate,
225
returnValueKind: ValueKind.Mutable,
226
+ noAlias: true,
227
}),
228
],
229
[
@@ -230,6 +237,7 @@ addObject(BUILTIN_SHAPES, BuiltInArrayId, [
237
// the array object itself is not modified
238
calleeEffect: Effect.ConditionallyMutate,
239
returnValueKind: ValueKind.Mutable,
240
+ noAlias: true,
241
}),
242
],
243
[
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PruneNonEscapingScopes.ts
+46
-10
@@ -24,6 +24,7 @@ import {
24
getHookKind,
25
isMutableEffect,
26
} from "../HIR";
27
+import { getFunctionCallSignature } from "../Inference/InferReferenceEffects";
28
import { log } from "../Utils/logger";
29
import { assertExhaustive } from "../Utils/utils";
30
import { getPlaceScope } from "./BuildReactiveBlocks";
@@ -119,7 +120,11 @@ export function pruneNonEscapingScopes(
120
for (const param of fn.params) {
121
state.declare(param.identifier.id);
122
}
122
- visitReactiveFunction(fn, new CollectDependenciesVisitor(options), state);
123
+ visitReactiveFunction(
124
+ fn,
125
+ new CollectDependenciesVisitor(fn.env, options),
126
+ state
127
+ );
128
129
log(() => prettyFormat(state));
130
@@ -343,6 +348,7 @@ type LValueMemoization = {
348
* - level: the level of memoization to apply to this value
349
*/
350
function computeMemoizationInputs(
351
+ env: Environment,
352
value: ReactiveValue,
353
lvalue: Place | null,
354
options: MemoizationOptions
@@ -361,8 +367,10 @@ function computeMemoizationInputs(
367
: [],
368
rvalues: [
369
// Conditionals do not alias their test value.
364
- ...computeMemoizationInputs(value.consequent, null, options).rvalues,
365
- ...computeMemoizationInputs(value.alternate, null, options).rvalues,
370
+ ...computeMemoizationInputs(env, value.consequent, null, options)
371
+ .rvalues,
372
+ ...computeMemoizationInputs(env, value.alternate, null, options)
373
+ .rvalues,
374
],
375
};
376
}
@@ -374,8 +382,8 @@ function computeMemoizationInputs(
382
? [{ place: lvalue, level: MemoizationLevel.Conditional }]
383
: [],
384
rvalues: [
377
- ...computeMemoizationInputs(value.left, null, options).rvalues,
378
- ...computeMemoizationInputs(value.right, null, options).rvalues,
385
+ ...computeMemoizationInputs(env, value.left, null, options).rvalues,
386
+ ...computeMemoizationInputs(env, value.right, null, options).rvalues,
387
],
388
};
389
}
@@ -389,7 +397,8 @@ function computeMemoizationInputs(
397
// Only the final value of the sequence is a true rvalue:
398
// values from the sequence's instructions are evaluated
399
// as separate nodes
392
- rvalues: computeMemoizationInputs(value.value, null, options).rvalues,
400
+ rvalues: computeMemoizationInputs(env, value.value, null, options)
401
+ .rvalues,
402
};
403
}
404
case "JsxExpression": {
@@ -595,18 +604,42 @@ function computeMemoizationInputs(
604
return {
605
lvalues: lvalues,
606
rvalues: [
598
- ...computeMemoizationInputs(value.value, null, options).rvalues,
607
+ ...computeMemoizationInputs(env, value.value, null, options).rvalues,
608
],
609
};
610
}
611
+ case "MethodCall": {
612
+ const signature = env.enableNoAliasOptimizations
613
+ ? getFunctionCallSignature(env, value.property.identifier.type)
614
+ : null;
615
+ const operands = [...eachReactiveValueOperand(value)];
616
+ let lvalues = [];
617
+ if (lvalue !== null) {
618
+ lvalues.push({ place: lvalue, level: MemoizationLevel.Memoized });
619
+ }
620
+ if (signature?.noAlias === true) {
621
+ return {
622
+ lvalues,
623
+ rvalues: [],
624
+ };
625
+ }
626
+ lvalues.push(
627
+ ...operands
628
+ .filter((operand) => isMutableEffect(operand.effect, operand.loc))
629
+ .map((place) => ({ place, level: MemoizationLevel.Memoized }))
630
+ );
631
+ return {
632
+ lvalues,
633
+ rvalues: operands,
634
+ };
635
+ }
636
case "RegExpLiteral":
637
case "FunctionExpression":
638
case "TaggedTemplateExpression":
605
- case "CallExpression":
639
case "ArrayExpression":
640
case "NewExpression":
641
case "ObjectExpression":
609
- case "MethodCall":
642
+ case "CallExpression":
643
case "PropertyStore": {
644
// All of these instructions may produce new values which must be memoized if
645
// reachable from a return value. Any mutable rvalue may alias any other rvalue
@@ -680,10 +713,12 @@ function computePatternLValues(pattern: Pattern): Array<LValueMemoization> {
713
* identifier's and scope's dependencies.
714
*/
715
class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
716
+ env: Environment;
717
options: MemoizationOptions;
718
685
- constructor(options: MemoizationOptions) {
719
+ constructor(env: Environment, options: MemoizationOptions) {
720
super();
721
+ this.env = env;
722
this.options = options;
723
}
724
@@ -695,6 +730,7 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
730
731
// Determe the level of memoization for this value and the lvalues/rvalues
732
const aliasing = computeMemoizationInputs(
733
+ this.env,
734
instruction.value,
735
instruction.lvalue,
736
this.options
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.expect.md
new
+60
@@ -0,0 +1,60 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enableNoAliasOptimizations
6
+function Component(props) {
7
+ const x = [];
8
+ <dif>{x}</dif>;
9
+ const y = x.map((item) => item);
10
+ return [x, y];
11
+}
12
+
13
+export const FIXTURE_ENTRYPOINT = {
14
+ fn: Component,
15
+ params: [{}],
16
+ isComponent: false,
17
+};
18
+
19
+```
20
+
21
+## Code
22
+
23
+```javascript
24
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableNoAliasOptimizations
25
+function Component(props) {
26
+ const $ = useMemoCache(3);
27
+ let t0;
28
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
29
+ t0 = [];
30
+ $[0] = t0;
31
+ } else {
32
+ t0 = $[0];
33
+ }
34
+ const x = t0;
35
+ let t1;
36
+ if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
37
+ t1 = x.map((item) => item);
38
+ $[1] = t1;
39
+ } else {
40
+ t1 = $[1];
41
+ }
42
+ const y = t1;
43
+ let t2;
44
+ if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
45
+ t2 = [x, y];
46
+ $[2] = t2;
47
+ } else {
48
+ t2 = $[2];
49
+ }
50
+ return t2;
51
+}
52
+
53
+export const FIXTURE_ENTRYPOINT = {
54
+ fn: Component,
55
+ params: [{}],
56
+ isComponent: false,
57
+};
58
+
59
+```
60
+
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-frozen-array-noAlias.js
new
+13
@@ -0,0 +1,13 @@
1
+// @enableNoAliasOptimizations
2
+function Component(props) {
3
+ const x = [];
4
+ <dif>{x}</dif>;
5
+ const y = x.map((item) => item);
6
+ return [x, y];
7
+}
8
+
9
+export const FIXTURE_ENTRYPOINT = {
10
+ fn: Component,
11
+ params: [{}],
12
+ isComponent: false,
13
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.expect.md
new
+61
@@ -0,0 +1,61 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @enableNoAliasOptimizations
6
+function Component(props) {
7
+ const x = [];
8
+ const y = x.map((item) => {
9
+ item.updated = true;
10
+ return item;
11
+ });
12
+ return [x, y];
13
+}
14
+
15
+export const FIXTURE_ENTRYPOINT = {
16
+ fn: Component,
17
+ params: [{}],
18
+ isComponent: false,
19
+};
20
+
21
+```
22
+
23
+## Code
24
+
25
+```javascript
26
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableNoAliasOptimizations
27
+function Component(props) {
28
+ const $ = useMemoCache(3);
29
+ let t0;
30
+ let x;
31
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
32
+ x = [];
33
+ t0 = x.map((item) => {
34
+ item.updated = true;
35
+ return item;
36
+ });
37
+ $[0] = t0;
38
+ $[1] = x;
39
+ } else {
40
+ t0 = $[0];
41
+ x = $[1];
42
+ }
43
+ const y = t0;
44
+ let t1;
45
+ if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
46
+ t1 = [x, y];
47
+ $[2] = t1;
48
+ } else {
49
+ t1 = $[2];
50
+ }
51
+ return t1;
52
+}
53
+
54
+export const FIXTURE_ENTRYPOINT = {
55
+ fn: Component,
56
+ params: [{}],
57
+ isComponent: false,
58
+};
59
+
60
+```
61
+
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-mutable-array-mutating-lambda-noAlias.js
new
+15
@@ -0,0 +1,15 @@
1
+// @enableNoAliasOptimizations
2
+function Component(props) {
3
+ const x = [];
4
+ const y = x.map((item) => {
5
+ item.updated = true;
6
+ return item;
7
+ });
8
+ return [x, y];
9
+}
10
+
11
+export const FIXTURE_ENTRYPOINT = {
12
+ fn: Component,
13
+ params: [{}],
14
+ isComponent: false,
15
+};
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.expect.md
new
+62
@@ -0,0 +1,62 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+function Component(props) {
6
+ const f = (item) => item;
7
+ const x = [...props.items].map(f); // `f` doesn't escape here...
8
+ return [x, f]; // ...but it does here so it's memoized
9
+}
10
+
11
+export const FIXTURE_ENTRYPOINT = {
12
+ fn: Component,
13
+ params: [{ items: [{ id: 1 }] }],
14
+ isComponent: false,
15
+};
16
+
17
+```
18
+
19
+## Code
20
+
21
+```javascript
22
+import { unstable_useMemoCache as useMemoCache } from "react";
23
+function Component(props) {
24
+ const $ = useMemoCache(5);
25
+ let t0;
26
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
27
+ t0 = (item) => item;
28
+ $[0] = t0;
29
+ } else {
30
+ t0 = $[0];
31
+ }
32
+ const f = t0;
33
+ const c_1 = $[1] !== props.items;
34
+ let t1;
35
+ if (c_1) {
36
+ t1 = [...props.items].map(f);
37
+ $[1] = props.items;
38
+ $[2] = t1;
39
+ } else {
40
+ t1 = $[2];
41
+ }
42
+ const x = t1;
43
+ const c_3 = $[3] !== x;
44
+ let t2;
45
+ if (c_3) {
46
+ t2 = [x, f];
47
+ $[3] = x;
48
+ $[4] = t2;
49
+ } else {
50
+ t2 = $[4];
51
+ }
52
+ return t2;
53
+}
54
+
55
+export const FIXTURE_ENTRYPOINT = {
56
+ fn: Component,
57
+ params: [{ items: [{ id: 1 }] }],
58
+ isComponent: false,
59
+};
60
+
61
+```
62
+
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/array-map-noAlias-escaping-function.js
new
+11
@@ -0,0 +1,11 @@
1
+function Component(props) {
2
+ const f = (item) => item;
3
+ const x = [...props.items].map(f); // `f` doesn't escape here...
4
+ return [x, f]; // ...but it does here so it's memoized
5
+}
6
+
7
+export const FIXTURE_ENTRYPOINT = {
8
+ fn: Component,
9
+ params: [{ items: [{ id: 1 }] }],
10
+ isComponent: false,
11
+};
compiler/packages/fixture-test-utils/src/compiler-utils.ts
+6
-1
@@ -28,6 +28,7 @@ export function transformFixtureInput(
28
let enableEmitFreeze = null;
29
let compilationMode: CompilationMode = "all";
30
let enableForest = false;
31
+ let enableNoAliasOptimizations = false;
32
33
if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
34
assert(
@@ -83,6 +84,9 @@ export function transformFixtureInput(
84
if (firstLine.includes("@enableForest true")) {
85
enableForest = true;
86
}
87
+ if (firstLine.includes("@enableNoAliasOptimizations")) {
88
+ enableNoAliasOptimizations = true;
89
+ }
90
91
return pluginFn(
92
input,
@@ -101,6 +105,7 @@ export function transformFixtureInput(
105
]),
106
enableAssumeHooksFollowRulesOfReact,
107
enableFunctionCallSignatureOptimizations: true,
108
+ enableNoAliasOptimizations,
109
disableAllMemoization,
110
enableTreatHooksAsFunctions,
111
inlineUseMemo: true,
@@ -117,7 +122,7 @@ export function transformFixtureInput(
122
logger: null,
123
gating,
124
instrumentForget,
120
- panicThreshold: 'ALL_ERRORS',
125
+ panicThreshold: "ALL_ERRORS",
126
noEmit: false,
127
},
128
includeAst