From 05a8b6a7171497265205120db943d7e558f86da2 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Tue, 27 May 2025 14:23:38 -0700 Subject: [PATCH] [compiler] Alternate pipeline for new mutability model This PR gets a first fixture working end-to-end with the new mutability and aliasing model. Key changes: * Add a feature flag to enable the model. When enabled we no longer call InferReferenceEffects or InferMutableRanges, and instead use the new equivalents. * Adds a pass that infers Place-specific effects based on mutable ranges and instruction effects. This is necessary to satisfy existing code that requires operand effects to be populated. * Adds a pass that infers the outwardly-visible capturing/aliasing behavior of a function expression. The idea is that this can bubble up and be used in conjunction with the `Apply` effect to get precise inference of things like `array.map(() => { ... })`. [ghstack-poisoned] --- .../src/Entrypoint/Pipeline.ts | 41 +++--- .../src/HIR/BuildHIR.ts | 1 + .../src/HIR/Environment.ts | 5 + .../src/HIR/HIR.ts | 1 + .../src/HIR/PrintHIR.ts | 4 +- .../src/Inference/AnalyseFunctions.ts | 33 ++++- .../Inference/InferMutationAliasingEffects.ts | 48 ++++--- .../InferMutationAliasingFunctionEffects.ts | 136 ++++++++++++++++++ .../Inference/InferMutationAliasingRanges.ts | 116 ++++++++++++++- .../src/Utils/utils.ts | 13 ++ compiler/packages/snap/src/sprout/index.ts | 1 + 11 files changed, 357 insertions(+), 42 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingFunctionEffects.ts diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts index 9053682f04a8f..7524f4fff5090 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -228,25 +228,27 @@ function runWithEnvironment( analyseFunctions(hir); log({kind: 'hir', name: 'AnalyseFunctions', value: hir}); - const mutabilityAliasingErrors = inferMutationAliasingEffects(hir); - log({kind: 'hir', name: 'InferMutationAliasingEffects', value: hir}); - if (env.isInferredMemoEnabled) { - if (mutabilityAliasingErrors.isErr()) { - throw mutabilityAliasingErrors.unwrapErr(); + if (!env.config.enableNewMutationAliasingModel) { + const fnEffectErrors = inferReferenceEffects(hir); + if (env.isInferredMemoEnabled) { + if (fnEffectErrors.length > 0) { + CompilerError.throw(fnEffectErrors[0]); + } } - } - inferMutationAliasingRanges(hir); - log({kind: 'hir', name: 'InferMutationAliasingRanges', value: hir}); - - const fnEffectErrors = inferReferenceEffects(hir); - if (env.isInferredMemoEnabled) { - if (fnEffectErrors.length > 0) { - CompilerError.throw(fnEffectErrors[0]); + log({kind: 'hir', name: 'InferReferenceEffects', value: hir}); + } else { + const mutabilityAliasingErrors = inferMutationAliasingEffects(hir); + log({kind: 'hir', name: 'InferMutationAliasingEffects', value: hir}); + if (env.isInferredMemoEnabled) { + if (mutabilityAliasingErrors.isErr()) { + throw mutabilityAliasingErrors.unwrapErr(); + } } } - log({kind: 'hir', name: 'InferReferenceEffects', value: hir}); - validateLocalsNotReassignedAfterRender(hir); + if (!env.config.enableNewMutationAliasingModel) { + validateLocalsNotReassignedAfterRender(hir); + } // Note: Has to come after infer reference effects because "dead" code may still affect inference deadCodeElimination(hir); @@ -260,8 +262,13 @@ function runWithEnvironment( pruneMaybeThrows(hir); log({kind: 'hir', name: 'PruneMaybeThrows', value: hir}); - inferMutableRanges(hir); - log({kind: 'hir', name: 'InferMutableRanges', value: hir}); + if (!env.config.enableNewMutationAliasingModel) { + inferMutableRanges(hir); + log({kind: 'hir', name: 'InferMutableRanges', value: hir}); + } else { + inferMutationAliasingRanges(hir); + log({kind: 'hir', name: 'InferMutationAliasingRanges', value: hir}); + } if (env.isInferredMemoEnabled) { if (env.config.assertValidMutableRanges) { diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts index 796cba9c5890a..813f1bb98a662 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts @@ -226,6 +226,7 @@ export function lower( loc: func.node.loc ?? GeneratedSource, env, effects: null, + aliasingEffects: null, directives, }); } diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts index 6e6643cd1d68f..8d2e72b22e4e9 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -243,6 +243,11 @@ export const EnvironmentConfigSchema = z.object({ */ enableUseTypeAnnotations: z.boolean().default(false), + /** + * Enable a new model for mutability and aliasing inference + */ + enableNewMutationAliasingModel: z.boolean().default(false), + /** * Enables inference of optional dependency chains. Without this flag * a property chain such as `props?.items?.foo` will infer as a dep on diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts index 4f8f84d7b492e..f1c7a6aebd07b 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts @@ -287,6 +287,7 @@ export type HIRFunction = { generator: boolean; async: boolean; directives: Array; + aliasingEffects?: Array | null; }; export type FunctionEffect = diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts index b597688e8db82..2fc93b1cfd082 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts @@ -75,9 +75,9 @@ export function printFunction(fn: HIRFunction): string { if (definition.length !== 0) { output.push(definition); } - output.push(printType(fn.returnType)); - output.push(printHIR(fn.body)); + output.push(`: ${printType(fn.returnType)} @ ${printPlace(fn.returns)}`); output.push(...fn.directives); + output.push(printHIR(fn.body)); return output.join('\n'); } diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/AnalyseFunctions.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/AnalyseFunctions.ts index 6a4da94d706e2..23cbf27a8d833 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/AnalyseFunctions.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/AnalyseFunctions.ts @@ -26,6 +26,12 @@ import { eachInstructionValueOperand, } from '../HIR/visitors'; import {Iterable_some} from '../Utils/utils'; +import { + AliasingEffect, + inferMutationAliasingEffects, +} from './InferMutationAliasingEffects'; +import {inferMutationAliasingFunctionEffects} from './InferMutationAliasingFunctionEffects'; +import {inferMutationAliasingRanges} from './InferMutationAliasingRanges'; export default function analyseFunctions(func: HIRFunction): void { for (const [_, block] of func.body.blocks) { @@ -33,8 +39,13 @@ export default function analyseFunctions(func: HIRFunction): void { switch (instr.value.kind) { case 'ObjectMethod': case 'FunctionExpression': { - const aliases = lower(instr.value.loweredFunc.func); - infer(instr.value.loweredFunc, aliases); + if (!func.env.config.enableNewMutationAliasingModel) { + const aliases = lower(instr.value.loweredFunc.func); + infer(instr.value.loweredFunc, aliases); + } else { + lowerWithMutationAliasing(instr.value.loweredFunc.func); + infer(instr.value.loweredFunc, new DisjointSet()); + } /** * Reset mutable range for outer inferReferenceEffects @@ -51,6 +62,22 @@ export default function analyseFunctions(func: HIRFunction): void { } } +function lowerWithMutationAliasing(fn: HIRFunction): void { + analyseFunctions(fn); + inferMutationAliasingEffects(fn, {isFunctionExpression: true}); + deadCodeElimination(fn); + inferMutationAliasingRanges(fn); + rewriteInstructionKindsBasedOnReassignment(fn); + inferReactiveScopeVariables(fn); + fn.env.logger?.debugLogIRs?.({ + kind: 'hir', + name: 'AnalyseFunction (inner)', + value: fn, + }); + const effects = inferMutationAliasingFunctionEffects(fn); + fn.aliasingEffects = effects; +} + function lower(func: HIRFunction): DisjointSet { analyseFunctions(func); inferReferenceEffects(func, {isFunctionExpression: true}); @@ -64,7 +91,7 @@ function lower(func: HIRFunction): DisjointSet { value: func, }); inferAliasesForCapturing(func, aliases); - return aliases; + return aliases ?? new DisjointSet(); } /** 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 1c471eac05d4f..2674d68768eed 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts @@ -38,9 +38,7 @@ import { Set_isSuperset, } from '../Utils/utils'; import { - printAliasingEffect, printIdentifier, - printInstruction, printInstructionValue, printPlace, printSourceLocation, @@ -51,7 +49,7 @@ export function inferMutationAliasingEffects( {isFunctionExpression}: {isFunctionExpression: boolean} = { isFunctionExpression: false, }, -): Result, CompilerError> { +): Result { const initialState = InferenceState.empty(fn.env, isFunctionExpression); // Map of blocks to the last (merged) incoming state that was processed @@ -169,7 +167,7 @@ export function inferMutationAliasingEffects( } } } - return Ok([]); + return Ok(undefined); } function inferParam( @@ -203,9 +201,11 @@ function inferBlock( instructionSignature = computeSignatureForInstruction(state.env, instr); instructionSignatureCache.set(instr, instructionSignature); } - // console.log( - // printInstruction({...instr, effects: [...instructionSignature.effects]}), - // ); + /* + * console.log( + * printInstruction({...instr, effects: [...instructionSignature.effects]}), + * ); + */ const effects = applySignature( state, instructionSignature, @@ -304,6 +304,14 @@ function applySignature( isCopyByReferenceValue = false; break; } + case ValueKind.Frozen: { + /* + * TODO: add a separate "ImmutableAlias" effect to downgrade to, that doesn't impact mutable ranges + * We want to remember that the data flow occurred for PruneNonEscapingScopes + */ + isCopyByReferenceValue = false; + break; + } default: { isCopyByReferenceValue = true; break; @@ -321,6 +329,12 @@ function applySignature( */ const fromKind = state.kind(effect.from).kind; switch (fromKind) { + /* + * TODO: add a separate "ImmutableAlias" effect to downgrade to, that doesn't impact mutable ranges + * We want to remember that the data flow occurred for PruneNonEscapingScopes + * (use this to replace the ValueKind.Frozen case) + */ + case ValueKind.Frozen: case ValueKind.Global: case ValueKind.Primitive: { let value = effectInstructionValueCache.get(effect); @@ -1184,8 +1198,8 @@ function computeEffectsForSignature( } // Build substitutions const substitutions: Map> = new Map(); - substitutions.set(signature.receiver.identifier.id, [receiver]); - substitutions.set(signature.returns.identifier.id, [lvalue]); + substitutions.set(signature.receiver, [receiver]); + substitutions.set(signature.returns, [lvalue]); const params = signature.params; for (let i = 0; i < args.length; i++) { const arg = args[i]; @@ -1194,14 +1208,10 @@ function computeEffectsForSignature( return null; } const place = arg.kind === 'Identifier' ? arg : arg.place; - getOrInsertWith( - substitutions, - signature.rest.identifier.id, - () => [], - ).push(place); + getOrInsertWith(substitutions, signature.rest, () => []).push(place); } else { const param = params[i]; - substitutions.set(param.identifier.id, [arg]); + substitutions.set(param, [arg]); } } @@ -1414,10 +1424,10 @@ export type AliasingEffect = }; export type AliasingSignature = { - receiver: Place; - params: Array; - rest: Place | null; - returns: Place; + receiver: IdentifierId; + params: Array; + rest: IdentifierId | null; + returns: IdentifierId; effects: Array; }; diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingFunctionEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingFunctionEffects.ts new file mode 100644 index 0000000000000..c179cd674538b --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingFunctionEffects.ts @@ -0,0 +1,136 @@ +/** + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +import {HIRFunction, IdentifierId, Place} from '../HIR'; +import {getOrInsertDefault} from '../Utils/utils'; +import {AliasingEffect} from './InferMutationAliasingEffects'; + +export function inferMutationAliasingFunctionEffects( + fn: HIRFunction, +): Array | null { + const effects: Array = []; + /* + * Quick hack to infer "mutation" of context vars and params. The problem is that we really want + * to figure out more precisely what changed. Is there a known mutation, or just a conditional + * mutation? + * once we assign scope ids, we can just look through all the effects in the entire body + * and find the maximum level of mutation on each scope (for scopes that are on context refs/params) + */ + const tracked = new Map(); + tracked.set(fn.returns.identifier.id, fn.returns); + for (const operand of fn.context) { + tracked.set(operand.identifier.id, operand); + if (operand.identifier.mutableRange.end > 0) { + effects.push({kind: 'MutateTransitiveConditionally', value: operand}); + } + } + for (const param of fn.params) { + const place = param.kind === 'Identifier' ? param : param.place; + tracked.set(place.identifier.id, place); + if (place.identifier.mutableRange.end > 0) { + effects.push({kind: 'MutateTransitiveConditionally', value: place}); + } + } + type AliasedIdentifier = { + kind: 'Capture' | 'Alias' | 'CreateFrom'; + place: Place; + }; + const dataFlow = new Map>(); + + function lookup( + place: Place, + kind: AliasedIdentifier['kind'], + ): Array | null { + if (tracked.has(place.identifier.id)) { + return [{kind, place}]; + } + return ( + dataFlow.get(place.identifier.id)?.map(aliased => ({ + kind: joinEffects(aliased.kind, kind), + place: aliased.place, + })) ?? null + ); + } + + // todo: fixpoint + for (const block of fn.body.blocks.values()) { + for (const phi of block.phis) { + const operands: Array = []; + for (const operand of phi.operands.values()) { + const inputs = lookup(operand, 'Alias'); + if (inputs != null) { + operands.push(...inputs); + } + } + if (operands.length !== 0) { + dataFlow.set(phi.place.identifier.id, operands); + } + } + for (const instr of block.instructions) { + if (instr.effects == null) continue; + for (const effect of instr.effects) { + if ( + effect.kind === 'Capture' || + effect.kind === 'Alias' || + effect.kind === 'CreateFrom' + ) { + const from = lookup(effect.from, effect.kind); + if (from == null) { + continue; + } + const into = lookup(effect.into, 'Alias'); + if (into == null) { + getOrInsertDefault(dataFlow, effect.into.identifier.id, []).push( + ...from, + ); + } else { + for (const aliased of into) { + getOrInsertDefault( + dataFlow, + aliased.place.identifier.id, + [], + ).push(...from); + } + } + } + } + } + if (block.terminal.kind === 'return') { + const from = lookup(block.terminal.value, 'Alias'); + if (from != null) { + getOrInsertDefault(dataFlow, fn.returns.identifier.id, []).push( + ...from, + ); + } + } + } + + for (const [into, from] of dataFlow) { + const input = tracked.get(into); + if (input == null) { + continue; + } + for (const aliased of from) { + const effect = {kind: aliased.kind, from: aliased.place, into: input}; + effects.push(effect); + } + } + + return effects; +} + +type AliasingKind = 'Alias' | 'Capture' | 'CreateFrom'; +function joinEffects( + effect1: AliasingKind, + effect2: AliasingKind, +): AliasingKind { + if (effect1 === 'Capture' || effect2 === 'Capture') { + return 'Capture'; + } else { + return 'Alias'; + } +} 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 adf9ea3534592..e9f2bce050ae6 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts @@ -5,8 +5,20 @@ * LICENSE file in the root directory of this source tree. */ -import {HIRFunction, Identifier, makeInstructionId} from '../HIR/HIR'; +import { + Effect, + HIRFunction, + Identifier, + IdentifierId, + makeInstructionId, +} from '../HIR/HIR'; +import { + eachInstructionLValue, + eachInstructionValueOperand, + eachTerminalOperand, +} from '../HIR/visitors'; import DisjointSet from '../Utils/DisjointSet'; +import {assertExhaustive} from '../Utils/utils'; import {inferMutableRangesForAlias} from './InferMutableRangesForAlias'; export function inferMutationAliasingRanges(fn: HIRFunction): void { @@ -27,6 +39,13 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void { } for (const instr of block.instructions) { + for (const lvalue of eachInstructionLValue(instr)) { + lvalue.identifier.mutableRange.start = instr.id; + lvalue.identifier.mutableRange.end = makeInstructionId( + Math.max(instr.id + 1, lvalue.identifier.mutableRange.end), + ); + } + if (instr.effects == null) continue; for (const effect of instr.effects) { if ( @@ -72,4 +91,99 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void { } } inferMutableRangesForAlias(fn, aliases); + + /** + * Part 3 + * Add legacy operand-specific effects based on instruction effects and mutable ranges + */ + for (const block of fn.body.blocks.values()) { + for (const phi of block.phis) { + // TODO: we don't actually set these effects today! + phi.place.effect = Effect.Store; + const isPhiMutatedAfterCreation: boolean = + phi.place.identifier.mutableRange.end > + (block.instructions.at(0)?.id ?? block.terminal.id); + for (const operand of phi.operands.values()) { + operand.effect = isPhiMutatedAfterCreation + ? Effect.Capture + : Effect.Read; + } + } + for (const instr of block.instructions) { + if (instr.effects == null) { + for (const lvalue of eachInstructionLValue(instr)) { + lvalue.effect = Effect.ConditionallyMutate; + } + for (const operand of eachInstructionValueOperand(instr.value)) { + operand.effect = Effect.Read; + } + continue; + } + const operandEffects = new Map(); + for (const effect of instr.effects) { + switch (effect.kind) { + case 'Alias': + case 'Capture': + case 'CreateFrom': { + const isMutatedOrReassigned = + effect.into.identifier.mutableRange.end > instr.id; + if (isMutatedOrReassigned) { + operandEffects.set(effect.from.identifier.id, Effect.Capture); + operandEffects.set(effect.into.identifier.id, Effect.Store); + } else { + operandEffects.set(effect.from.identifier.id, Effect.Read); + operandEffects.set(effect.into.identifier.id, Effect.Store); + } + break; + } + case 'Create': { + break; + } + case 'Mutate': { + operandEffects.set(effect.value.identifier.id, Effect.Store); + break; + } + case 'Apply': { + operandEffects.set( + effect.function.place.identifier.id, + Effect.ConditionallyMutate, + ); + break; + } + case 'MutateTransitive': + case 'MutateConditionally': + case 'MutateTransitiveConditionally': { + operandEffects.set( + effect.value.identifier.id, + Effect.ConditionallyMutate, + ); + break; + } + case 'Freeze': { + operandEffects.set(effect.value.identifier.id, Effect.Freeze); + break; + } + default: { + assertExhaustive( + effect, + `Unexpected effect kind ${(effect as any).kind}`, + ); + } + } + } + for (const lvalue of eachInstructionLValue(instr)) { + const effect = + operandEffects.get(lvalue.identifier.id) ?? + Effect.ConditionallyMutate; + lvalue.effect = effect; + } + for (const operand of eachInstructionValueOperand(instr.value)) { + const effect = operandEffects.get(operand.identifier.id) ?? Effect.Read; + operand.effect = effect; + } + } + for (const operand of eachTerminalOperand(block.terminal)) { + operand.effect = Effect.Read; + } + } } diff --git a/compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts b/compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts index e5fbacfc772df..6283be66c10c3 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts @@ -148,6 +148,19 @@ export function Iterable_some( return false; } +export function Iterable_filter( + iter: Iterable, + pred: (item: T) => boolean, +): Array { + const result: Array = []; + for (const item of iter) { + if (pred(item)) { + result.push(item); + } + } + return result; +} + export function nonNull, U>( value: T | null | undefined, ): value is T { diff --git a/compiler/packages/snap/src/sprout/index.ts b/compiler/packages/snap/src/sprout/index.ts index 04748bed28f4f..614f84b3ea238 100644 --- a/compiler/packages/snap/src/sprout/index.ts +++ b/compiler/packages/snap/src/sprout/index.ts @@ -42,6 +42,7 @@ export function runSprout( (globalThis as any).__SNAP_EVALUATOR_MODE = undefined; } if (forgetResult.kind === 'UnexpectedError') { + console.log(forgetCode); return makeError('Unexpected error in Forget runner', forgetResult.value); } if (originalCode.indexOf('@disableNonForgetInSprout') === -1) {