fix(compiler): preserve runNumber across StoreGlobal mutations (#34899)

compiler/packages/babel-plugin-react-compiler/src/Optimization/ConstantPropagation.ts

@@ -57,12 +57,37 @@ import {eliminateRedundantPhi} from '../SSA';
  */
 export function constantPropagation(fn: HIRFunction): void {
   const constants: Constants = new Map();
-  constantPropagationImpl(fn, constants);
+  const mutatedGlobals = collectMutatedGlobals(fn);
+  constantPropagationImpl(fn, constants, mutatedGlobals);
 }
 
-function constantPropagationImpl(fn: HIRFunction, constants: Constants): void {
+/**
+ * Collect all globals that are reassigned anywhere in the function.
+ * This includes StoreGlobal instructions generated by compound assignments.
+ */
+function collectMutatedGlobals(fn: HIRFunction): Set<string> {
+  const mutatedGlobals = new Set<string>();
+  for (const [, block] of fn.body.blocks) {
+    for (const instr of block.instructions) {
+      if (instr.value.kind === 'StoreGlobal') {
+        mutatedGlobals.add(instr.value.name);
+      }
+    }
+  }
+  return mutatedGlobals;
+}
+
+function constantPropagationImpl(
+  fn: HIRFunction,
+  constants: Constants,
+  mutatedGlobals: Set<string>,
+): void {
   while (true) {
-    const haveTerminalsChanged = applyConstantPropagation(fn, constants);
+    const haveTerminalsChanged = applyConstantPropagation(
+      fn,
+      constants,
+      mutatedGlobals,
+    );
     if (!haveTerminalsChanged) {
       break;
     }
@@ -106,6 +131,7 @@ function constantPropagationImpl(fn: HIRFunction, constants: Constants): void {
 function applyConstantPropagation(
   fn: HIRFunction,
   constants: Constants,
+  mutatedGlobals: Set<string>,
 ): boolean {
   let hasChanges = false;
   for (const [, block] of fn.body.blocks) {
@@ -115,9 +141,13 @@ function applyConstantPropagation(
      * phi values for blocks that have a back-edge.
      */
     for (const phi of block.phis) {
-      let value = evaluatePhi(phi, constants);
+      let value = evaluatePhi(phi, constants, mutatedGlobals);
       if (value !== null) {
+        const previousValue = constants.get(phi.place.identifier.id);
         constants.set(phi.place.identifier.id, value);
+        if (previousValue !== value) {
+          hasChanges = true;
+        }
       }
     }
 
@@ -130,9 +160,13 @@ function applyConstantPropagation(
         continue;
       }
       const instr = block.instructions[i]!;
-      const value = evaluateInstruction(constants, instr);
+      const previousConstant = constants.get(instr.lvalue.identifier.id);
+      const value = evaluateInstruction(constants, instr, mutatedGlobals);
       if (value !== null) {
         constants.set(instr.lvalue.identifier.id, value);
+        if (previousConstant !== value) {
+          hasChanges = true;
+        }
       }
     }
 
@@ -164,7 +198,11 @@ function applyConstantPropagation(
   return hasChanges;
 }
 
-function evaluatePhi(phi: Phi, constants: Constants): Constant | null {
+function evaluatePhi(
+  phi: Phi,
+  constants: Constants,
+  mutatedGlobals: Set<string>,
+): Constant | null {
   let value: Constant | null = null;
   for (const [, operand] of phi.operands) {
     const operandValue = constants.get(operand.identifier.id) ?? null;
@@ -226,6 +264,11 @@ function evaluatePhi(phi: Phi, constants: Constants): Constant | null {
         if (operandValue.binding.name !== value.binding.name) {
           return null;
         }
+
+        // If this global is mutated anywhere, don't constant propagate it
+        if (mutatedGlobals.has(operandValue.binding.name)) {
+          return null;
+        }
         break;
       }
       default:
@@ -239,13 +282,18 @@ function evaluatePhi(phi: Phi, constants: Constants): Constant | null {
 function evaluateInstruction(
   constants: Constants,
   instr: Instruction,
+  mutatedGlobals: Set<string>,
 ): Constant | null {
   const value = instr.value;
   switch (value.kind) {
     case 'Primitive': {
       return value;
     }
     case 'LoadGlobal': {
+      // If this global is mutated anywhere, don't treat it as a constant
+      if (mutatedGlobals.has(value.binding.name)) {
+        return null;
+      }
       return value;
     }
     case 'ComputedLoad': {
@@ -606,7 +654,17 @@ function evaluateInstruction(
     }
     case 'ObjectMethod':
     case 'FunctionExpression': {
-      constantPropagationImpl(value.loweredFunc.func, constants);
+      // Collect mutated globals from nested function and merge with parent
+      const nestedMutatedGlobals = collectMutatedGlobals(value.loweredFunc.func);
+      const combinedMutatedGlobals = new Set([
+        ...mutatedGlobals,
+        ...nestedMutatedGlobals,
+      ]);
+      constantPropagationImpl(
+        value.loweredFunc.func,
+        constants,
+        combinedMutatedGlobals,
+      );
       return null;
     }
     case 'StartMemoize': {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/global-capture-before-assignment.tsx

@@ -0,0 +1,29 @@
+import {useEffect, StrictMode} from 'react';
+
+let i = 0;
+
+function App() {
+  useEffect(() => {
+    const runNumber = i;
+    console.log('effect run', runNumber);
+    i += 1;
+    return () => {
+      console.log('cleanup run', runNumber);
+    };
+  }, []);
+  return <div>OK</div>;
+}
+
+export function Main() {
+  return (
+    <StrictMode>
+      <App />
+    </StrictMode>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Main,
+  params: [{}],
+};
+