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 3e6f82999a9e..651ea451633f 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 000000000000..4c04ae197292 --- /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 000000000000..923d0b59bb81 --- /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}, + ], +};