From 87dc463d3abf5e1b8b2f13c2483fc8abf6653ccf Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Wed, 25 Jun 2025 10:53:38 -0700 Subject: [PATCH] [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 b91b606d50..e875aab68a 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts @@ -2470,10 +2470,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 79f8cf8c0e..0dce5f81d0 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts @@ -140,7 +140,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); @@ -155,7 +155,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') { @@ -465,6 +465,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 @@ -472,28 +565,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( @@ -572,9 +649,14 @@ type Node = { local: {kind: MutationKind; loc: SourceLocation} | null; lastMutated: number; value: + | {kind: 'Assign'} + | {kind: 'CreateFrom'} | {kind: 'Object'} | {kind: 'Phi'} - | {kind: 'Function'; function: HIRFunction}; + | { + kind: 'Function'; + effect: Extract; + }; }; class AliasingState { nodes: Map = new Map(); @@ -594,7 +676,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) { @@ -644,7 +726,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) { @@ -704,7 +789,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 da6b57defe..f7742f8f8b 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 = $[2]; + t1 = $[1]; } - return t1; + const arr = t1; + fn(arr); + let t2; + if ($[2] !== arr || $[3] !== myVar) { + t2 = ; + $[2] = arr; + $[3] = myVar; + $[4] = t2; + } else { + t2 = $[4]; + } + return t2; } function _temp2() { return console.log("b");