From 1671142ea8322410868aa143dbfef6a98fcda238 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 28 Aug 2025 14:54:27 -0700 Subject: [PATCH 1/3] [compiler] InferEffects uses effects as value keys In InferReferenceEffects we used `InstructionValue` as the key to represent values, since each time we process an instruction this object will be the same. However this was always a bit of a hack, and in the new model and InferMutationAliasingEffects we can instead use the (creation) effect as the stable value. This avoids an extra layer of memoization since the effects are already interned anyway. --- .../Inference/InferMutationAliasingEffects.ts | 147 +++++++----------- 1 file changed, 52 insertions(+), 95 deletions(-) diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts index 8dd79409ce15..77cb972ee516 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts @@ -25,7 +25,6 @@ import { IdentifierId, Instruction, InstructionKind, - InstructionValue, isArrayType, isJsxType, isMapType, @@ -57,7 +56,6 @@ import { printAliasingSignature, printIdentifier, printInstruction, - printInstructionValue, printPlace, printSourceLocation, } from '../HIR/PrintHIR'; @@ -108,11 +106,11 @@ export function inferMutationAliasingEffects( const statesByBlock: Map = new Map(); for (const ref of fn.context) { - // TODO: using InstructionValue as a bit of a hack, but it's pragmatic - const value: InstructionValue = { - kind: 'ObjectExpression', - properties: [], - loc: ref.loc, + const value: AliasingEffect = { + kind: 'Create', + into: ref, + value: ValueKind.Context, + reason: ValueReason.Other, }; initialState.initialize(value, { kind: ValueKind.Context, @@ -145,10 +143,11 @@ export function inferMutationAliasingEffects( } if (ref != null) { const place = ref.kind === 'Identifier' ? ref : ref.place; - const value: InstructionValue = { - kind: 'ObjectExpression', - properties: [], - loc: place.loc, + const value: AliasingEffect = { + kind: 'Create', + into: place, + value: ValueKind.Mutable, + reason: ValueReason.Other, }; initialState.initialize(value, { kind: ValueKind.Mutable, @@ -265,8 +264,6 @@ function findHoistedContextDeclarations( class Context { internedEffects: Map = new Map(); instructionSignatureCache: Map = new Map(); - effectInstructionValueCache: Map = - new Map(); applySignatureCache: Map< AliasingSignature, Map | null> @@ -318,10 +315,11 @@ function inferParam( paramKind: AbstractValue, ): void { const place = param.kind === 'Identifier' ? param : param.place; - const value: InstructionValue = { - kind: 'Primitive', - loc: place.loc, - value: undefined, + const value: AliasingEffect = { + kind: 'Create', + into: place, + value: paramKind.kind, + reason: ValueReason.Other, }; initialState.initialize(value, paramKind); initialState.define(place, value); @@ -542,20 +540,11 @@ function applyEffect( }); initialized.add(effect.into.identifier.id); - let value = context.effectInstructionValueCache.get(effect); - if (value == null) { - value = { - kind: 'ObjectExpression', - properties: [], - loc: effect.into.loc, - }; - context.effectInstructionValueCache.set(effect, value); - } - state.initialize(value, { + state.initialize(effect, { kind: effect.value, reason: new Set([effect.reason]), }); - state.define(effect.into, value); + state.define(effect.into, effect); effects.push(effect); break; } @@ -582,20 +571,11 @@ function applyEffect( initialized.add(effect.into.identifier.id); const fromValue = state.kind(effect.from); - let value = context.effectInstructionValueCache.get(effect); - if (value == null) { - value = { - kind: 'ObjectExpression', - properties: [], - loc: effect.into.loc, - }; - context.effectInstructionValueCache.set(effect, value); - } - state.initialize(value, { + state.initialize(effect, { kind: fromValue.kind, reason: new Set(fromValue.reason), }); - state.define(effect.into, value); + state.define(effect.into, effect); switch (fromValue.kind) { case ValueKind.Primitive: case ValueKind.Global: { @@ -683,11 +663,11 @@ function applyEffect( operand.effect = Effect.Read; } } - state.initialize(effect.function, { + state.initialize(effect, { kind: isMutable ? ValueKind.Mutable : ValueKind.Frozen, reason: new Set([]), }); - state.define(effect.into, effect.function); + state.define(effect.into, effect); for (const capture of effect.captures) { applyEffect( context, @@ -793,38 +773,20 @@ function applyEffect( initialized, effects, ); - let value = context.effectInstructionValueCache.get(effect); - if (value == null) { - value = { - kind: 'Primitive', - value: undefined, - loc: effect.from.loc, - }; - context.effectInstructionValueCache.set(effect, value); - } - state.initialize(value, { + state.initialize(effect, { kind: fromKind, reason: new Set(fromValue.reason), }); - state.define(effect.into, value); + state.define(effect.into, effect); break; } case ValueKind.Global: case ValueKind.Primitive: { - let value = context.effectInstructionValueCache.get(effect); - if (value == null) { - value = { - kind: 'Primitive', - value: undefined, - loc: effect.from.loc, - }; - context.effectInstructionValueCache.set(effect, value); - } - state.initialize(value, { + state.initialize(effect, { kind: fromKind, reason: new Set(fromValue.reason), }); - state.define(effect.into, value); + state.define(effect.into, effect); break; } default: { @@ -839,14 +801,15 @@ function applyEffect( const functionValues = state.values(effect.function); if ( functionValues.length === 1 && - functionValues[0].kind === 'FunctionExpression' && - functionValues[0].loweredFunc.func.aliasingEffects != null + functionValues[0].kind === 'CreateFunction' && + functionValues[0].function.kind === 'FunctionExpression' && + functionValues[0].function.loweredFunc.func.aliasingEffects != null ) { /* * We're calling a locally declared function, we already know it's effects! * We just have to substitute in the args for the params */ - const functionExpr = functionValues[0]; + const functionExpr = functionValues[0].function; let signature = context.functionSignatureCache.get(functionExpr); if (signature == null) { signature = buildSignatureFromFunctionExpression( @@ -1136,19 +1099,19 @@ class InferenceState { #isFunctionExpression: boolean; // The kind of each value, based on its allocation site - #values: Map; + #values: Map; /* * The set of values pointed to by each identifier. This is a set * to accomodate phi points (where a variable may have different * values from different control flow paths). */ - #variables: Map>; + #variables: Map>; constructor( env: Environment, isFunctionExpression: boolean, - values: Map, - variables: Map>, + values: Map, + variables: Map>, ) { this.env = env; this.#isFunctionExpression = isFunctionExpression; @@ -1168,18 +1131,11 @@ class InferenceState { } // (Re)initializes a @param value with its default @param kind. - initialize(value: InstructionValue, kind: AbstractValue): void { - CompilerError.invariant(value.kind !== 'LoadLocal', { - reason: - '[InferMutationAliasingEffects] Expected all top-level identifiers to be defined as variables, not values', - description: null, - loc: value.loc, - suggestions: null, - }); + initialize(value: AliasingEffect, kind: AbstractValue): void { this.#values.set(value, kind); } - values(place: Place): Array { + values(place: Place): Array { const values = this.#variables.get(place.identifier.id); CompilerError.invariant(values != null, { reason: `[InferMutationAliasingEffects] Expected value kind to be initialized`, @@ -1242,13 +1198,13 @@ class InferenceState { } // Defines (initializing or updating) a variable with a specific kind of value. - define(place: Place, value: InstructionValue): void { + define(place: Place, value: AliasingEffect): void { CompilerError.invariant(this.#values.has(value), { reason: `[InferMutationAliasingEffects] Expected value to be initialized at '${printSourceLocation( - value.loc, + place.loc, )}'`, - description: printInstructionValue(value), - loc: value.loc, + description: printAliasingEffect(value), + loc: place.loc, suggestions: null, }); this.#variables.set(place.identifier.id, new Set([value])); @@ -1288,17 +1244,17 @@ class InferenceState { } } - freezeValue(value: InstructionValue, reason: ValueReason): void { + freezeValue(value: AliasingEffect, reason: ValueReason): void { this.#values.set(value, { kind: ValueKind.Frozen, reason: new Set([reason]), }); if ( - value.kind === 'FunctionExpression' && + value.kind === 'CreateFunction' && (this.env.config.enablePreserveExistingMemoizationGuarantees || this.env.config.enableTransitivelyFreezeFunctionExpressions) ) { - for (const place of value.loweredFunc.func.context) { + for (const place of value.function.loweredFunc.func.context) { this.freeze(place, reason); } } @@ -1373,8 +1329,8 @@ class InferenceState { * termination. */ merge(other: InferenceState): InferenceState | null { - let nextValues: Map | null = null; - let nextVariables: Map> | null = null; + let nextValues: Map | null = null; + let nextVariables: Map> | null = null; for (const [id, thisValue] of this.#values) { const otherValue = other.#values.get(id); @@ -1398,7 +1354,7 @@ class InferenceState { for (const [id, thisValues] of this.#variables) { const otherValues = other.#variables.get(id); if (otherValues !== undefined) { - let mergedValues: Set | null = null; + let mergedValues: Set | null = null; for (const otherValue of otherValues) { if (!thisValues.has(otherValue)) { mergedValues = mergedValues ?? new Set(thisValues); @@ -1451,8 +1407,8 @@ class InferenceState { */ debug(): any { const result: any = {values: {}, variables: {}}; - const objects: Map = new Map(); - function identify(value: InstructionValue): number { + const objects: Map = new Map(); + function identify(value: AliasingEffect): number { let id = objects.get(value); if (id == null) { id = objects.size; @@ -1464,7 +1420,7 @@ class InferenceState { const id = identify(value); result.values[id] = { abstract: this.debugAbstractValue(kind), - value: printInstructionValue(value), + value: printAliasingEffect(value), }; } for (const [variable, values] of this.#variables) { @@ -1481,7 +1437,7 @@ class InferenceState { } inferPhi(phi: Phi): void { - const values: Set = new Set(); + const values: Set = new Set(); for (const [_, operand] of phi.operands) { const operandValues = this.#variables.get(operand.identifier.id); // This is a backedge that will be handled later by State.merge @@ -2347,8 +2303,9 @@ function areArgumentsImmutableAndNonMutating( const values = state.values(place); for (const value of values) { if ( - value.kind === 'FunctionExpression' && - value.loweredFunc.func.params.some(param => { + value.kind === 'CreateFunction' && + value.function.kind === 'FunctionExpression' && + value.function.loweredFunc.func.params.some(param => { const place = param.kind === 'Identifier' ? param : param.place; const range = place.identifier.mutableRange; return range.end > range.start + 1; From d2b2a1b195f096474cc47ae44a0731b19ab3c6c4 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 28 Aug 2025 14:54:27 -0700 Subject: [PATCH 2/3] [compiler] Propagate CreateFunction effects for functions that return functions If you have a local helper function that itself returns a function (`() => () => { ... }`), we currently infer the return effect of the outer function as `Create mutable`. We correctly track the aliasing, but we lose some precision because we don't understand that a function specifically is being returned. Here, we do some extra analysis of which values are returned in InferMutationAliasingRanges, and if the sole return value is a function we infer a `CreateFunction` effect. We also infer an `Assign` (instead of a Create) if the sole return value was one of the context variables or parameters. --- .../Inference/InferMutationAliasingEffects.ts | 43 +++++- .../Inference/InferMutationAliasingRanges.ts | 132 +++++++++++++++--- ...urned-inner-fn-reassigns-context.expect.md | 28 ++-- 3 files changed, 168 insertions(+), 35 deletions(-) diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts index 77cb972ee516..2bb80c0e90cc 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts @@ -2498,10 +2498,47 @@ function computeEffectsForSignature( break; } case 'CreateFunction': { - CompilerError.throwTodo({ - reason: `Support CreateFrom effects in signatures`, - loc: receiver.loc, + const applyInto = substitutions.get(effect.into.identifier.id); + if (applyInto == null || applyInto.length !== 1) { + return null; + } + const captures: Array = []; + for (let i = 0; i < effect.captures.length; i++) { + const substitution = substitutions.get( + effect.captures[i].identifier.id, + ); + if (substitution == null || substitution.length !== 1) { + return null; + } + captures.push(substitution[0]); + } + const context: Array = []; + const originalContext = effect.function.loweredFunc.func.context; + for (let i = 0; i < originalContext.length; i++) { + const substitution = substitutions.get( + originalContext[i].identifier.id, + ); + if (substitution == null || substitution.length !== 1) { + return null; + } + context.push(substitution[0]); + } + effects.push({ + kind: 'CreateFunction', + into: applyInto[0], + function: { + ...effect.function, + loweredFunc: { + ...effect.function.loweredFunc, + func: { + ...effect.function.loweredFunc.func, + context, + }, + }, + }, + captures, }); + break; } default: { assertExhaustive( diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts index b53026a4d4b8..e2e974ddb52f 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts @@ -141,7 +141,7 @@ export function inferMutationAliasingRanges( } else if (effect.kind === 'CreateFunction') { state.create(effect.into, { kind: 'Function', - function: effect.function.loweredFunc.func, + effect, }); } else if (effect.kind === 'CreateFrom') { state.createFrom(index++, effect.from, effect.into); @@ -156,7 +156,7 @@ export function inferMutationAliasingRanges( * invariant here. */ if (!state.nodes.has(effect.into.identifier)) { - state.create(effect.into, {kind: 'Object'}); + state.create(effect.into, {kind: 'Assign'}); } state.assign(index++, effect.from, effect.into); } else if (effect.kind === 'Alias') { @@ -474,6 +474,99 @@ export function inferMutationAliasingRanges( } } + const tracked: Array = []; + for (const param of [...fn.params, ...fn.context, fn.returns]) { + const place = param.kind === 'Identifier' ? param : param.place; + tracked.push(place); + } + + const returned: Set = new Set(); + const queue: Array = [state.nodes.get(fn.returns.identifier)!]; + const seen: Set = new Set(); + while (queue.length !== 0) { + const node = queue.pop()!; + if (seen.has(node)) { + continue; + } + seen.add(node); + for (const id of node.aliases.keys()) { + queue.push(state.nodes.get(id)!); + } + for (const id of node.createdFrom.keys()) { + queue.push(state.nodes.get(id)!); + } + if (node.id.id === fn.returns.identifier.id) { + continue; + } + switch (node.value.kind) { + case 'Assign': + case 'CreateFrom': { + break; + } + case 'Phi': + case 'Object': + case 'Function': { + returned.add(node); + break; + } + default: { + assertExhaustive( + node.value, + `Unexpected node value kind '${(node.value as any).kind}'`, + ); + } + } + } + const returnedValues = [...returned]; + if ( + returnedValues.length === 1 && + returnedValues[0].value.kind === 'Object' && + tracked.some(place => place.identifier.id === returnedValues[0].id.id) + ) { + const from = tracked.find( + place => place.identifier.id === returnedValues[0].id.id, + )!; + functionEffects.push({ + kind: 'Assign', + from, + into: fn.returns, + }); + } else if ( + returnedValues.length === 1 && + returnedValues[0].value.kind === 'Function' + ) { + const outerContext = new Set(fn.context.map(p => p.identifier.id)); + const effect = returnedValues[0].value.effect; + functionEffects.push({ + kind: 'CreateFunction', + function: { + ...effect.function, + loweredFunc: { + func: { + ...effect.function.loweredFunc.func, + context: effect.function.loweredFunc.func.context.filter(p => + outerContext.has(p.identifier.id), + ), + }, + }, + }, + captures: effect.captures.filter(p => outerContext.has(p.identifier.id)), + into: fn.returns, + }); + } else { + const returns = fn.returns.identifier; + functionEffects.push({ + kind: 'Create', + into: fn.returns, + value: isPrimitiveType(returns) + ? ValueKind.Primitive + : isJsxType(returns.type) + ? ValueKind.Frozen + : ValueKind.Mutable, + reason: ValueReason.KnownReturnSignature, + }); + } + /** * Part 3 * Finish populating the externally visible effects. Above we bubble-up the side effects @@ -481,28 +574,12 @@ export function inferMutationAliasingRanges( * Here we populate an effect to create the return value as well as populating alias/capture * effects for how data flows between the params, context vars, and return. */ - const returns = fn.returns.identifier; - functionEffects.push({ - kind: 'Create', - into: fn.returns, - value: isPrimitiveType(returns) - ? ValueKind.Primitive - : isJsxType(returns.type) - ? ValueKind.Frozen - : ValueKind.Mutable, - reason: ValueReason.KnownReturnSignature, - }); /** * Determine precise data-flow effects by simulating transitive mutations of the params/ * captures and seeing what other params/context variables are affected. Anything that * would be transitively mutated needs a capture relationship. */ - const tracked: Array = []; const ignoredErrors = new CompilerError(); - for (const param of [...fn.params, ...fn.context, fn.returns]) { - const place = param.kind === 'Identifier' ? param : param.place; - tracked.push(place); - } for (const into of tracked) { const mutationIndex = index++; state.mutate( @@ -588,9 +665,14 @@ type Node = { lastMutated: number; mutationReason: MutationReason | null; value: + | {kind: 'Assign'} + | {kind: 'CreateFrom'} | {kind: 'Object'} | {kind: 'Phi'} - | {kind: 'Function'; function: HIRFunction}; + | { + kind: 'Function'; + effect: Extract; + }; }; class AliasingState { nodes: Map = new Map(); @@ -612,7 +694,7 @@ class AliasingState { } createFrom(index: number, from: Place, into: Place): void { - this.create(into, {kind: 'Object'}); + this.create(into, {kind: 'CreateFrom'}); const fromNode = this.nodes.get(from.identifier); const toNode = this.nodes.get(into.identifier); if (fromNode == null || toNode == null) { @@ -674,7 +756,10 @@ class AliasingState { continue; } if (node.value.kind === 'Function') { - appendFunctionErrors(errors, node.value.function); + appendFunctionErrors( + errors, + node.value.effect.function.loweredFunc.func, + ); } for (const [alias, when] of node.createdFrom) { if (when >= index) { @@ -738,7 +823,10 @@ class AliasingState { node.transitive == null && node.local == null ) { - appendFunctionErrors(errors, node.value.function); + appendFunctionErrors( + errors, + node.value.effect.function.loweredFunc.func, + ); } if (transitive) { if (node.transitive == null || node.transitive.kind < kind) { diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md index da6b57defe0b..f7742f8f8b66 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md @@ -55,7 +55,7 @@ import { makeArray, Stringify, useIdentity } from "shared-runtime"; */ function Foo(t0) { "use memo"; - const $ = _c(3); + const $ = _c(5); const { b } = t0; const fnFactory = () => () => { @@ -66,18 +66,26 @@ function Foo(t0) { useIdentity(); const fn = fnFactory(); - const arr = makeArray(b); - fn(arr); let t1; - if ($[0] !== arr || $[1] !== myVar) { - t1 = ; - $[0] = arr; - $[1] = myVar; - $[2] = t1; + if ($[0] !== b) { + t1 = makeArray(b); + $[0] = b; + $[1] = t1; + } else { + t1 = $[1]; + } + const arr = t1; + fn(arr); + let t2; + if ($[2] !== arr || $[3] !== myVar) { + t2 = ; + $[2] = arr; + $[3] = myVar; + $[4] = t2; } else { - t1 = $[2]; + t2 = $[4]; } - return t1; + return t2; } function _temp2() { return console.log("b"); From 5008fd1d18e4dcfbb04c10a5c4ebdfba1816f54e Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 28 Aug 2025 14:54:27 -0700 Subject: [PATCH 3/3] [compiler] Repro for missed case of global mutation I realized a pattern of global mutations that we don't currently detect (both in the old and new inference models). If you hide the mutation inside a function returned from another function, we lose track of it: ```js const f = () => () => { global.property = true; }; f()(); ``` Roughly speaking, we need to track that if the return value of `f` is mutated, that it should count as triggering some effects. Right now we encode the idea that a function specifically can have side effects if it is mutated, but other values don't have a way to represent this. I'm thinking that we change the shape of the `Create` effect a bit, and allow room for an optional "mutation effects" array. Then, InferMutationAliasingRanges can visit these effects like it does when trying to find transitive function effects. --- .../error.invalid-assign-global-in-function-factory.js | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-assign-global-in-function-factory.js diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-assign-global-in-function-factory.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-assign-global-in-function-factory.js new file mode 100644 index 000000000000..83b332297f8d --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-assign-global-in-function-factory.js @@ -0,0 +1,7 @@ +function Component() { + const f = () => () => { + global.property = true; + }; + f()(); + return
Ooops
; +}