diff --git a/compiler/packages/babel-plugin-react-compiler/src/Optimization/ConstantPropagation.ts b/compiler/packages/babel-plugin-react-compiler/src/Optimization/ConstantPropagation.ts index 7330b63ddce..0bdfc58307a 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Optimization/ConstantPropagation.ts +++ b/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 { + const mutatedGlobals = new Set(); + 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, +): 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, ): 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, +): 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,6 +282,7 @@ function evaluatePhi(phi: Phi, constants: Constants): Constant | null { function evaluateInstruction( constants: Constants, instr: Instruction, + mutatedGlobals: Set, ): Constant | null { const value = instr.value; switch (value.kind) { @@ -246,6 +290,10 @@ function evaluateInstruction( 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': { diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/global-capture-before-assignment.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/global-capture-before-assignment.tsx new file mode 100644 index 00000000000..a92fb6111c4 --- /dev/null +++ b/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
OK
; +} + +export function Main() { + return ( + + + + ); +} + +export const FIXTURE_ENTRYPOINT = { + fn: Main, + params: [{}], +}; +