From 996b3be844d4f644d5b2cf177fdf6b121adbd811 Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Sun, 1 Jun 2025 17:12:48 -0700 Subject: [PATCH] [compiler] InferMutationAliasingRanges precisely models which values mutate when It turns out that InferMutationAliasingRanges does need a fixpoint loop, but the approach is arguably simpler overall and more precise than the previous implementation. Like InferMutationAliasingEffects (which is the new InferReferenceEffects), we build an abstract model of the heap. But here we know what the effects are, so we can do abstract interpretation of the effects. Each abstract value stores a set of values that it has captured (for transitive mutation), while each variable keeps a set of values it may directly mutate (for assign/alias/capturefrom). This means that at each mutation, we can mark _exactly_ the set of variables/values that are affected by that specific instruction. This means we can correctly infer that `mutate(b)` can't impact `a` here: ``` a = make(); b = make(); mutate(b); // when we interpret the mutation here, a isn't captured yet b.a = a; ``` We will need to make this a fixpoint, but only if there are backedges in the CFG. ghstack-source-id: 6fd16624699774535c21c57228e46b53ce5e5a29 Pull Request resolved: https://github.com/facebook/react/pull/33401 --- .../Inference/InferMutationAliasingRanges.ts | 269 +++++++++++++----- ...ssa-renaming-ternary-destruction.expect.md | 70 +++++ .../ssa-renaming-ternary-destruction.js | 21 ++ 3 files changed, 281 insertions(+), 79 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/ssa-renaming-ternary-destruction.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/ssa-renaming-ternary-destruction.js 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 3e6f82999a..651ea45163 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingRanges.ts @@ -5,46 +5,61 @@ * LICENSE file in the root directory of this source tree. */ -import {CompilerError} from '..'; +import prettyFormat from 'pretty-format'; +import {CompilerError, SourceLocation} from '..'; import { Effect, HIRFunction, Identifier, IdentifierId, + InstructionId, makeInstructionId, + Place, } from '../HIR/HIR'; import { eachInstructionLValue, eachInstructionValueOperand, eachTerminalOperand, } from '../HIR/visitors'; -import DisjointSet from '../Utils/DisjointSet'; -import {assertExhaustive} from '../Utils/utils'; -import {debugAliases} from './InferMutableRanges'; -import {inferMutableRangesForAlias} from './InferMutableRangesForAlias'; +import {assertExhaustive, getOrInsertWith} from '../Utils/utils'; +import {printFunction} from '../HIR'; +import {printInstruction} from '../HIR/PrintHIR'; + +const DEBUG = false; /** * Infers mutable ranges for all values. */ export function inferMutationAliasingRanges(fn: HIRFunction): void { + if (DEBUG) { + console.log(); + } /** - * Part 1 - * Infer ranges for transitive mutations, which includes mutations that affect - * captured references and not just direct aliases. We build a distjoing set - * that tracks capturing and direct aliasing, and look at transitive mutations - * only. + * Part 1: Infer mutable ranges for values. We build an abstract model + * of the effect types, distinguishing values which can capture references + * to other values and variables, which can point to one or more values. + * + * Transitive mutation marks value identifiers as mutated and also walks + * into the identifiers captured by that abstract value: local mutation + * only marks the top-level identifiers as mutated. + * + * TODO: add a fixpoint. */ - const captures = new DisjointSet(); + const state = new AliasingState(); + for (const param of fn.context) { + state.create(param); + } for (const block of fn.body.blocks.values()) { for (const phi of block.phis) { - captures.union([ - phi.place.identifier, - ...[...phi.operands.values()].map(place => place.identifier), - ]); + state.create(phi.place); + for (const operand of phi.operands.values()) { + state.alias(operand, phi.place); + } } for (const instr of block.instructions) { for (const lvalue of eachInstructionLValue(instr)) { + state.create(lvalue); if (lvalue.identifier.mutableRange.start === 0) { lvalue.identifier.mutableRange.start = instr.id; } @@ -55,84 +70,46 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void { if (instr.effects == null) continue; for (const effect of instr.effects) { - if ( - effect.kind === 'Assign' || - effect.kind === 'Alias' || - effect.kind === 'CreateFrom' || - effect.kind === 'Capture' - ) { - captures.union([effect.from.identifier, effect.into.identifier]); + if (effect.kind === 'Create' || effect.kind === 'CreateFunction') { + state.create(effect.into); + } else if (effect.kind === 'Assign') { + state.assign(effect.from, effect.into); + } else if (effect.kind === 'Alias' || effect.kind === 'CreateFrom') { + state.alias(effect.from, effect.into); + } else if (effect.kind === 'Capture') { + state.capture(effect.from, effect.into); } else if ( effect.kind === 'MutateTransitive' || effect.kind === 'MutateTransitiveConditionally' ) { - const value = effect.value; - value.identifier.mutableRange.end = makeInstructionId(instr.id + 1); - } - } - } - } - /** - * TODO: this will incorrectly mark values as mutated in the following case: - * 1. Create value x - * 2. Create value y - * 3. Transitively mutate y - * 4. Capture x -> y - * - * We will put these all into one alias set in captures, and then InferMutableRangesForAlias - * will look at all identifiers in the set that start before the mutation. Thus it will see - * that x is created before the mutation, and consider it mutated. But the capture doesn't - * occur until after the mutation! - * - * The fix is to use a graph to precisely describe what is captured where at each instruction, - * so that on a transitive mutation we can iterate all the captured values and mark them. - * - * This then needs a fixpoint: although we would track captures for phi operands, the operands' - * own capture values won't be populated until we do a fixpoint. - */ - inferMutableRangesForAlias(fn, captures); - - /** - * Part 2 - * Infer ranges for local (non-transitive) mutations. We build a disjoint set - * that only tracks direct value aliasing, and look only at local mutations - * to extend ranges - * - * TODO: similar design to the above todo for captures, use a directed graph instead of disjoint union, - * with fixpoint. - */ - const aliases = new DisjointSet(); - for (const block of fn.body.blocks.values()) { - for (const phi of block.phis) { - aliases.union([ - phi.place.identifier, - ...[...phi.operands.values()].map(place => place.identifier), - ]); - } - - for (const instr of block.instructions) { - if (instr.effects == null) continue; - for (const effect of instr.effects) { - if ( - effect.kind === 'Assign' || - effect.kind === 'Alias' || - effect.kind === 'CreateFrom' - ) { - aliases.union([effect.from.identifier, effect.into.identifier]); + state.mutateTransitive( + effect.value.identifier, + makeInstructionId(instr.id + 1), + effect.value.loc, + ); } else if ( effect.kind === 'Mutate' || effect.kind === 'MutateConditionally' ) { - const value = effect.value; - value.identifier.mutableRange.end = makeInstructionId(instr.id + 1); + state.mutate( + effect.value.identifier, + makeInstructionId(instr.id + 1), + effect.value.loc, + ); } } + if (DEBUG) { + console.log(printInstruction(instr)); + console.log(state.debug()); + } } } - inferMutableRangesForAlias(fn, aliases); + if (DEBUG) { + console.log(state.debug()); + } /** - * Part 3 + * Part 2 * Add legacy operand-specific effects based on instruction effects and mutable ranges. * Also fixes up operand mutable ranges, making sure that start is non-zero if the value * is mutated (depended on by later passes like InferReactiveScopeVariables which uses this @@ -239,4 +216,138 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void { operand.effect = Effect.Read; } } + + if (DEBUG) { + console.log(printFunction(fn)); + } +} + +/** + * TODO: similar to alias effect inference state: + * - values represent the conceptual values (context vars, things that got Create*-ed) + * into which other values are captured + * - variables deals with aliases (Assign, Alias, CreateFrom) + * + * Ex: + * `Capture a -> b`: + * - find the values represented by `a` (aValues) by looking up a in .variables, then mapping to .values + * - find the values represented by `b` (bValues) by looking up b in .variables, then mapping to .values + * - add aValues into bValues + */ +class AliasingState { + values: Map> = new Map(); + variables: Map> = new Map(); + + create(place: Place): void { + this.values.set(place.identifier, new Set()); + this.variables.set(place.identifier.id, new Set([place.identifier])); + } + + clear(lvalue: Place): void { + this.variables.set(lvalue.identifier.id, new Set()); + } + + assign(from: Place, into: Place): void { + const fromVariables = this.variables.get(from.identifier.id); + if (fromVariables == null) { + return; + } + this.variables.set( + into.identifier.id, + new Set([...fromVariables, from.identifier, into.identifier]), + // fromVariables, + ); + } + + alias(from: Place, into: Place): void { + const intoVariables = getOrInsertWith( + this.variables, + into.identifier.id, + () => new Set(), + ); + intoVariables.add(from.identifier); + } + + capture(from: Place, into: Place): void { + const intoVariables = this.variables.get(into.identifier.id)!; + for (const v of intoVariables) { + const values = this.values.get(v)!; + values.add(from.identifier); + } + } + + mutateTransitive( + place: Identifier, + end: InstructionId, + loc: SourceLocation, + ): void { + const variables = this.variables.get(place.id); + if (variables == null) { + return; + } + for (const value of variables) { + value.mutableRange.end = makeInstructionId( + Math.max(value.mutableRange.end, end), + ); + const captured = this.values.get(value)!; + for (const capture of captured) { + this.mutateTransitive(capture, end, loc); + } + } + } + + mutate(place: Identifier, end: InstructionId, _loc: SourceLocation): void { + const variables = this.variables.get(place.id); + if (variables == null) { + return; + } + place.mutableRange.end = makeInstructionId( + Math.max(place.mutableRange.end, end), + ); + for (const value of variables) { + // Unlike mutateTransitive, we don't traverse into captured values + value.mutableRange.end = makeInstructionId( + Math.max(value.mutableRange.end, end), + ); + } + } + + canonicalize(): Map { + const map = new Map(); + for (const value of this.values.keys()) { + map.set( + value.id, + `${value.mutableRange.start}:${value.mutableRange.end}`, + ); + } + return map; + } + + debug(): string { + const values = new Map(); + for (const [k, v] of this.values) { + values.set( + `${k.name?.value ?? ''}$${k.id}:${k.mutableRange.start}:${k.mutableRange.end}`, + [...v] + .map( + v => + `${v.name?.value ?? ''}$${v.id}:${v.mutableRange.start}:${v.mutableRange.end}`, + ) + .join(', '), + ); + } + const variables = new Map(); + for (const [k, v] of this.variables) { + variables.set( + `$${k}`, + [...v] + .map( + v => + `${v.name?.value ?? ''}$${v.id}:${v.mutableRange.start}:${v.mutableRange.end}`, + ) + .join(', '), + ); + } + return prettyFormat({values, variables}); + } } diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/ssa-renaming-ternary-destruction.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/ssa-renaming-ternary-destruction.expect.md new file mode 100644 index 0000000000..4c04ae1972 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/ssa-renaming-ternary-destruction.expect.md @@ -0,0 +1,70 @@ + +## Input + +```javascript +// @enablePropagateDepsInHIR @enableNewMutationAliasingModel +function useFoo(props) { + let x = []; + x.push(props.bar); + // todo: the below should memoize separately from the above + // my guess is that the phi causes the different `x` identifiers + // to get added to an alias group. this is where we need to track + // the actual state of the alias groups at the time of the mutation + props.cond ? (({x} = {x: {}}), ([x] = [[]]), x.push(props.foo)) : null; + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [{cond: false, foo: 2, bar: 55}], + sequentialRenders: [ + {cond: false, foo: 2, bar: 55}, + {cond: false, foo: 3, bar: 55}, + {cond: true, foo: 3, bar: 55}, + ], +}; + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR @enableNewMutationAliasingModel +function useFoo(props) { + const $ = _c(5); + let x; + if ($[0] !== props.bar) { + x = []; + x.push(props.bar); + $[0] = props.bar; + $[1] = x; + } else { + x = $[1]; + } + if ($[2] !== props.cond || $[3] !== props.foo) { + props.cond ? (([x] = [[]]), x.push(props.foo)) : null; + $[2] = props.cond; + $[3] = props.foo; + $[4] = x; + } else { + x = $[4]; + } + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [{ cond: false, foo: 2, bar: 55 }], + sequentialRenders: [ + { cond: false, foo: 2, bar: 55 }, + { cond: false, foo: 3, bar: 55 }, + { cond: true, foo: 3, bar: 55 }, + ], +}; + +``` + +### Eval output +(kind: ok) [55] +[55] +[3] \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/ssa-renaming-ternary-destruction.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/ssa-renaming-ternary-destruction.js new file mode 100644 index 0000000000..923d0b59bb --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/ssa-renaming-ternary-destruction.js @@ -0,0 +1,21 @@ +// @enablePropagateDepsInHIR @enableNewMutationAliasingModel +function useFoo(props) { + let x = []; + x.push(props.bar); + // todo: the below should memoize separately from the above + // my guess is that the phi causes the different `x` identifiers + // to get added to an alias group. this is where we need to track + // the actual state of the alias groups at the time of the mutation + props.cond ? (({x} = {x: {}}), ([x] = [[]]), x.push(props.foo)) : null; + return x; +} + +export const FIXTURE_ENTRYPOINT = { + fn: useFoo, + params: [{cond: false, foo: 2, bar: 55}], + sequentialRenders: [ + {cond: false, foo: 2, bar: 55}, + {cond: false, foo: 3, bar: 55}, + {cond: true, foo: 3, bar: 55}, + ], +};