Skip to content

Commit

Permalink
[compiler][hir] Only hoist always-accessed PropertyLoads from functio…
Browse files Browse the repository at this point in the history 
…n decls

ghstack-source-id: e18f46a30e92e85d6e3103c504a13cc980eb90ad
Pull Request resolved: #31066
  • Loading branch information
@mofeiZ
mofeiZ committed Sep 30, 2024
1 parent 97608a6 commit 1fe2f2b
Show file tree
Hide file tree
Showing 33 changed files with 1,592 additions and 75 deletions.
231 changes: 161 additions & 70 deletions compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import {
Set_union,
getOrInsertDefault,
} from '../Utils/utils';
import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies';
import {
BasicBlock,
BlockId,
Expand All @@ -15,10 +16,12 @@ import {
HIRFunction,
Identifier,
IdentifierId,
InstructionId,
InstructionValue,
ReactiveScopeDependency,
ScopeId,
} from './HIR';
import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR';

/**
* Helper function for `PropagateScopeDependencies`. Uses control flow graph
Expand Down Expand Up @@ -83,28 +86,57 @@ export function collectHoistablePropertyLoads(
fn: HIRFunction,
temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
): ReadonlyMap<ScopeId, BlockInfo> {
nestedFnImmutableContext: ReadonlySet<IdentifierId> | null,
): ReadonlyMap<BlockId, BlockInfo> {
const registry = new PropertyPathRegistry();

const nodes = collectNonNullsInBlocks(
fn,
temporaries,
const functionExpressionLoads = collectFunctionExpressionFakeLoads(fn);
const actuallyEvaluatedTemporaries = new Map(
[...temporaries].filter(([id]) => !functionExpressionLoads.has(id)),
);

/**
* Due to current limitations of mutable range inference, there are edge cases in
* which we infer known-immutable values (e.g. props or hook params) to have a
* mutable range and scope.
* (see `destructure-array-declaration-to-context-var` fixture)
* We track known immutable identifiers to reduce regressions (as PropagateScopeDeps
* is being rewritten to HIR).
*/
const knownImmutableIdentifiers = new Set<IdentifierId>();
if (fn.fnType === 'Component' || fn.fnType === 'Hook') {
for (const p of fn.params) {
if (p.kind === 'Identifier') {
knownImmutableIdentifiers.add(p.identifier.id);
}
}
}
const nodes = collectNonNullsInBlocks(fn, {
temporaries: actuallyEvaluatedTemporaries,
knownImmutableIdentifiers,
hoistableFromOptionals,
registry,
);
nestedFnImmutableContext,
});
propagateNonNull(fn, nodes, registry);

const nodesKeyedByScopeId = new Map<ScopeId, BlockInfo>();
return nodes;
}

