From 84347ecff2ad2d81714829e65ffbfd4bc73e097f Mon Sep 17 00:00:00 2001 From: Mofei Zhang Date: Wed, 21 May 2025 14:11:54 -0400 Subject: [PATCH] [compiler] Inferred effect dependencies now include optional chains Inferred effect dependencies now include optional chains. This is a temporary solution while https://github.com/facebook/react/pull/32099 and its followups are worked on. Ideally, we should model reactive scope dependencies in the IR similarly to `ComputeIR` -- dependencies should be hoisted and all references rewritten to use the hoisted dependencies. (will add more test cases) --- .../src/HIR/ScopeDependencyUtils.ts | 280 ++++++++++++++++++ .../src/Inference/InferEffectDependencies.ts | 278 +++++++++++------ .../reactive-optional-chain.expect.md | 31 +- .../reactive-optional-chain.js | 8 +- 4 files changed, 496 insertions(+), 101 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/HIR/ScopeDependencyUtils.ts diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/ScopeDependencyUtils.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/ScopeDependencyUtils.ts new file mode 100644 index 0000000000..b986cd211b --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/ScopeDependencyUtils.ts @@ -0,0 +1,280 @@ +import { + Place, + ReactiveScopeDependency, + Identifier, + makeInstructionId, + InstructionKind, + GeneratedSource, + BlockId, + makeTemporaryIdentifier, + Effect, + GotoVariant, + HIR, +} from './HIR'; +import {CompilerError} from '../CompilerError'; +import {Environment} from './Environment'; +import HIRBuilder from './HIRBuilder'; +import {lowerValueToTemporary} from './BuildHIR'; +import {printDependency} from '../ReactiveScopes/PrintReactiveFunction'; + +type DependencyInstructions = { + place: Place; + value: HIR; + exitBlockId: BlockId; +}; + +export function buildDependencyInstructions( + dep: ReactiveScopeDependency, + env: Environment, +): DependencyInstructions { + const builder = new HIRBuilder(env, { + entryBlockKind: 'value', + }); + let result: Place; + if (dep.path.every(path => !path.optional)) { + const last = writeNonOptionalDependency(dep, env, builder); + result = { + kind: 'Identifier', + identifier: last, + effect: Effect.Freeze, + reactive: dep.reactive, + loc: GeneratedSource, + }; + } else { + const last = writeOptionalDependency( + dep.path.length - 1, + dep, + builder, + null, + ); + result = { + kind: 'Identifier', + identifier: last, + effect: Effect.Freeze, + reactive: dep.reactive, + loc: GeneratedSource, + }; + } + + const exitBlockId = builder.terminate( + { + kind: 'unsupported', + loc: GeneratedSource, + id: makeInstructionId(0), + }, + null, + ); + return { + place: result, + value: builder.build(), + exitBlockId, + }; +} +function writeNonOptionalDependency( + dep: ReactiveScopeDependency, + env: Environment, + builder: HIRBuilder, +): Identifier { + const loc = dep.identifier.loc; + let last: Identifier = makeTemporaryIdentifier(env.nextIdentifierId, loc); + builder.push({ + lvalue: { + identifier: last, + kind: 'Identifier', + effect: Effect.Mutate, + reactive: dep.reactive, + loc, + }, + value: { + kind: 'LoadLocal', + place: { + identifier: dep.identifier, + kind: 'Identifier', + effect: Effect.Freeze, + reactive: dep.reactive, + loc, + }, + loc, + }, + id: makeInstructionId(1), + loc: loc, + }); + + for (const path of dep.path) { + const next = makeTemporaryIdentifier(env.nextIdentifierId, loc); + builder.push({ + lvalue: { + identifier: next, + kind: 'Identifier', + effect: Effect.Mutate, + reactive: dep.reactive, + loc, + }, + value: { + kind: 'PropertyLoad', + object: { + identifier: last, + kind: 'Identifier', + effect: Effect.Freeze, + reactive: dep.reactive, + loc, + }, + property: path.property, + loc, + }, + id: makeInstructionId(1), + loc: loc, + }); + last = next; + } + return last; +} + +function writeOptionalDependency( + idx: number, + dep: ReactiveScopeDependency, + builder: HIRBuilder, + parentAlternate: BlockId | null, +): Identifier { + const env = builder.environment; + CompilerError.invariant( + idx >= 0 && !dep.path.slice(0, idx + 1).every(path => !path.optional), + { + reason: '[WriteOptional] Expected optional path', + description: `${idx} ${printDependency(dep)}`, + loc: GeneratedSource, + }, + ); + const continuationBlock = builder.reserve(builder.currentBlockKind()); + const consequent = builder.reserve('value'); + + const returnPlace: Place = { + kind: 'Identifier', + identifier: makeTemporaryIdentifier(env.nextIdentifierId, GeneratedSource), + effect: Effect.Mutate, + reactive: dep.reactive, + loc: GeneratedSource, + }; + + let alternate; + if (parentAlternate != null) { + alternate = parentAlternate; + } else { + /** + * Make outermost alternate block + * $N = Primitive undefined + * $M = StoreLocal $OptionalResult = $N + * goto fallthrough + */ + alternate = builder.enter('value', () => { + const temp = lowerValueToTemporary(builder, { + kind: 'Primitive', + value: undefined, + loc: GeneratedSource, + }); + lowerValueToTemporary(builder, { + kind: 'StoreLocal', + lvalue: {kind: InstructionKind.Const, place: {...returnPlace}}, + value: {...temp}, + type: null, + loc: GeneratedSource, + }); + return { + kind: 'goto', + variant: GotoVariant.Break, + block: continuationBlock.id, + id: makeInstructionId(0), + loc: GeneratedSource, + }; + }); + } + + let testIdentifier: Identifier | null = null; + const testBlock = builder.enter('value', () => { + const firstOptional = dep.path.findIndex(path => path.optional); + if (idx === firstOptional) { + // Lower test block + testIdentifier = writeNonOptionalDependency( + { + identifier: dep.identifier, + reactive: dep.reactive, + path: dep.path.slice(0, idx), + }, + env, + builder, + ); + } else { + testIdentifier = writeOptionalDependency( + idx - 1, + dep, + builder, + alternate, + ); + } + + return { + kind: 'branch', + test: { + identifier: testIdentifier, + effect: Effect.Freeze, + kind: 'Identifier', + loc: GeneratedSource, + reactive: dep.reactive, + }, + consequent: consequent.id, + alternate, + id: makeInstructionId(0), + loc: GeneratedSource, + fallthrough: continuationBlock.id, + }; + }); + + builder.enterReserved(consequent, () => { + CompilerError.invariant(testIdentifier !== null, { + reason: 'Satisfy type checker', + description: null, + loc: null, + suggestions: null, + }); + + const tmpConsequent = lowerValueToTemporary(builder, { + kind: 'PropertyLoad', + object: { + identifier: testIdentifier, + kind: 'Identifier', + effect: Effect.Freeze, + reactive: dep.reactive, + loc: GeneratedSource, + }, + property: dep.path[idx].property, + loc: GeneratedSource, + }); + lowerValueToTemporary(builder, { + kind: 'StoreLocal', + lvalue: {kind: InstructionKind.Const, place: {...returnPlace}}, + value: {...tmpConsequent}, + type: null, + loc: GeneratedSource, + }); + return { + kind: 'goto', + variant: GotoVariant.Break, + block: continuationBlock.id, + id: makeInstructionId(0), + loc: GeneratedSource, + }; + }); + builder.terminateWithContinuation( + { + kind: 'optional', + optional: dep.path[idx].optional, + test: testBlock, + fallthrough: continuationBlock.id, + id: makeInstructionId(0), + loc: GeneratedSource, + }, + continuationBlock, + ); + + return returnPlace.identifier; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferEffectDependencies.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferEffectDependencies.ts index f1a5843419..722f9f7587 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferEffectDependencies.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferEffectDependencies.ts @@ -10,7 +10,6 @@ import {CompilerError, SourceLocation} from '..'; import { ArrayExpression, Effect, - Environment, FunctionExpression, GeneratedSource, HIRFunction, @@ -29,6 +28,9 @@ import { isSetStateType, isFireFunctionType, makeScopeId, + HIR, + BasicBlock, + BlockId, } from '../HIR'; import {collectHoistablePropertyLoadsInInnerFn} from '../HIR/CollectHoistablePropertyLoads'; import {collectOptionalChainSidemap} from '../HIR/CollectOptionalChainDependencies'; @@ -44,7 +46,12 @@ import { DependencyCollectionContext, handleInstruction, } from '../HIR/PropagateScopeDependenciesHIR'; -import {eachInstructionOperand, eachTerminalOperand} from '../HIR/visitors'; +import {buildDependencyInstructions} from '../HIR/ScopeDependencyUtils'; +import { + eachInstructionOperand, + eachTerminalOperand, + terminalFallthrough, +} from '../HIR/visitors'; import {empty} from '../Utils/Stack'; import {getOrInsertWith} from '../Utils/utils'; @@ -53,7 +60,6 @@ import {getOrInsertWith} from '../Utils/utils'; * a second argument to the useEffect call if no dependency array is provided. */ export function inferEffectDependencies(fn: HIRFunction): void { - let hasRewrite = false; const fnExpressions = new Map< IdentifierId, TInstruction @@ -86,6 +92,7 @@ export function inferEffectDependencies(fn: HIRFunction): void { * reactive(Identifier i) = Union_{reference of i}(reactive(reference)) */ const reactiveIds = inferReactiveIdentifiers(fn); + const rewriteBlocks: Array = []; for (const [, block] of fn.body.blocks) { if (block.terminal.kind === 'scope') { @@ -101,7 +108,7 @@ export function inferEffectDependencies(fn: HIRFunction): void { ); } } - const rewriteInstrs = new Map>(); + const rewriteInstrs: Array = []; for (const instr of block.instructions) { const {value, lvalue} = instr; if (value.kind === 'FunctionExpression') { @@ -165,7 +172,6 @@ export function inferEffectDependencies(fn: HIRFunction): void { ) { // We have a useEffect call with no deps array, so we need to infer the deps const effectDeps: Array = []; - const newInstructions: Array = []; const deps: ArrayExpression = { kind: 'ArrayExpression', elements: effectDeps, @@ -196,24 +202,28 @@ export function inferEffectDependencies(fn: HIRFunction): void { */ const usedDeps = []; - for (const dep of minimalDeps) { + for (const maybeDep of minimalDeps) { if ( - ((isUseRefType(dep.identifier) || - isSetStateType(dep.identifier)) && - !reactiveIds.has(dep.identifier.id)) || - isFireFunctionType(dep.identifier) + ((isUseRefType(maybeDep.identifier) || + isSetStateType(maybeDep.identifier)) && + !reactiveIds.has(maybeDep.identifier.id)) || + isFireFunctionType(maybeDep.identifier) ) { // exclude non-reactive hook results, which will never be in a memo block continue; } - const {place, instructions} = writeDependencyToInstructions( + const dep = truncateDepAtCurrent(maybeDep); + const {place, value, exitBlockId} = buildDependencyInstructions( dep, - reactiveIds.has(dep.identifier.id), fn.env, - fnExpr.loc, ); - newInstructions.push(...instructions); + rewriteInstrs.push({ + kind: 'block', + location: instr.id, + value, + exitBlockId: exitBlockId, + }); effectDeps.push(place); usedDeps.push(dep); } @@ -234,27 +244,32 @@ export function inferEffectDependencies(fn: HIRFunction): void { }); } - newInstructions.push({ - id: makeInstructionId(0), - loc: GeneratedSource, - lvalue: {...depsPlace, effect: Effect.Mutate}, - value: deps, - }); - // Step 2: push the inferred deps array as an argument of the useEffect + rewriteInstrs.push({ + kind: 'instr', + location: instr.id, + value: { + id: makeInstructionId(0), + loc: GeneratedSource, + lvalue: {...depsPlace, effect: Effect.Mutate}, + value: deps, + }, + }); value.args.push({...depsPlace, effect: Effect.Freeze}); - rewriteInstrs.set(instr.id, newInstructions); fn.env.inferredEffectLocations.add(callee.loc); } else if (loadGlobals.has(value.args[0].identifier.id)) { // Global functions have no reactive dependencies, so we can insert an empty array - newInstructions.push({ - id: makeInstructionId(0), - loc: GeneratedSource, - lvalue: {...depsPlace, effect: Effect.Mutate}, - value: deps, + rewriteInstrs.push({ + kind: 'instr', + location: instr.id, + value: { + id: makeInstructionId(0), + loc: GeneratedSource, + lvalue: {...depsPlace, effect: Effect.Mutate}, + value: deps, + }, }); value.args.push({...depsPlace, effect: Effect.Freeze}); - rewriteInstrs.set(instr.id, newInstructions); fn.env.inferredEffectLocations.add(callee.loc); } } else if ( @@ -285,21 +300,13 @@ export function inferEffectDependencies(fn: HIRFunction): void { } } } - if (rewriteInstrs.size > 0) { - hasRewrite = true; - const newInstrs = []; - for (const instr of block.instructions) { - const newInstr = rewriteInstrs.get(instr.id); - if (newInstr != null) { - newInstrs.push(...newInstr, instr); - } else { - newInstrs.push(instr); - } - } - block.instructions = newInstrs; - } + rewriteSplices(block, rewriteInstrs, rewriteBlocks); } - if (hasRewrite) { + + if (rewriteBlocks.length > 0) { + for (const block of rewriteBlocks) { + fn.body.blocks.set(block.id, block); + } // Renumber instructions and fix scope ranges markInstructionIds(fn.body); fixScopeAndIdentifierRanges(fn.body); @@ -307,63 +314,144 @@ export function inferEffectDependencies(fn: HIRFunction): void { } } -function writeDependencyToInstructions( +function truncateDepAtCurrent( dep: ReactiveScopeDependency, - reactive: boolean, - env: Environment, - loc: SourceLocation, -): {place: Place; instructions: Array} { - const instructions: Array = []; - let currValue = createTemporaryPlace(env, GeneratedSource); - currValue.reactive = reactive; - instructions.push({ - id: makeInstructionId(0), - loc: GeneratedSource, - lvalue: {...currValue, effect: Effect.Mutate}, - value: { - kind: 'LoadLocal', - place: { - kind: 'Identifier', - identifier: dep.identifier, - effect: Effect.Capture, - reactive, - loc: loc, - }, - loc: loc, - }, - }); - for (const path of dep.path) { - if (path.optional) { - /** - * TODO: instead of truncating optional paths, reuse - * instructions from hoisted dependencies block(s) - */ - break; +): ReactiveScopeDependency { + const idx = dep.path.findIndex(path => path.property === 'current'); + if (idx === -1) { + return dep; + } else { + return {...dep, path: dep.path.slice(0, idx)}; + } +} + +type SpliceInfo = + | {kind: 'instr'; location: InstructionId; value: Instruction} + | { + kind: 'block'; + location: InstructionId; + value: HIR; + exitBlockId: BlockId; + }; + +function rewriteSplices( + originalBlock: BasicBlock, + splices: Array, + rewriteBlocks: Array, +): void { + if (splices.length === 0) { + return; + } + /** + * Splice instructions or value blocks into the original block. + * --- original block --- + * bb_original + * instr1 + * ... + * instr2 <-- splice location + * instr3 + * ... + * + * + * If there is more than one block in the splice, this means that we're + * splicing in a set of value-blocks of the following structure: + * --- blocks we're splicing in --- + * bb_entry: + * instrEntry + * ... + * fallthrough=bb_exit + * + * bb1(value): + * ... + * + * bb_exit: + * instrExit + * ... + * + * + * + * --- rewritten blocks --- + * bb_original + * instr1 + * ... (original instructions) + * instr2 + * instrEntry + * ... (spliced instructions) + * fallthrough=bb_exit + * + * bb1(value): + * ... + * + * bb_exit: + * instrExit + * ... (spliced instructions) + * instr3 + * ... (original instructions) + * + */ + const originalInstrs = originalBlock.instructions; + let currBlock: BasicBlock = {...originalBlock, instructions: []}; + rewriteBlocks.push(currBlock); + + let cursor = 0; + for (const rewrite of splices) { + while (originalInstrs[cursor].id < rewrite.location) { + CompilerError.invariant( + originalInstrs[cursor].id < originalInstrs[cursor + 1].id, + { + reason: + '[InferEffectDependencies] Internal invariant broken: expected block instructions to be sorted', + loc: originalInstrs[cursor].loc, + }, + ); + currBlock.instructions.push(originalInstrs[cursor]); + cursor++; } - if (path.property === 'current') { - /* - * Prune ref.current accesses. This may over-capture for non-ref values with - * a current property, but that's fine. - */ - break; + CompilerError.invariant(originalInstrs[cursor].id === rewrite.location, { + reason: + '[InferEffectDependencies] Internal invariant broken: splice location not found', + loc: originalInstrs[cursor].loc, + }); + + if (rewrite.kind === 'instr') { + currBlock.instructions.push(rewrite.value); + } else { + const {entry, blocks} = rewrite.value; + const entryBlock = blocks.get(entry)!; + // splice in all instructions from the entry block + currBlock.instructions.push(...entryBlock.instructions); + if (blocks.size > 1) { + /** + * We're splicing in a set of value-blocks, which means we need + * to push new blocks and update terminals. + */ + CompilerError.invariant( + terminalFallthrough(entryBlock.terminal) === rewrite.exitBlockId, + { + reason: + '[InferEffectDependencies] Internal invariant broken: expected entry block to have a fallthrough', + loc: entryBlock.terminal.loc, + }, + ); + const originalTerminal = currBlock.terminal; + currBlock.terminal = entryBlock.terminal; + + for (const [id, block] of blocks) { + if (id === entry) { + continue; + } + if (id === rewrite.exitBlockId) { + block.terminal = originalTerminal; + currBlock = block; + } + rewriteBlocks.push(block); + } + } } - const nextValue = createTemporaryPlace(env, GeneratedSource); - nextValue.reactive = reactive; - instructions.push({ - id: makeInstructionId(0), - loc: GeneratedSource, - lvalue: {...nextValue, effect: Effect.Mutate}, - value: { - kind: 'PropertyLoad', - object: {...currValue, effect: Effect.Capture}, - property: path.property, - loc: loc, - }, - }); - currValue = nextValue; + } + for (let i = cursor; i < originalInstrs.length; i++) { + currBlock.instructions.push(originalInstrs[i]); } - currValue.effect = Effect.Freeze; - return {place: currValue, instructions}; } function inferReactiveIdentifiers(fn: HIRFunction): Set { diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.expect.md index 7c9f21b85c..d2414f3f9b 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.expect.md +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.expect.md @@ -8,10 +8,16 @@ import {print} from 'shared-runtime'; // TODO: take optional chains as dependencies function ReactiveMemberExpr({cond, propVal}) { - const obj = {a: cond ? {b: propVal} : null}; + const obj = {a: cond ? {b: propVal} : null, c: null}; useEffect(() => print(obj.a?.b)); + useEffect(() => print(obj.c?.d)); } +export const FIXTURE_ENTRYPOINT = { + fn: ReactiveMemberExpr, + params: [{cond: true, propVal: 1}], +}; + ``` ## Code @@ -23,7 +29,7 @@ import { print } from "shared-runtime"; // TODO: take optional chains as dependencies function ReactiveMemberExpr(t0) { - const $ = _c(7); + const $ = _c(9); const { cond, propVal } = t0; let t1; if ($[0] !== cond || $[1] !== propVal) { @@ -36,7 +42,7 @@ function ReactiveMemberExpr(t0) { } let t2; if ($[3] !== t1) { - t2 = { a: t1 }; + t2 = { a: t1, c: null }; $[3] = t1; $[4] = t2; } else { @@ -51,10 +57,25 @@ function ReactiveMemberExpr(t0) { } else { t3 = $[6]; } - useEffect(t3, [obj.a]); + useEffect(t3, [obj.a?.b]); + let t4; + if ($[7] !== obj.c?.d) { + t4 = () => print(obj.c?.d); + $[7] = obj.c?.d; + $[8] = t4; + } else { + t4 = $[8]; + } + useEffect(t4, [obj.c?.d]); } +export const FIXTURE_ENTRYPOINT = { + fn: ReactiveMemberExpr, + params: [{ cond: true, propVal: 1 }], +}; + ``` ### Eval output -(kind: exception) Fixture not implemented \ No newline at end of file +(kind: ok) +logs: [1,undefined] \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.js index 8a76784e24..4f567fa152 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.js +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/infer-effect-dependencies/reactive-optional-chain.js @@ -4,6 +4,12 @@ import {print} from 'shared-runtime'; // TODO: take optional chains as dependencies function ReactiveMemberExpr({cond, propVal}) { - const obj = {a: cond ? {b: propVal} : null}; + const obj = {a: cond ? {b: propVal} : null, c: null}; useEffect(() => print(obj.a?.b)); + useEffect(() => print(obj.c?.d)); } + +export const FIXTURE_ENTRYPOINT = { + fn: ReactiveMemberExpr, + params: [{cond: true, propVal: 1}], +};