diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts index 456425aecaae9..d3c919a6d8afe 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts @@ -8,7 +8,6 @@ import { Set_union, getOrInsertDefault, } from '../Utils/utils'; -import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies'; import { BasicBlock, BlockId, @@ -22,7 +21,6 @@ import { ReactiveScopeDependency, ScopeId, } from './HIR'; -import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR'; const DEBUG_PRINT = false; @@ -373,17 +371,10 @@ function collectNonNullsInBlocks( !fn.env.config.enableTreatFunctionDepsAsConditional ) { const innerFn = instr.value.loweredFunc; - const innerTemporaries = collectTemporariesSidemap( - innerFn.func, - new Set(), - ); - const innerOptionals = collectOptionalChainSidemap(innerFn.func); const innerHoistableMap = collectHoistablePropertyLoadsImpl( innerFn.func, { ...context, - temporaries: innerTemporaries, // TODO: remove in later PR - hoistableFromOptionals: innerOptionals.hoistableObjects, // TODO: remove in later PR nestedFnImmutableContext: context.nestedFnImmutableContext ?? new Set( diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectOptionalChainDependencies.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectOptionalChainDependencies.ts index 453294784246f..0167c996b177d 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectOptionalChainDependencies.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectOptionalChainDependencies.ts @@ -1,4 +1,5 @@ import {CompilerError} from '..'; +import {getOrInsertDefault} from '../Utils/utils'; import {assertNonNull} from './CollectHoistablePropertyLoads'; import { BlockId, @@ -22,25 +23,14 @@ export function collectOptionalChainSidemap( fn: HIRFunction, ): OptionalChainSidemap { const context: OptionalTraversalContext = { + currFn: fn, blocks: fn.body.blocks, seenOptionals: new Set(), - processedInstrsInOptional: new Set(), + processedInstrsInOptional: new Map(), temporariesReadInOptional: new Map(), hoistableObjects: new Map(), }; - for (const [_, block] of fn.body.blocks) { - if ( - block.terminal.kind === 'optional' && - !context.seenOptionals.has(block.id) - ) { - traverseOptionalBlock( - block as TBasicBlock, - context, - null, - ); - } - } - + traverseFunction(fn, context); return { temporariesReadInOptional: context.temporariesReadInOptional, processedInstrsInOptional: context.processedInstrsInOptional, @@ -96,8 +86,13 @@ export type OptionalChainSidemap = { * bb5: * $5 = MethodCall $2.$4() <--- here, we want to take a dep on $2 and $4! * ``` + * + * Also note that InstructionIds are not unique across inner functions. */ - processedInstrsInOptional: ReadonlySet; + processedInstrsInOptional: ReadonlyMap< + HIRFunction, + ReadonlySet + >; /** * Records optional chains for which we can safely evaluate non-optional * PropertyLoads. e.g. given `a?.b.c`, we can evaluate any load from `a?.b` at @@ -115,16 +110,46 @@ export type OptionalChainSidemap = { }; type OptionalTraversalContext = { + currFn: HIRFunction; blocks: ReadonlyMap; // Track optional blocks to avoid outer calls into nested optionals seenOptionals: Set; - processedInstrsInOptional: Set; + processedInstrsInOptional: Map>; temporariesReadInOptional: Map; hoistableObjects: Map; }; +function traverseFunction( + fn: HIRFunction, + context: OptionalTraversalContext, +): void { + for (const [_, block] of fn.body.blocks) { + for (const instr of block.instructions) { + if ( + instr.value.kind === 'FunctionExpression' || + instr.value.kind === 'ObjectMethod' + ) { + traverseFunction(instr.value.loweredFunc.func, { + ...context, + currFn: instr.value.loweredFunc.func, + blocks: instr.value.loweredFunc.func.body.blocks, + }); + } + } + if ( + block.terminal.kind === 'optional' && + !context.seenOptionals.has(block.id) + ) { + traverseOptionalBlock( + block as TBasicBlock, + context, + null, + ); + } + } +} /** * Match the consequent and alternate blocks of an optional. * @returns propertyload computed by the consequent block, or null if the @@ -369,10 +394,13 @@ function traverseOptionalBlock( }, ], }; - context.processedInstrsInOptional.add( - matchConsequentResult.storeLocalInstrId, + const processedInstrsInOptionalByFn = getOrInsertDefault( + context.processedInstrsInOptional, + context.currFn, + new Set(), ); - context.processedInstrsInOptional.add(test.id); + processedInstrsInOptionalByFn.add(matchConsequentResult.storeLocalInstrId); + processedInstrsInOptionalByFn.add(test.id); context.temporariesReadInOptional.set( matchConsequentResult.consequentId, load, diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts index 0178aea6e4c56..8f4abdf6dac0d 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts @@ -176,8 +176,10 @@ function findTemporariesUsedOutsideDeclaringScope( * $2 = LoadLocal 'foo' * $3 = CallExpression $2($1) * ``` - * Only map LoadLocal and PropertyLoad lvalues to their source if we know that - * reordering the read (from the time-of-load to time-of-use) is valid. + * @param usedOutsideDeclaringScope is used to check the correctness of + * reordering LoadLocal / PropertyLoad calls. We only track a LoadLocal / + * PropertyLoad in the returned temporaries map if reordering the read (from the + * time-of-load to time-of-use) is valid. * * If a LoadLocal or PropertyLoad instruction is within the reactive scope range * (a proxy for mutable range) of the load source, later instructions may @@ -215,7 +217,29 @@ export function collectTemporariesSidemap( fn: HIRFunction, usedOutsideDeclaringScope: ReadonlySet, ): ReadonlyMap { - const temporaries = new Map(); + const temporaries = new Map(); + collectTemporariesSidemapImpl( + fn, + usedOutsideDeclaringScope, + temporaries, + false, + ); + return temporaries; +} + +/** + * Recursive collect a sidemap of all `LoadLocal` and `PropertyLoads` with a + * function and all nested functions. + * + * Note that IdentifierIds are currently unique, so we can use a single + * Map across all nested functions. + */ +function collectTemporariesSidemapImpl( + fn: HIRFunction, + usedOutsideDeclaringScope: ReadonlySet, + temporaries: Map, + isInnerFn: boolean, +): void { for (const [_, block] of fn.body.blocks) { for (const instr of block.instructions) { const {value, lvalue} = instr; @@ -224,27 +248,51 @@ export function collectTemporariesSidemap( ); if (value.kind === 'PropertyLoad' && !usedOutside) { - const property = getProperty( - value.object, - value.property, - false, - temporaries, - ); - temporaries.set(lvalue.identifier.id, property); + if (!isInnerFn || temporaries.has(value.object.identifier.id)) { + /** + * All dependencies of a inner / nested function must have a base + * identifier from the outermost component / hook. This is because the + * compiler cannot break an inner function into multiple granular + * scopes. + */ + const property = getProperty( + value.object, + value.property, + false, + temporaries, + ); + temporaries.set(lvalue.identifier.id, property); + } } else if ( value.kind === 'LoadLocal' && lvalue.identifier.name == null && value.place.identifier.name !== null && !usedOutside ) { - temporaries.set(lvalue.identifier.id, { - identifier: value.place.identifier, - path: [], - }); + if ( + !isInnerFn || + fn.context.some( + context => context.identifier.id === value.place.identifier.id, + ) + ) { + temporaries.set(lvalue.identifier.id, { + identifier: value.place.identifier, + path: [], + }); + } + } else if ( + value.kind === 'FunctionExpression' || + value.kind === 'ObjectMethod' + ) { + collectTemporariesSidemapImpl( + value.loweredFunc.func, + usedOutsideDeclaringScope, + temporaries, + true, + ); } } } - return temporaries; } function getProperty( @@ -310,6 +358,12 @@ class Context { #temporaries: ReadonlyMap; #temporariesUsedOutsideScope: ReadonlySet; + /** + * Tracks the traversal state. See Context.declare for explanation of why this + * is needed. + */ + inInnerFn: boolean = false; + constructor( temporariesUsedOutsideScope: ReadonlySet, temporaries: ReadonlyMap, @@ -360,12 +414,23 @@ class Context { } /* - * Records where a value was declared, and optionally, the scope where the value originated from. - * This is later used to determine if a dependency should be added to a scope; if the current - * scope we are visiting is the same scope where the value originates, it can't be a dependency - * on itself. + * Records where a value was declared, and optionally, the scope where the + * value originated from. This is later used to determine if a dependency + * should be added to a scope; if the current scope we are visiting is the + * same scope where the value originates, it can't be a dependency on itself. + * + * Note that we do not track declarations or reassignments within inner + * functions for the following reasons: + * - inner functions cannot be split by scope boundaries and are guaranteed + * to consume their own declarations + * - reassignments within inner functions are tracked as context variables, + * which already have extended mutable ranges to account for reassignments + * - *most importantly* it's currently simply incorrect to compare inner + * function instruction ids (tracked by `decl`) with outer ones (as stored + * by root identifier mutable ranges). */ declare(identifier: Identifier, decl: Decl): void { + if (this.inInnerFn) return; if (!this.#declarations.has(identifier.declarationId)) { this.#declarations.set(identifier.declarationId, decl); } @@ -575,7 +640,10 @@ function collectDependencies( fn: HIRFunction, usedOutsideDeclaringScope: ReadonlySet, temporaries: ReadonlyMap, - processedInstrsInOptional: ReadonlySet, + processedInstrsInOptional: ReadonlyMap< + HIRFunction, + ReadonlySet + >, ): Map> { const context = new Context(usedOutsideDeclaringScope, temporaries); @@ -595,6 +663,12 @@ function collectDependencies( const scopeTraversal = new ScopeBlockTraversal(); + const shouldSkipInstructionDependencies = ( + fn: HIRFunction, + id: InstructionId, + ): boolean => { + return processedInstrsInOptional.get(fn)?.has(id) ?? false; + }; for (const [blockId, block] of fn.body.blocks) { scopeTraversal.recordScopes(block); const scopeBlockInfo = scopeTraversal.blockInfos.get(blockId); @@ -614,12 +688,12 @@ function collectDependencies( } } for (const instr of block.instructions) { - if (!processedInstrsInOptional.has(instr.id)) { + if (!shouldSkipInstructionDependencies(fn, instr.id)) { handleInstruction(instr, context); } } - if (!processedInstrsInOptional.has(block.terminal.id)) { + if (!shouldSkipInstructionDependencies(fn, block.terminal.id)) { for (const place of eachTerminalOperand(block.terminal)) { context.visitOperand(place); } diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts index 217bc3132bd14..c9ee803bfaffd 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts @@ -1215,9 +1215,17 @@ export class ScopeBlockTraversal { } } + /** + * @returns if the given scope is currently 'active', i.e. if the scope start + * block but not the scope fallthrough has been recorded. + */ isScopeActive(scopeId: ScopeId): boolean { return this.#activeScopes.indexOf(scopeId) !== -1; } + + /** + * The current, innermost active scope. + */ get currentScope(): ScopeId | null { return this.#activeScopes.at(-1) ?? null; }