From 82b8f612de10d687c78a56adefa1aa182df5bfa8 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 7 Jul 2025 14:59:27 -0700
Subject: [PATCH] [compiler] Fix for consecutive DCE'd branches with phis

This is an optimized version of @asmjmp0's fix in https://github.com/facebook/react/pull/31940. When we merge consecutive blocks we need to take care to rewrite later phis whose operands will now be different blocks due to merging. Rather than iterate all the blocks on each merge as in #31940, we can do a single iteration over all the phis at the end to fix them up.

Note: this is a redo of #31959
---
 .../src/HIR/MergeConsecutiveBlocks.ts         | 11 ++++
 ...ssing-phi-after-dce-merge-scopes.expect.md | 52 +++++++++++++++++++
 ...epro-missing-phi-after-dce-merge-scopes.js | 20 +++++++
 3 files changed, 83 insertions(+)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-missing-phi-after-dce-merge-scopes.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-missing-phi-after-dce-merge-scopes.js

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/MergeConsecutiveBlocks.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/MergeConsecutiveBlocks.ts
index 3d6ae4e6b211d..881e4e93ffcdb 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/MergeConsecutiveBlocks.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/MergeConsecutiveBlocks.ts
@@ -107,6 +107,17 @@ export function mergeConsecutiveBlocks(fn: HIRFunction): void {
     merged.merge(block.id, predecessorId);
     fn.body.blocks.delete(block.id);
   }
+  for (const [, block] of fn.body.blocks) {
+    for (const phi of block.phis) {
+      for (const [predecessorId, operand] of phi.operands) {
+        const mapped = merged.get(predecessorId);
+        if (mapped !== predecessorId) {
+          phi.operands.delete(predecessorId);
+          phi.operands.set(mapped, operand);
+        }
+      }
+    }
+  }
   markPredecessors(fn.body);
   for (const [, {terminal}] of fn.body.blocks) {
     if (terminalHasFallthrough(terminal)) {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-missing-phi-after-dce-merge-scopes.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-missing-phi-after-dce-merge-scopes.expect.md
new file mode 100644
index 0000000000000..da7220c73c9b7
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-missing-phi-after-dce-merge-scopes.expect.md
@@ -0,0 +1,52 @@
+
+## Input
+
+```javascript
+function Component() {
+  let v3, v4, acc;
+  v3 = false;
+  v4 = v3;
+  acc = v3;
+  if (acc) {
+    acc = true;
+    v3 = acc;
+  }
+  if (acc) {
+    v3 = v4;
+  }
+  v4 = v3;
+  return [acc, v3, v4];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = [false, false, false];
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) [false,false,false]
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-missing-phi-after-dce-merge-scopes.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-missing-phi-after-dce-merge-scopes.js
new file mode 100644
index 0000000000000..298ccbebaf27d
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-missing-phi-after-dce-merge-scopes.js
@@ -0,0 +1,20 @@
+function Component() {
+  let v3, v4, acc;
+  v3 = false;
+  v4 = v3;
+  acc = v3;
+  if (acc) {
+    acc = true;
+    v3 = acc;
+  }
+  if (acc) {
+    v3 = v4;
+  }
+  v4 = v3;
+  return [acc, v3, v4];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};