[compiler] Fix assignment within for update expression

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveFunction.ts

@@ -921,14 +921,26 @@ class Driver {
         });
       } else if (defaultBlock.instructions.length === 1) {
         const instr = defaultBlock.instructions[0]!;
-        let place: Place = instr.lvalue!;
+        let place: Place = instr.lvalue;
         let value: ReactiveValue = instr.value;
-        if (instr.value.kind === "StoreLocal") {
-          place = instr.value.lvalue.place;
+        if (
+          /*
+           * Value blocks generally end in a StoreLocal to assign the value of the
+           * expression for this branch. These StoreLocal instructions can be pruned,
+           * since we represent the value blocks as a compund value in ReactiveFunction
+           * (no phis). However, it's also possible to have a value block that ends in
+           * an AssignmentExpression, which we need to keep. So we only prune
+           * StoreLocal for temporaries — any named/promoted values must be used
+           * elsewhere and aren't safe to prune.
+           */
+          value.kind === "StoreLocal" &&
+          value.lvalue.place.identifier.name === null
+        ) {
+          place = value.lvalue.place;
           value = {
             kind: "LoadLocal",
-            place: instr.value.value,
-            loc: instr.value.value.loc,
+            place: value.value,
+            loc: value.value.loc,
           };
         }
         return {
@@ -939,14 +951,26 @@ class Driver {
         };
       } else {
         const instr = defaultBlock.instructions.at(-1)!;
-        let place: Place = instr.lvalue!;
+        let place: Place = instr.lvalue;
         let value: ReactiveValue = instr.value;
-        if (instr.value.kind === "StoreLocal") {
-          place = instr.value.lvalue.place;
+        if (
+          /*
+           * Value blocks generally end in a StoreLocal to assign the value of the
+           * expression for this branch. These StoreLocal instructions can be pruned,
+           * since we represent the value blocks as a compund value in ReactiveFunction
+           * (no phis). However, it's also possible to have a value block that ends in
+           * an AssignmentExpression, which we need to keep. So we only prune
+           * StoreLocal for temporaries — any named/promoted values must be used
+           * elsewhere and aren't safe to prune.
+           */
+          value.kind === "StoreLocal" &&
+          value.lvalue.place.identifier.name === null
+        ) {
+          place = value.lvalue.place;
           value = {
             kind: "LoadLocal",
-            place: instr.value.value,
-            loc: instr.value.value.loc,
+            place: value.value,
+            loc: value.value.loc,
           };
         }
         const sequence: ReactiveSequenceValue = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/for-with-assignment-as-update.expect.md

@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  let x = props.init;
+  for (let i = 0; i < 100; i = i + 1) {
+    x += i;
+  }
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ init: 0 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component(props) {
+  const $ = _c(2);
+  let x = props.init;
+  for (let i = 0; i < 100; i = i + 1) {
+    x = x + i;
+  }
+  let t0;
+  if ($[0] !== x) {
+    t0 = [x];
+    $[0] = x;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ init: 0 }],
+};
+
+```
+
+### Eval output
+(kind: ok) [4950]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/for-with-assignment-as-update.js

@@ -0,0 +1,12 @@
+function Component(props) {
+  let x = props.init;
+  for (let i = 0; i < 100; i = i + 1) {
+    x += i;
+  }
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ init: 0 }],
+};