From ab2469c4936f3a9cf1e90297d0b5f4eff9083708 Mon Sep 17 00:00:00 2001
From: Mofei Zhang <feifei0@meta.com>
Date: Tue, 18 Feb 2025 12:26:55 -0700
Subject: [PATCH] [compiler] remove invariant to account for backedges

Fixes https://github.com/facebook/react/issues/32269, see comments for details.

Added test fixture for repro
---
 .../src/Inference/InferReferenceEffects.ts    | 18 ++---
 .../repro-backedge-reference-effect.expect.md | 71 +++++++++++++++++++
 .../repro-backedge-reference-effect.js        | 23 ++++++
 3 files changed, 101 insertions(+), 11 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-backedge-reference-effect.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-backedge-reference-effect.js

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
index 6420e74a2fd20..e6883250705d4 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
@@ -1598,18 +1598,14 @@ function inferBlock(
         break;
       }
       case 'LoadLocal': {
+        /**
+         * Due to backedges in the CFG, we may revisit LoadLocal lvalues
+         * multiple times. Unlike StoreLocal which may reassign to existing
+         * identifiers, LoadLocal always evaluates to store a new temporary.
+         * This means that we should always model LoadLocal as a Capture effect
+         * on the rvalue.
+         */
         const lvalue = instr.lvalue;
-        CompilerError.invariant(
-          !(
-            state.isDefined(lvalue) &&
-            state.kind(lvalue).kind === ValueKind.Context
-          ),
-          {
-            reason:
-              '[InferReferenceEffects] Unexpected LoadLocal with context kind',
-            loc: lvalue.loc,
-          },
-        );
         state.referenceAndRecordEffects(
           freezeActions,
           instrValue.place,
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-backedge-reference-effect.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-backedge-reference-effect.expect.md
new file mode 100644
index 0000000000000..cc33cfc0846ce
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-backedge-reference-effect.expect.md
@@ -0,0 +1,71 @@
+
+## Input
+
+```javascript
+import {Stringify} from 'shared-runtime';
+
+function Foo({userIds}) {
+  return (
+    <Stringify
+      fn={() => {
+        const arr = [];
+
+        for (const selectedUser of userIds) {
+          arr.push(selectedUser);
+        }
+        return arr;
+      }}
+      shouldInvokeFns={true}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{userIds: [1, 2, 3]}],
+  sequentialRenders: [{userIds: [1, 2, 4]}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { Stringify } from "shared-runtime";
+
+function Foo(t0) {
+  const $ = _c(2);
+  const { userIds } = t0;
+  let t1;
+  if ($[0] !== userIds) {
+    t1 = (
+      <Stringify
+        fn={() => {
+          const arr = [];
+          for (const selectedUser of userIds) {
+            arr.push(selectedUser);
+          }
+          return arr;
+        }}
+        shouldInvokeFns={true}
+      />
+    );
+    $[0] = userIds;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ userIds: [1, 2, 3] }],
+  sequentialRenders: [{ userIds: [1, 2, 4] }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"fn":{"kind":"Function","result":[1,2,4]},"shouldInvokeFns":true}</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-backedge-reference-effect.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-backedge-reference-effect.js
new file mode 100644
index 0000000000000..03145445ec5b2
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-backedge-reference-effect.js
@@ -0,0 +1,23 @@
+import {Stringify} from 'shared-runtime';
+
+function Foo({userIds}) {
+  return (
+    <Stringify
+      fn={() => {
+        const arr = [];
+
+        for (const selectedUser of userIds) {
+          arr.push(selectedUser);
+        }
+        return arr;
+      }}
+      shouldInvokeFns={true}
+    />
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{userIds: [1, 2, 3]}],
+  sequentialRenders: [{userIds: [1, 2, 4]}],
+};