[compiler] Fix assignment within for update expression

When converting value blocks from HIR to ReactiveFunction, we have to drop StoreLocal assignments that represent the assignment of the phi, since ReactiveFunction supports compound expressions. These StoreLocals are only present to represent the conditional assignment of the value itself - but it's also possible for the expression to have contained an assignment expression. Before, in trying to strip the first category of StoreLocal we also accidentally stripped the second category. Now we check that the assignment is for a temporary, and don't strip otherwise.

ghstack-source-id: e7759c963b
Pull Request resolved: https://github.com/facebook/react/pull/30067

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 }],
+};