From f8026d06f561494f1dd46f2f199afd2fdcf09998 Mon Sep 17 00:00:00 2001
From: Mofei Zhang <feifei0@meta.com>
Date: Mon, 24 Jun 2024 19:19:42 -0400
Subject: [PATCH] [compiler][hir] Correctly remove non-existent terminal preds
 when pruning labels

[ghstack-poisoned]
---
 .../src/Entrypoint/Pipeline.ts                |  3 +++
 ...sExist.ts => AssertTerminalBlocksExist.ts} | 22 ++++++++++++++++++-
 .../src/HIR/PruneUnusedLabelsHIR.ts           | 10 +++++++++
 .../src/HIR/index.ts                          |  5 ++++-
 4 files changed, 38 insertions(+), 2 deletions(-)
 rename compiler/packages/babel-plugin-react-compiler/src/HIR/{AssertTerminalSuccessorsExist.ts => AssertTerminalBlocksExist.ts} (54%)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
index 863eca5dcf3..2e7613f0a2a 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts
@@ -13,6 +13,7 @@ import {
   HIRFunction,
   ReactiveFunction,
   assertConsistentIdentifiers,
+  assertTerminalPredsExist,
   assertTerminalSuccessorsExist,
   assertValidBlockNesting,
   assertValidMutableRanges,
@@ -303,6 +304,8 @@ function* runWithEnvironment(
       name: "FlattenScopesWithHooksOrUseHIR",
       value: hir,
     });
+    assertTerminalSuccessorsExist(hir);
+    assertTerminalPredsExist(hir);
   }
 
   const reactiveFunction = buildReactiveFunction(hir);
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalSuccessorsExist.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalBlocksExist.ts
similarity index 54%
rename from compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalSuccessorsExist.ts
rename to compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalBlocksExist.ts
index 493ff54c03f..e5dbb1b8dbf 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalSuccessorsExist.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/AssertTerminalBlocksExist.ts
@@ -8,7 +8,7 @@
 import { CompilerError } from "../CompilerError";
 import { GeneratedSource, HIRFunction } from "./HIR";
 import { printTerminal } from "./PrintHIR";
-import { mapTerminalSuccessors } from "./visitors";
+import { eachTerminalSuccessor, mapTerminalSuccessors } from "./visitors";
 
 export function assertTerminalSuccessorsExist(fn: HIRFunction): void {
   for (const [, block] of fn.body.blocks) {
@@ -25,3 +25,23 @@ export function assertTerminalSuccessorsExist(fn: HIRFunction): void {
     });
   }
 }
+
+export function assertTerminalPredsExist(fn: HIRFunction): void {
+  for (const [, block] of fn.body.blocks) {
+    for (const pred of block.preds) {
+      const predBlock = fn.body.blocks.get(pred);
+      CompilerError.invariant(predBlock != null, {
+        reason: "Expected predecessor block to exist",
+        description: `Block ${block.id} references non-existent ${pred}`,
+        loc: GeneratedSource,
+      });
+      CompilerError.invariant(
+        [...eachTerminalSuccessor(predBlock.terminal)].includes(block.id),
+        {
+          reason: "Terminal successor does not reference correct predecessor",
+          loc: GeneratedSource,
+        }
+      );
+    }
+  }
+}
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/PruneUnusedLabelsHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/PruneUnusedLabelsHIR.ts
index 10714d5d879..fd1ea1b14b1 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/PruneUnusedLabelsHIR.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/PruneUnusedLabelsHIR.ts
@@ -67,4 +67,14 @@ export function pruneUnusedLabelsHIR(fn: HIRFunction): void {
     fn.body.blocks.delete(fallthroughId);
     rewrites.set(fallthroughId, labelId);
   }
+
+  for (const [_, block] of fn.body.blocks) {
+    for (const pred of [...block.preds]) {
+      const rewritten = rewrites.get(pred);
+      if (rewritten != null) {
+        block.preds.delete(pred);
+        block.preds.add(rewritten);
+      }
+    }
+  }
 }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/index.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/index.ts
index b17d3a09d08..bef4f1cb95d 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/index.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/index.ts
@@ -6,7 +6,10 @@
  */
 
 export { assertConsistentIdentifiers } from "./AssertConsistentIdentifiers";
-export { assertTerminalSuccessorsExist } from "./AssertTerminalSuccessorsExist";
+export {
+  assertTerminalSuccessorsExist,
+  assertTerminalPredsExist,
+} from "./AssertTerminalBlocksExist";
 export { assertValidBlockNesting } from "./AssertValidBlockNesting";
 export { assertValidMutableRanges } from "./AssertValidMutableRanges";
 export { lower } from "./BuildHIR";