From ce62ebf501c64e85bbc6341d2a60d30190f3b774 Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Thu, 2 Jan 2025 11:20:50 -0800
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.
---
 .../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 ea132b772aa44..081934042c142 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/MergeConsecutiveBlocks.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/MergeConsecutiveBlocks.ts
@@ -104,6 +104,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: [],
+};