[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

compiler/packages/babel-plugin-react-forget/src/HIR/PruneUnusedLabelsHIR.ts

@@ -7,7 +7,7 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
     next: BlockId;
     fallthrough: BlockId;
   }> = [];
-
+  const rewrites: Map<BlockId, BlockId> = new Map();
   for (const [blockId, block] of fn.body.blocks) {
     const terminal = block.terminal;
     if (terminal.kind === "label") {
@@ -32,10 +32,11 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
   }
 
   for (const {
-    label: labelId,
+    label: originalLabelId,
     next: nextId,
     fallthrough: fallthroughId,
   } of merged) {
+    const labelId = rewrites.get(originalLabelId) ?? originalLabelId;
     const label = fn.body.blocks.get(labelId)!;
     const next = fn.body.blocks.get(nextId)!;
     const fallthrough = fn.body.blocks.get(fallthroughId)!;
@@ -52,7 +53,7 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
     CompilerError.invariant(
       next.preds.size === 1 &&
         fallthrough.preds.size === 1 &&
-        next.preds.has(labelId) &&
+        next.preds.has(originalLabelId) &&
         fallthrough.preds.has(nextId),
       {
         reason: "Unexpected block predecessors when merging label blocks",
@@ -64,5 +65,6 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
     label.terminal = fallthrough.terminal;
     fn.body.blocks.delete(nextId);
     fn.body.blocks.delete(fallthroughId);
+    rewrites.set(fallthroughId, labelId);
   }
 }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.expect.md

@@ -0,0 +1,72 @@
+
+## Input
+
+```javascript
+import { useMemo } from "react";
+import { identity } from "shared-runtime";
+
+function useHook({ a, b }) {
+  const valA = useMemo(() => identity({ a }), [a]);
+  const valB = useMemo(() => identity([b]), [b]);
+  return [valA, valB];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useHook,
+  params: [{ a: 2, b: 3 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
+import { identity } from "shared-runtime";
+
+function useHook(t0) {
+  const $ = useMemoCache(7);
+  const { a, b } = t0;
+  let t1;
+  let t2;
+  if ($[0] !== a) {
+    t2 = identity({ a });
+    $[0] = a;
+    $[1] = t2;
+  } else {
+    t2 = $[1];
+  }
+  t1 = t2;
+  const valA = t1;
+  let t3;
+  let t4;
+  if ($[2] !== b) {
+    t4 = identity([b]);
+    $[2] = b;
+    $[3] = t4;
+  } else {
+    t4 = $[3];
+  }
+  t3 = t4;
+  const valB = t3;
+  let t5;
+  if ($[4] !== valA || $[5] !== valB) {
+    t5 = [valA, valB];
+    $[4] = valA;
+    $[5] = valB;
+    $[6] = t5;
+  } else {
+    t5 = $[6];
+  }
+  return t5;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useHook,
+  params: [{ a: 2, b: 3 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) [{"a":2},[3]]

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/consecutive-use-memo.ts

@@ -0,0 +1,13 @@
+import { useMemo } from "react";
+import { identity } from "shared-runtime";
+
+function useHook({ a, b }) {
+  const valA = useMemo(() => identity({ a }), [a]);
+  const valB = useMemo(() => identity([b]), [b]);
+  return [valA, valB];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useHook,
+  params: [{ a: 2, b: 3 }],
+};