[hir] Add a DropMemoCall pass
This drops the memo hook calls from the IR
Sathya Gunasekaran committed
Feb 14, 2023 at 23:07 UTC
fb8f293c321050817cb1c1fdb201cc33954d089b
8 files changed
+165
compiler/forget/src/CompilerPipeline.ts
+4
@@ -14,6 +14,7 @@ import {
14
} from "./HIR";
15
import {
16
analyseFunctions,
17
+ dropMemoCalls,
18
inferMutableRanges,
19
inferReferenceEffects,
20
} from "./Inference";
@@ -65,6 +66,9 @@ export function* run(
66
inferTypes(hir);
67
yield log({ kind: "hir", name: "InferTypes", value: hir });
68
69
+ dropMemoCalls(hir);
70
+ yield log({ kind: "hir", name: "DropMemoCalls", value: hir });
71
+
72
analyseFunctions(hir);
73
yield log({ kind: "hir", name: "AnalyseFunctions", value: hir });
74
compiler/forget/src/HIR/HIR.ts
+4
@@ -745,3 +745,7 @@ export function isObjectType(id: Identifier): boolean {
745
export function isPrimitiveType(id: Identifier): boolean {
746
return id.type.kind === "Primitive";
747
}
748
+
749
+export function isHookType(id: Identifier): boolean {
750
+ return id.type.kind === "Hook";
751
+}
compiler/forget/src/Inference/DropMemoCalls.ts
new
+60
@@ -0,0 +1,60 @@
1
+import invariant from "invariant";
2
+import {
3
+ Effect,
4
+ HIRFunction,
5
+ HookType,
6
+ InstructionValue,
7
+ isHookType,
8
+} from "../HIR";
9
+
10
+export default function (func: HIRFunction) {
11
+ for (const [_, block] of func.body.blocks) {
12
+ for (const instr of block.instructions) {
13
+ switch (instr.value.kind) {
14
+ case "CallExpression": {
15
+ if (isHookType(instr.value.callee.identifier)) {
16
+ const name = (instr.value.callee.identifier.type as HookType).name;
17
+ if (name === "useMemo") {
18
+ const [fn] = instr.value.args;
19
+
20
+ // TODO(gsn): Consider inlining the function passed to useMemo,
21
+ // rather than just calling it directly.
22
+ //
23
+ // Replace the hook callee with the fn arg.
24
+ //
25
+ // before:
26
+ // foo = Call useMemo$2($9, $10)
27
+ //
28
+ // after:
29
+ // foo = Call $9()
30
+ instr.value = {
31
+ kind: "CallExpression",
32
+ callee: fn,
33
+ // Drop the args, including the deps array which DCE will remove
34
+ // later.
35
+ args: [],
36
+ loc: instr.value.loc,
37
+ };
38
+ } else if (name === "useCallback") {
39
+ const [fn] = instr.value.args;
40
+
41
+ // Instead of a Call, just alias the callback directly.
42
+ //
43
+ // before:
44
+ // foo = Call useCallback$8($19)
45
+ //
46
+ // after:
47
+ // foo = $19
48
+ instr.value = {
49
+ kind: "Identifier",
50
+ identifier: fn.identifier,
51
+ effect: Effect.Unknown,
52
+ loc: instr.value.loc,
53
+ };
54
+ }
55
+ }
56
+ }
57
+ }
58
+ }
59
+ }
60
+}
compiler/forget/src/Inference/index.ts
+1
@@ -5,6 +5,7 @@
5
* LICENSE file in the root directory of this source tree.
6
*/
7
8
+export { default as dropMemoCalls } from "./DropMemoCalls";
9
export { inferMutableRanges } from "./InferMutableRanges";
10
export { default as analyseFunctions } from "./AnalyseFunctions";
11
export { default as inferReferenceEffects } from "./InferReferenceEffects";
compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.expect.md
new
+47
@@ -0,0 +1,47 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+function component() {
6
+ const [count, setCount] = useState(0);
7
+ const increment = useCallback(() => setCount(count + 1));
8
+
9
+ return <Foo onClick={increment}></Foo>;
10
+}
11
+
12
+```
13
+
14
+## Code
15
+
16
+```javascript
17
+function component() {
18
+ const $ = React.unstable_useMemoCache();
19
+ const t2 = useState(0);
20
+ const count = t2[0];
21
+ const setCount = t2[1];
22
+ const c_0 = $[0] !== setCount;
23
+ const c_1 = $[1] !== count;
24
+ let t0;
25
+ if (c_0 || c_1) {
26
+ t0 = () => setCount(count + 1);
27
+ $[0] = setCount;
28
+ $[1] = count;
29
+ $[2] = t0;
30
+ } else {
31
+ t0 = $[2];
32
+ }
33
+ const increment = t0;
34
+ const c_3 = $[3] !== increment;
35
+ let t1;
36
+ if (c_3) {
37
+ t1 = <Foo onClick={increment}></Foo>;
38
+ $[3] = increment;
39
+ $[4] = t1;
40
+ } else {
41
+ t1 = $[4];
42
+ }
43
+ return t1;
44
+}
45
+
46
+```
47
+
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/hir/use-callback-simple.js
new
+6
@@ -0,0 +1,6 @@
1
+function component() {
2
+ const [count, setCount] = useState(0);
3
+ const increment = useCallback(() => setCount(count + 1));
4
+
5
+ return <Foo onClick={increment}></Foo>;
6
+}
compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.expect.md
new
+39
@@ -0,0 +1,39 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+function component(a) {
6
+ let x = useMemo(() => [a], [a]);
7
+ return <Foo x={x}></Foo>;
8
+}
9
+
10
+```
11
+
12
+## Code
13
+
14
+```javascript
15
+function component(a) {
16
+ const $ = React.unstable_useMemoCache();
17
+ const c_0 = $[0] !== a;
18
+ let x;
19
+ if (c_0) {
20
+ x = (() => [a])();
21
+ $[0] = a;
22
+ $[1] = x;
23
+ } else {
24
+ x = $[1];
25
+ }
26
+ const c_2 = $[2] !== x;
27
+ let t0;
28
+ if (c_2) {
29
+ t0 = <Foo x={x}></Foo>;
30
+ $[2] = x;
31
+ $[3] = t0;
32
+ } else {
33
+ t0 = $[3];
34
+ }
35
+ return t0;
36
+}
37
+
38
+```
39
+
\ No newline at end of file
compiler/forget/src/__tests__/fixtures/hir/useMemo-simple.js
new
+4
@@ -0,0 +1,4 @@
1
+function component(a) {
2
+ let x = useMemo(() => [a], [a]);
3
+ return <Foo x={x}></Foo>;
4
+}