Memoize hook args (treat as escaping)
Updates `PruneNonEscapingScopes` to consider hook arguments as potentially escaping. This is because hook inputs are "owned" by React — for example, closures passed to `useEffect`, or a value that is passed to a custom hook and which then becomes a memoized input.
Joe Savona committed
Apr 4, 2023 at 12:30 UTC
2e3aa3954ce69db9b1e488d09cd49dcb2e5b7bf7
4 files changed
+131
-7
compiler/forget/src/HIR/Globals.ts
+18
@@ -169,6 +169,24 @@ const BUILTIN_HOOKS: Array<[string, Hook]> = [
169
valueKind: ValueKind.Frozen,
170
},
171
],
172
+ [
173
+ "useEffect",
174
+ {
175
+ kind: "Memo",
176
+ name: "useEffect",
177
+ effectKind: Effect.Freeze,
178
+ valueKind: ValueKind.Frozen,
179
+ },
180
+ ],
181
+ [
182
+ "useLayoutEffect",
183
+ {
184
+ kind: "Memo",
185
+ name: "useLayoutEffect",
186
+ effectKind: Effect.Freeze,
187
+ valueKind: ValueKind.Frozen,
188
+ },
189
+ ],
190
];
191
192
export type Global = BuiltInType | HookType | PolyType;
compiler/forget/src/ReactiveScopes/PruneNonEscapingScopes.ts
+22
-7
@@ -12,6 +12,7 @@ import {
12
Effect,
13
IdentifierId,
14
InstructionId,
15
+ isHookType,
16
Pattern,
17
Place,
18
ReactiveFunction,
@@ -23,6 +24,7 @@ import {
24
ReactiveValue,
25
ScopeId,
26
} from "../HIR";
27
+import { eachInstructionValueOperand } from "../HIR/visitors";
28
import { log } from "../Utils/logger";
29
import { assertExhaustive } from "../Utils/utils";
30
import { getPlaceScope } from "./BuildReactiveBlocks";
@@ -37,8 +39,13 @@ import {
39
40
/**
41
* This pass prunes reactive scopes that are not necessary to bound downstream computation.
40
- * Specifically, the pass identifies the set of identifiers which are directly returned by
41
- * the function and/or transitively aliased by a return value - ie, values that "escape".
42
+ * Specifically, the pass identifies the set of identifiers which may "escape". Values can
43
+ * escape in one of two ways:
44
+ * * They are directly returned by the function and/or transitively aliased by a return
45
+ * value.
46
+ * * They are passed as input to a hook. This is because any value passed to a hook may
47
+ * have its referenced ultimately stored by React (ie, be aliased by an external value).
48
+ * For example, the closure passed to useEffect escapes.
49
*
50
* Example to build intuition:
51
*
@@ -96,7 +103,8 @@ import {
103
* b. Conditional and logical expressions (and a few others) are conditinally aliased,
104
* depending on whether their result value is aliased.
105
* c. JSX is always unaliased (though its props children may be)
99
- * 2. The same pass which builds the graph also stores the set of returned identifiers.
106
+ * 2. The same pass which builds the graph also stores the set of returned identifiers and set of
107
+ * identifiers passed as arguments to hooks.
108
* 3. We traverse the graph starting from the returned identifiers and mark reachable dependencies
109
* as escaping, based on the combination of the parent node's type and its children (eg a
110
* conditional node with an aliased dep promotes to aliased).
@@ -198,7 +206,7 @@ class State {
206
207
identifiers: Map<IdentifierId, IdentifierNode> = new Map();
208
scopes: Map<ScopeId, ScopeNode> = new Map();
201
- returned: Set<IdentifierId> = new Set();
209
+ escapingValues: Set<IdentifierId> = new Set();
210
211
/**
212
* Declare a new identifier, used for function id and params
@@ -302,8 +310,8 @@ function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {
310
}
311
312
// Walk from the "roots" aka returned identifiers.
305
- for (const returned of state.returned) {
306
- visit(returned);
313
+ for (const value of state.escapingValues) {
314
+ visit(value);
315
}
316
317
return memoized;
@@ -649,6 +657,13 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
657
instruction.lvalue.identifier.id,
658
instruction.value.place.identifier.id
659
);
660
+ } else if (instruction.value.kind === "CallExpression") {
661
+ const callee = instruction.value.callee;
662
+ if (isHookType(callee.identifier)) {
663
+ for (const operand of eachInstructionValueOperand(instruction.value)) {
664
+ state.escapingValues.add(operand.identifier.id);
665
+ }
666
+ }
667
}
668
}
669
@@ -659,7 +674,7 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
674
this.traverseTerminal(stmt, state);
675
676
if (stmt.terminal.kind === "return" && stmt.terminal.value !== null) {
662
- state.returned.add(stmt.terminal.value.identifier.id);
677
+ state.escapingValues.add(stmt.terminal.value.identifier.id);
678
}
679
}
680
}
compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.expect.md
new
+75
@@ -0,0 +1,75 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+function Component(props) {
6
+ const dispatch = useDispatch();
7
+ useFreeze(dispatch);
8
+
9
+ // onUpdate should be memoized even though it doesn't
10
+ // flow into the return value
11
+ const onUpdate = () => {
12
+ dispatch({ kind: "update" });
13
+ };
14
+
15
+ useEffect(() => {
16
+ onUpdate();
17
+ }, [onUpdate]);
18
+
19
+ return <div />;
20
+}
21
+
22
+```
23
+
24
+## Code
25
+
26
+```javascript
27
+function Component(props) {
28
+ const $ = React.unstable_useMemoCache(7);
29
+ const dispatch = useDispatch();
30
+ useFreeze(dispatch);
31
+ const c_0 = $[0] !== dispatch;
32
+ let t0;
33
+ if (c_0) {
34
+ t0 = () => {
35
+ dispatch({ kind: "update" });
36
+ };
37
+ $[0] = dispatch;
38
+ $[1] = t0;
39
+ } else {
40
+ t0 = $[1];
41
+ }
42
+ const onUpdate = t0;
43
+ const c_2 = $[2] !== onUpdate;
44
+ let t1;
45
+ if (c_2) {
46
+ t1 = () => {
47
+ onUpdate();
48
+ };
49
+ $[2] = onUpdate;
50
+ $[3] = t1;
51
+ } else {
52
+ t1 = $[3];
53
+ }
54
+ const c_4 = $[4] !== onUpdate;
55
+ let t2;
56
+ if (c_4) {
57
+ t2 = [onUpdate];
58
+ $[4] = onUpdate;
59
+ $[5] = t2;
60
+ } else {
61
+ t2 = $[5];
62
+ }
63
+ useEffect(t1, t2);
64
+ let t3;
65
+ if ($[6] === Symbol.for("react.memo_cache_sentinel")) {
66
+ t3 = <div />;
67
+ $[6] = t3;
68
+ } else {
69
+ t3 = $[6];
70
+ }
71
+ return t3;
72
+}
73
+
74
+```
75
+
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/compiler/useEffect-arg-memoized.js
new
+16
@@ -0,0 +1,16 @@
1
+function Component(props) {
2
+ const dispatch = useDispatch();
3
+ useFreeze(dispatch);
4
+
5
+ // onUpdate should be memoized even though it doesn't
6
+ // flow into the return value
7
+ const onUpdate = () => {
8
+ dispatch({ kind: "update" });
9
+ };
10
+
11
+ useEffect(() => {
12
+ onUpdate();
13
+ }, [onUpdate]);
14
+
15
+ return <div />;
16
+}