export function keyByScopeId<T>(
fn: HIRFunction,
source: ReadonlyMap<BlockId, T>,
): ReadonlyMap<ScopeId, T> {
const keyedByScopeId = new Map<ScopeId, T>();
for (const [_, block] of fn.body.blocks) {
if (block.terminal.kind === 'scope') {
nodesKeyedByScopeId.set(
keyedByScopeId.set(
block.terminal.scope.id,
nodes.get(block.terminal.block)!,
source.get(block.terminal.block)!,
);
}
}

return nodesKeyedByScopeId;
return keyedByScopeId;
}

export type BlockInfo = {
Expand Down Expand Up @@ -211,45 +243,75 @@ class PropertyPathRegistry {

function getMaybeNonNullInInstruction(
instr: InstructionValue,
temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
registry: PropertyPathRegistry,
context: CollectNonNullsInBlocksContext,
): PropertyPathNode | null {
let path = null;
if (instr.kind === 'PropertyLoad') {
path = temporaries.get(instr.object.identifier.id) ?? {
path = context.temporaries.get(instr.object.identifier.id) ?? {
identifier: instr.object.identifier,
path: [],
};
} else if (instr.kind === 'Destructure') {
path = temporaries.get(instr.value.identifier.id) ?? null;
path = context.temporaries.get(instr.value.identifier.id) ?? null;
} else if (instr.kind === 'ComputedLoad') {
path = temporaries.get(instr.object.identifier.id) ?? null;
path = context.temporaries.get(instr.object.identifier.id) ?? null;
}
return path != null ? context.registry.getOrCreateProperty(path) : null;
}

function isImmutableAtInstr(
identifier: Identifier,
instr: InstructionId,
context: CollectNonNullsInBlocksContext,
): boolean {
if (context.nestedFnImmutableContext != null) {
/**
* Comparing instructions ids across inner-outer function bodies is not valid, as they are numbered
*/
return context.nestedFnImmutableContext.has(identifier.id);
} else {
/**
* Since this runs *after* buildReactiveScopeTerminals, identifier mutable ranges
* are not valid with respect to current instruction id numbering.
* We use attached reactive scope ranges as a proxy for mutable range, but this
* is an overestimate as (1) scope ranges merge and align to form valid program
* blocks and (2) passes like MemoizeFbtAndMacroOperands may assign scopes to
* non-mutable identifiers.
*
* See comment in exported function for why we track known immutable identifiers.
*/
const mutableAtInstr =
identifier.mutableRange.end > identifier.mutableRange.start + 1 &&
identifier.scope != null &&
inRange(
{
id: instr,
},
identifier.scope.range,
);
return (
!mutableAtInstr || context.knownImmutableIdentifiers.has(identifier.id)
);
}
return path != null ? registry.getOrCreateProperty(path) : null;
}

type CollectNonNullsInBlocksContext = {
temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
knownImmutableIdentifiers: ReadonlySet<IdentifierId>;
hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>;
registry: PropertyPathRegistry;
/**
* (For nested / inner function declarations)
* Context variables (i.e. captured from an outer scope) that are immutable.
* Note that this technically could be merged into `knownImmutableIdentifiers`,
* but are currently kept separate for readability.
*/
nestedFnImmutableContext: ReadonlySet<IdentifierId> | null;
};
function collectNonNullsInBlocks(
fn: HIRFunction,
temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
registry: PropertyPathRegistry,
context: CollectNonNullsInBlocksContext,
): ReadonlyMap<BlockId, BlockInfo> {
/**
* Due to current limitations of mutable range inference, there are edge cases in
* which we infer known-immutable values (e.g. props or hook params) to have a
* mutable range and scope.
* (see `destructure-array-declaration-to-context-var` fixture)
* We track known immutable identifiers to reduce regressions (as PropagateScopeDeps
* is being rewritten to HIR).
*/
const knownImmutableIdentifiers = new Set<IdentifierId>();
if (fn.fnType === 'Component' || fn.fnType === 'Hook') {
for (const p of fn.params) {
if (p.kind === 'Identifier') {
knownImmutableIdentifiers.add(p.identifier.id);
}
}
}
/**
* Known non-null objects such as functional component props can be safely
* read from any block.
Expand All @@ -261,7 +323,9 @@ function collectNonNullsInBlocks(
fn.params[0].kind === 'Identifier'
) {
const identifier = fn.params[0].identifier;
knownNonNullIdentifiers.add(registry.getOrCreateIdentifier(identifier));
knownNonNullIdentifiers.add(
context.registry.getOrCreateIdentifier(identifier),
);
}
const nodes = new Map<BlockId, BlockInfo>();
for (const [_, block] of fn.body.blocks) {
Expand All @@ -273,46 +337,49 @@ function collectNonNullsInBlocks(
block,
assumedNonNullObjects,
});
const maybeOptionalChain = hoistableFromOptionals.get(block.id);
const maybeOptionalChain = context.hoistableFromOptionals.get(block.id);
if (maybeOptionalChain != null) {
assumedNonNullObjects.add(
registry.getOrCreateProperty(maybeOptionalChain),
context.registry.getOrCreateProperty(maybeOptionalChain),
);
continue;
}
for (const instr of block.instructions) {
const maybeNonNull = getMaybeNonNullInInstruction(
instr.value,
temporaries,
registry,
);
if (maybeNonNull != null) {
const baseIdentifier = maybeNonNull.fullPath.identifier;
/**
* Since this runs *after* buildReactiveScopeTerminals, identifier mutable ranges
* are not valid with respect to current instruction id numbering.
* We use attached reactive scope ranges as a proxy for mutable range, but this
* is an overestimate as (1) scope ranges merge and align to form valid program
* blocks and (2) passes like MemoizeFbtAndMacroOperands may assign scopes to
* non-mutable identifiers.
*
* See comment at top of function for why we track known immutable identifiers.
*/
const isMutableAtInstr =
baseIdentifier.mutableRange.end >
baseIdentifier.mutableRange.start + 1 &&
baseIdentifier.scope != null &&
inRange(
{
id: instr.id,
},
baseIdentifier.scope.range,
);
if (
!isMutableAtInstr ||
knownImmutableIdentifiers.has(baseIdentifier.id)
) {
assumedNonNullObjects.add(maybeNonNull);
const maybeNonNull = getMaybeNonNullInInstruction(instr.value, context);
if (
maybeNonNull != null &&
isImmutableAtInstr(maybeNonNull.fullPath.identifier, instr.id, context)
) {
assumedNonNullObjects.add(maybeNonNull);
}
if (
instr.value.kind === 'FunctionExpression' &&
!fn.env.config.enableTreatFunctionDepsAsConditional
) {
const innerFn = instr.value.loweredFunc;
const innerTemporaries = collectTemporariesSidemap(
innerFn.func,
new Set(),
);
const innerOptionals = collectOptionalChainSidemap(innerFn.func);
const innerHoistableMap = collectHoistablePropertyLoads(
innerFn.func,
innerTemporaries,
innerOptionals.hoistableObjects,
context.nestedFnImmutableContext ??
new Set(
innerFn.func.context
.filter(place =>
isImmutableAtInstr(place.identifier, instr.id, context),
)
.map(place => place.identifier.id),
),
);
const innerHoistables = assertNonNull(
innerHoistableMap.get(innerFn.func.body.entry),
);
for (const entry of innerHoistables.assumedNonNullObjects) {
assumedNonNullObjects.add(entry);
}
}
}
Expand Down Expand Up @@ -515,3 +582,27 @@ function reduceMaybeOptionalChains(
}
} while (changed);
}

function collectFunctionExpressionFakeLoads(
fn: HIRFunction,
): Set<IdentifierId> {
const sources = new Map<IdentifierId, IdentifierId>();
const functionExpressionReferences = new Set<IdentifierId>();

for (const [_, block] of fn.body.blocks) {
for (const {lvalue, value} of block.instructions) {
if (value.kind === 'FunctionExpression') {
for (const reference of value.loweredFunc.dependencies) {
let curr: IdentifierId | undefined = reference.identifier.id;
while (curr != null) {
functionExpressionReferences.add(curr);
curr = sources.get(curr);
}
}
} else if (value.kind === 'PropertyLoad') {
sources.set(lvalue.identifier.id, value.object.identifier.id);
}
}
}
return functionExpressionReferences;
}
12 changes: 7 additions & 5 deletions compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,10 @@ import {
areEqualPaths,
IdentifierId,
} from './HIR';
import {collectHoistablePropertyLoads} from './CollectHoistablePropertyLoads';
import {
collectHoistablePropertyLoads,
keyByScopeId,
} from './CollectHoistablePropertyLoads';
import {
ScopeBlockTraversal,
eachInstructionOperand,
Expand All @@ -41,10 +44,9 @@ export function propagateScopeDependenciesHIR(fn: HIRFunction): void {
hoistableObjects,
} = collectOptionalChainSidemap(fn);

const hoistablePropertyLoads = collectHoistablePropertyLoads(
const hoistablePropertyLoads = keyByScopeId(
fn,
temporaries,
hoistableObjects,
collectHoistablePropertyLoads(fn, temporaries, hoistableObjects, null),
);

const scopeDeps = collectDependencies(
Expand Down Expand Up @@ -209,7 +211,7 @@ function findTemporariesUsedOutsideDeclaringScope(
* of $1, as the evaluation of `arr.length` changes between instructions $1 and
* $3. We do not track $1 -> arr.length in this case.
*/
function collectTemporariesSidemap(
export function collectTemporariesSidemap(
fn: HIRFunction,
usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
): ReadonlyMap<IdentifierId, ReactiveScopeDependency> {
Expand Down
Loading

0 comments on commit 1fe2f2b

Please sign in to comment.