diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/ControlDominators.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/ControlDominators.ts new file mode 100644 index 0000000000000..1fab651947a18 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/ControlDominators.ts @@ -0,0 +1,114 @@ +/** + * Copyright (c) Meta Platforms, Inc. and affiliates. + * + * This source code is licensed under the MIT license found in the + * LICENSE file in the root directory of this source tree. + */ + +import {BlockId, computePostDominatorTree, HIRFunction, Place} from '../HIR'; +import {PostDominator} from '../HIR/Dominator'; + +export type ControlDominators = (id: BlockId) => boolean; + +/** + * Returns an object that lazily calculates whether particular blocks are controlled + * by values of interest. Which values matter are up to the caller. + */ +export function createControlDominators( + fn: HIRFunction, + isControlVariable: (place: Place) => boolean, +): ControlDominators { + const postDominators = computePostDominatorTree(fn, { + includeThrowsAsExitNode: false, + }); + const postDominatorFrontierCache = new Map>(); + + function isControlledBlock(id: BlockId): boolean { + let controlBlocks = postDominatorFrontierCache.get(id); + if (controlBlocks === undefined) { + controlBlocks = postDominatorFrontier(fn, postDominators, id); + postDominatorFrontierCache.set(id, controlBlocks); + } + for (const blockId of controlBlocks) { + const controlBlock = fn.body.blocks.get(blockId)!; + switch (controlBlock.terminal.kind) { + case 'if': + case 'branch': { + if (isControlVariable(controlBlock.terminal.test)) { + return true; + } + break; + } + case 'switch': { + if (isControlVariable(controlBlock.terminal.test)) { + return true; + } + for (const case_ of controlBlock.terminal.cases) { + if (case_.test !== null && isControlVariable(case_.test)) { + return true; + } + } + break; + } + } + } + return false; + } + + return isControlledBlock; +} + +/* + * Computes the post-dominator frontier of @param block. These are immediate successors of nodes that + * post-dominate @param targetId and from which execution may not reach @param block. Intuitively, these + * are the earliest blocks from which execution branches such that it may or may not reach the target block. + */ +function postDominatorFrontier( + fn: HIRFunction, + postDominators: PostDominator, + targetId: BlockId, +): Set { + const visited = new Set(); + const frontier = new Set(); + const targetPostDominators = postDominatorsOf(fn, postDominators, targetId); + for (const blockId of [...targetPostDominators, targetId]) { + if (visited.has(blockId)) { + continue; + } + visited.add(blockId); + const block = fn.body.blocks.get(blockId)!; + for (const pred of block.preds) { + if (!targetPostDominators.has(pred)) { + // The predecessor does not always reach this block, we found an item on the frontier! + frontier.add(pred); + } + } + } + return frontier; +} + +function postDominatorsOf( + fn: HIRFunction, + postDominators: PostDominator, + targetId: BlockId, +): Set { + const result = new Set(); + const visited = new Set(); + const queue = [targetId]; + while (queue.length) { + const currentId = queue.shift()!; + if (visited.has(currentId)) { + continue; + } + visited.add(currentId); + const current = fn.body.blocks.get(currentId)!; + for (const pred of current.preds) { + const predPostDominator = postDominators.get(pred) ?? pred; + if (predPostDominator === targetId || result.has(predPostDominator)) { + result.add(pred); + } + queue.push(pred); + } + } + return result; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReactivePlaces.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReactivePlaces.ts index 271a76e92c125..d7ebcfe2fbc81 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReactivePlaces.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReactivePlaces.ts @@ -7,7 +7,6 @@ import {CompilerError} from '..'; import { - BlockId, Effect, Environment, HIRFunction, @@ -15,14 +14,12 @@ import { IdentifierId, Instruction, Place, - computePostDominatorTree, evaluatesToStableTypeOrContainer, getHookKind, isStableType, isStableTypeContainer, isUseOperator, } from '../HIR'; -import {PostDominator} from '../HIR/Dominator'; import { eachInstructionLValue, eachInstructionOperand, @@ -35,6 +32,7 @@ import { } from '../ReactiveScopes/InferReactiveScopeVariables'; import DisjointSet from '../Utils/DisjointSet'; import {assertExhaustive} from '../Utils/utils'; +import {createControlDominators} from './ControlDominators'; /** * Side map to track and propagate sources of stability (i.e. hook calls such as @@ -212,45 +210,9 @@ export function inferReactivePlaces(fn: HIRFunction): void { reactiveIdentifiers.markReactive(place); } - const postDominators = computePostDominatorTree(fn, { - includeThrowsAsExitNode: false, - }); - const postDominatorFrontierCache = new Map>(); - - function isReactiveControlledBlock(id: BlockId): boolean { - let controlBlocks = postDominatorFrontierCache.get(id); - if (controlBlocks === undefined) { - controlBlocks = postDominatorFrontier(fn, postDominators, id); - postDominatorFrontierCache.set(id, controlBlocks); - } - for (const blockId of controlBlocks) { - const controlBlock = fn.body.blocks.get(blockId)!; - switch (controlBlock.terminal.kind) { - case 'if': - case 'branch': { - if (reactiveIdentifiers.isReactive(controlBlock.terminal.test)) { - return true; - } - break; - } - case 'switch': { - if (reactiveIdentifiers.isReactive(controlBlock.terminal.test)) { - return true; - } - for (const case_ of controlBlock.terminal.cases) { - if ( - case_.test !== null && - reactiveIdentifiers.isReactive(case_.test) - ) { - return true; - } - } - break; - } - } - } - return false; - } + const isReactiveControlledBlock = createControlDominators(fn, place => + reactiveIdentifiers.isReactive(place), + ); do { for (const [, block] of fn.body.blocks) { @@ -411,61 +373,6 @@ export function inferReactivePlaces(fn: HIRFunction): void { propagateReactivityToInnerFunctions(fn, true); } -/* - * Computes the post-dominator frontier of @param block. These are immediate successors of nodes that - * post-dominate @param targetId and from which execution may not reach @param block. Intuitively, these - * are the earliest blocks from which execution branches such that it may or may not reach the target block. - */ -function postDominatorFrontier( - fn: HIRFunction, - postDominators: PostDominator, - targetId: BlockId, -): Set { - const visited = new Set(); - const frontier = new Set(); - const targetPostDominators = postDominatorsOf(fn, postDominators, targetId); - for (const blockId of [...targetPostDominators, targetId]) { - if (visited.has(blockId)) { - continue; - } - visited.add(blockId); - const block = fn.body.blocks.get(blockId)!; - for (const pred of block.preds) { - if (!targetPostDominators.has(pred)) { - // The predecessor does not always reach this block, we found an item on the frontier! - frontier.add(pred); - } - } - } - return frontier; -} - -function postDominatorsOf( - fn: HIRFunction, - postDominators: PostDominator, - targetId: BlockId, -): Set { - const result = new Set(); - const visited = new Set(); - const queue = [targetId]; - while (queue.length) { - const currentId = queue.shift()!; - if (visited.has(currentId)) { - continue; - } - visited.add(currentId); - const current = fn.body.blocks.get(currentId)!; - for (const pred of current.preds) { - const predPostDominator = postDominators.get(pred) ?? pred; - if (predPostDominator === targetId || result.has(predPostDominator)) { - result.add(pred); - } - queue.push(pred); - } - } - return result; -} - class ReactivityMap { hasChanges: boolean = false; reactive: Set = new Set();