[HIR] Followup to pruneUnusedLabelsHIR (#2866)
Followup to https://github.com/facebook/react-forget/pull/2866 ghstack-source-id: 087d5da53787cb7cff6495b6a791326ff8a952b4 Pull Request resolved: https://github.com/facebook/react-forget/pull/2896
Mofei Zhang committed
Apr 23, 2024 at 19:58 UTC
f22dd62ab8e2056b89fb29193d0ed30a98fd5ff4
3 files changed
+90
-3
compiler/packages/babel-plugin-react-forget/src/HIR/PruneUnusedLabelsHIR.ts
+5
-3
@@ -7,7 +7,7 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
7
next: BlockId;
8
fallthrough: BlockId;
9
}> = [];
10
-
10
+ const rewrites: Map<BlockId, BlockId> = new Map();
11
for (const [blockId, block] of fn.body.blocks) {
12
const terminal = block.terminal;
13
if (terminal.kind === "label") {
@@ -32,10 +32,11 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
32
}
33
34
for (const {
35
- label: labelId,
35
+ label: originalLabelId,
36
next: nextId,
37
fallthrough: fallthroughId,
38
} of merged) {
39
+ const labelId = rewrites.get(originalLabelId) ?? originalLabelId;
40
const label = fn.body.blocks.get(labelId)!;
41
const next = fn.body.blocks.get(nextId)!;
42
const fallthrough = fn.body.blocks.get(fallthroughId)!;
@@ -52,7 +53,7 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
53
CompilerError.invariant(
54
next.preds.size === 1 &&
55
fallthrough.preds.size === 1 &&
55
- next.preds.has(labelId) &&
56
+ next.preds.has(originalLabelId) &&
57
fallthrough.preds.has(nextId),
58
{
59
reason: "Unexpected block predecessors when merging label blocks",
@@ -64,5 +65,6 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
65
label.terminal = fallthrough.terminal;
66
fn.body.blocks.delete(nextId);
67
fn.body.blocks.delete(fallthroughId);
68
+ rewrites.set(fallthroughId, labelId);
69
}
70
}
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.expect.md
new
+72
@@ -0,0 +1,72 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+import { useMemo } from "react";
6
+import { identity } from "shared-runtime";
7
+
8
+function useHook({ a, b }) {
9
+ const valA = useMemo(() => identity({ a }), [a]);
10
+ const valB = useMemo(() => identity([b]), [b]);
11
+ return [valA, valB];
12
+}
13
+
14
+export const FIXTURE_ENTRYPOINT = {
15
+ fn: useHook,
16
+ params: [{ a: 2, b: 3 }],
17
+};
18
+
19
+```
20
+
21
+## Code
22
+
23
+```javascript
24
+import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
25
+import { identity } from "shared-runtime";
26
+
27
+function useHook(t0) {
28
+ const $ = useMemoCache(7);
29
+ const { a, b } = t0;
30
+ let t1;
31
+ let t2;
32
+ if ($[0] !== a) {
33
+ t2 = identity({ a });
34
+ $[0] = a;
35
+ $[1] = t2;
36
+ } else {
37
+ t2 = $[1];
38
+ }
39
+ t1 = t2;
40
+ const valA = t1;
41
+ let t3;
42
+ let t4;
43
+ if ($[2] !== b) {
44
+ t4 = identity([b]);
45
+ $[2] = b;
46
+ $[3] = t4;
47
+ } else {
48
+ t4 = $[3];
49
+ }
50
+ t3 = t4;
51
+ const valB = t3;
52
+ let t5;
53
+ if ($[4] !== valA || $[5] !== valB) {
54
+ t5 = [valA, valB];
55
+ $[4] = valA;
56
+ $[5] = valB;
57
+ $[6] = t5;
58
+ } else {
59
+ t5 = $[6];
60
+ }
61
+ return t5;
62
+}
63
+
64
+export const FIXTURE_ENTRYPOINT = {
65
+ fn: useHook,
66
+ params: [{ a: 2, b: 3 }],
67
+};
68
+
69
+```
70
+
71
+### Eval output
72
+(kind: ok) [{"a":2},[3]]
\ No newline at end of file
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.ts
new
+13
@@ -0,0 +1,13 @@
1
+import { useMemo } from "react";
2
+import { identity } from "shared-runtime";
3
+
4
+function useHook({ a, b }) {
5
+ const valA = useMemo(() => identity({ a }), [a]);
6
+ const valB = useMemo(() => identity([b]), [b]);
7
+ return [valA, valB];
8
+}
9
+
10
+export const FIXTURE_ENTRYPOINT = {
11
+ fn: useHook,
12
+ params: [{ a: 2, b: 3 }],
13
+};