diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts index ca6abc0748eb9..f5052cb1bcebc 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -99,6 +99,8 @@ import {propagateScopeDependenciesHIR} from '../HIR/PropagateScopeDependenciesHI import {outlineJSX} from '../Optimization/OutlineJsx'; import {optimizePropsMethodCalls} from '../Optimization/OptimizePropsMethodCalls'; import {transformFire} from '../Transform'; +import {buildReactiveGraph} from '../ReactiveIR/BuildReactiveGraph'; +import {printReactiveGraph} from '../ReactiveIR/ReactiveIR'; export type CompilerPipelineValue = | {kind: 'ast'; name: string; value: CodegenFunction} @@ -373,6 +375,15 @@ function runWithEnvironment( }); } + if (env.config.enableReactiveGraph) { + const reactiveGraph = buildReactiveGraph(hir); + log({ + kind: 'debug', + name: 'BuildReactiveGraph', + value: printReactiveGraph(reactiveGraph), + }); + } + const reactiveFunction = buildReactiveFunction(hir); log({ kind: 'reactive', diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts index f3f426df56e44..f6543e9337371 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -395,6 +395,12 @@ const EnvironmentConfigSchema = z.object({ */ enableInstructionReordering: z.boolean().default(false), + /** + * Enables ReactiveGraph-based optimizations including reordering across terminal + * boundaries + */ + enableReactiveGraph: z.boolean().default(false), + /** * Enables function outlinining, where anonymous functions that do not close over * local variables can be extracted into top-level helper functions. diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts new file mode 100644 index 0000000000000..c9e5e98b0e1ef --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/BuildReactiveGraph.ts @@ -0,0 +1,228 @@ +/** + * 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 {CompilerError, SourceLocation} from '..'; +import {BlockId, HIRFunction, Identifier, IdentifierId, Place} from '../HIR'; +import {printIdentifier, printPlace} from '../HIR/PrintHIR'; +import { + eachInstructionLValue, + eachInstructionValueOperand, + terminalFallthrough, +} from '../HIR/visitors'; +import { + IfNode, + InstructionNode, + LoadArgumentNode, + makeReactiveId, + NodeDependencies, + NodeReference, + populateReactiveGraphNodeOutputs, + printReactiveNodes, + ReactiveGraph, + ReactiveId, + ReactiveNode, + ReturnNode, + reversePostorderReactiveGraph, + ScopeNode, +} from './ReactiveIR'; + +export function buildReactiveGraph(fn: HIRFunction): ReactiveGraph { + const builder = new Builder(); + for (const param of fn.params) { + const place = param.kind === 'Identifier' ? param : param.place; + const node: LoadArgumentNode = { + kind: 'LoadArgument', + id: builder.nextReactiveId, + loc: place.loc, + outputs: [], + place: {...place}, + }; + builder.nodes.set(node.id, node); + builder.declare(node.id, place); + } + + const exitNode = buildBlockScope(fn, builder, fn.body.entry); + + const graph: ReactiveGraph = { + async: fn.async, + directives: fn.directives, + env: fn.env, + exit: exitNode, + fnType: fn.fnType, + generator: fn.generator, + id: fn.id, + loc: fn.loc, + nextNodeId: builder._nextNodeId, + nodes: builder.nodes, + params: fn.params, + }; + populateReactiveGraphNodeOutputs(graph); + reversePostorderReactiveGraph(graph); + return graph; +} + +class Builder { + _nextNodeId: number = 0; + #environment: Map = new Map(); + nodes: Map = new Map(); + args: Set = new Set(); + + get nextReactiveId(): ReactiveId { + return makeReactiveId(this._nextNodeId++); + } + + declare(node: ReactiveId, place: Place): void { + this.#environment.set(place.identifier.id, {node, from: place}); + } + + lookup( + identifier: Identifier, + loc: SourceLocation, + ): {node: ReactiveId; from: Place} { + const dep = this.#environment.get(identifier.id); + if (dep == null) { + console.log(printReactiveNodes(this.nodes)); + for (const [id, dep] of this.#environment) { + console.log(`t#${id} => £${dep.node} . ${printPlace(dep.from)}`); + } + + console.log(); + console.log(`could not find ${printIdentifier(identifier)}`); + } + CompilerError.invariant(dep != null, { + reason: `No source node for identifier ${printIdentifier(identifier)}`, + loc, + }); + return dep; + } +} + +function buildBlockScope( + fn: HIRFunction, + builder: Builder, + entry: BlockId, +): ReactiveId { + let block = fn.body.blocks.get(entry)!; + let lastNode: ReactiveNode = { + kind: 'Empty', + id: builder.nextReactiveId, + loc: block.terminal.loc, + outputs: [], + }; + builder.nodes.set(lastNode.id, lastNode); + while (true) { + // iterate instructions of the block + for (const instr of block.instructions) { + const dependencies: NodeDependencies = new Map(); + for (const operand of eachInstructionValueOperand(instr.value)) { + const dep = builder.lookup(operand.identifier, operand.loc); + dependencies.set(dep.node, { + from: {...dep.from}, + as: {...operand}, + }); + } + const node: InstructionNode = { + kind: 'Value', + controlDependency: null, + dependencies, + id: builder.nextReactiveId, + loc: instr.loc, + outputs: [], + value: instr, + }; + builder.nodes.set(node.id, node); + lastNode = node; + for (const lvalue of eachInstructionLValue(instr)) { + builder.declare(node.id, lvalue); + } + } + + // handle the terminal + const terminal = block.terminal; + switch (terminal.kind) { + case 'if': { + const testDep = builder.lookup( + terminal.test.identifier, + terminal.test.loc, + ); + const test: NodeReference = { + node: testDep.node, + from: testDep.from, + as: {...terminal.test}, + }; + const consequent = buildBlockScope(fn, builder, terminal.consequent); + const alternate = buildBlockScope(fn, builder, terminal.alternate); + const ifNode: IfNode = { + kind: 'If', + alternate, + consequent, + id: builder.nextReactiveId, + loc: terminal.loc, + outputs: [], + test, + }; + builder.nodes.set(ifNode.id, ifNode); + lastNode = ifNode; + break; + } + case 'return': { + const valueDep = builder.lookup( + terminal.value.identifier, + terminal.value.loc, + ); + const value: NodeReference = { + node: valueDep.node, + from: valueDep.from, + as: {...terminal.value}, + }; + const returnNode: ReturnNode = { + kind: 'Return', + id: builder.nextReactiveId, + loc: terminal.loc, + outputs: [], + value, + }; + builder.nodes.set(returnNode.id, returnNode); + lastNode = returnNode; + break; + } + case 'scope': { + const body = buildBlockScope(fn, builder, terminal.block); + const scopeNode: ScopeNode = { + kind: 'Scope', + body, + dependencies: new Map(), + id: builder.nextReactiveId, + loc: terminal.scope.loc, + outputs: [], + scope: terminal.scope, + }; + builder.nodes.set(scopeNode.id, scopeNode); + lastNode = scopeNode; + break; + } + case 'goto': { + break; + } + default: { + CompilerError.throwTodo({ + reason: `Support ${terminal.kind} nodes`, + loc: terminal.loc, + }); + } + } + + // Continue iteration in the fallthrough + const fallthrough = terminalFallthrough(terminal); + if (fallthrough != null) { + block = fn.body.blocks.get(fallthrough)!; + } else { + break; + } + } + return lastNode.id; +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts new file mode 100644 index 0000000000000..ae7b14a3455b5 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveIR/ReactiveIR.ts @@ -0,0 +1,352 @@ +/** + * 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 {CompilerError} from '..'; +import { + Environment, + Instruction, + Place, + ReactiveScope, + SourceLocation, + SpreadPattern, +} from '../HIR'; +import {ReactFunctionType} from '../HIR/Environment'; +import {printInstruction, printPlace} from '../HIR/PrintHIR'; +import {assertExhaustive} from '../Utils/utils'; + +export type ReactiveGraph = { + nodes: Map; + nextNodeId: number; + exit: ReactiveId; + loc: SourceLocation; + id: string | null; + params: Array; + generator: boolean; + async: boolean; + env: Environment; + directives: Array; + fnType: ReactFunctionType; +}; + +/* + * Simulated opaque type for Reactive IDs to prevent using normal numbers as ids + * accidentally. + */ +const opaqueReactiveId = Symbol(); +export type ReactiveId = number & {[opaqueReactiveId]: 'ReactiveId'}; + +export function makeReactiveId(id: number): ReactiveId { + CompilerError.invariant(id >= 0 && Number.isInteger(id), { + reason: 'Expected reactive node id to be a non-negative integer', + description: null, + loc: null, + suggestions: null, + }); + return id as ReactiveId; +} + +export type ReactiveNode = + | LoadArgumentNode + | InstructionNode + | IfNode + | EmptyNode + | ReturnNode + | ScopeNode; + +export type NodeReference = { + node: ReactiveId; + from: Place; + as: Place; +}; + +export type NodeDependencies = Map; +export type NodeDependency = {from: Place; as: Place}; + +export type LoadArgumentNode = { + kind: 'LoadArgument'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; + place: Place; +}; + +// An individual instruction +export type InstructionNode = { + kind: 'Value'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; + dependencies: NodeDependencies; + controlDependency: ReactiveId | null; + value: Instruction; +}; + +export type ReturnNode = { + kind: 'Return'; + id: ReactiveId; + loc: SourceLocation; + value: NodeReference; + outputs: Array; +}; + +export type IfNode = { + kind: 'If'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; + test: NodeReference; + consequent: ReactiveId; + alternate: ReactiveId; +}; + +export type EmptyNode = { + kind: 'Empty'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; +}; + +export type ScopeNode = { + kind: 'Scope'; + id: ReactiveId; + loc: SourceLocation; + outputs: Array; + scope: ReactiveScope; + /** + * The hoisted dependencies of the scope. Instructions "within" the scope + * (ie, the declarations or their deps) will also depend on these same values + * but we explicitly describe them here to ensure that all deps come before the scope + */ + dependencies: NodeDependencies; + /** + * The nodes that produce the values declared by the scope + */ + // declarations: NodeDependencies; + body: ReactiveId; +}; + +function _staticInvariantReactiveNodeHasIdLocationAndOutputs( + node: ReactiveNode, +): [ReactiveId, SourceLocation, Array] { + // If this fails, it is because a variant of ReactiveNode is missing a .id and/or .loc - add it! + return [node.id, node.loc, node.outputs]; +} + +/** + * Populates the outputs of each node in the graph + */ +export function populateReactiveGraphNodeOutputs(graph: ReactiveGraph): void { + // Populate node outputs + for (const [, node] of graph.nodes) { + node.outputs.length = 0; + } + for (const [, node] of graph.nodes) { + for (const dep of eachNodeDependency(node)) { + const sourceNode = graph.nodes.get(dep); + CompilerError.invariant(sourceNode != null, { + reason: `Expected source dependency ${dep} to exist`, + loc: node.loc, + }); + sourceNode.outputs.push(node.id); + } + } + const exitNode = graph.nodes.get(graph.exit)!; + exitNode.outputs.push(graph.exit); +} + +/** + * Puts the nodes of the graph into reverse postorder, such that nodes + * appear before any of their "successors" (consumers/dependents). + */ +export function reversePostorderReactiveGraph(graph: ReactiveGraph): void { + const nodes: Map = new Map(); + function visit(id: ReactiveId): void { + if (nodes.has(id)) { + return; + } + const node = graph.nodes.get(id); + CompilerError.invariant(node != null, { + reason: `Missing definition for ID ${id}`, + loc: null, + }); + for (const dep of eachNodeDependency(node)) { + visit(dep); + } + nodes.set(id, node); + } + for (const [_id, node] of graph.nodes) { + if (node.outputs.length === 0 && node.kind !== 'Empty') { + visit(node.id); + } + } + visit(graph.exit); + graph.nodes = nodes; +} + +export function* eachNodeDependency(node: ReactiveNode): Iterable { + switch (node.kind) { + case 'LoadArgument': + case 'Empty': { + break; + } + case 'If': { + yield node.test.node; + yield node.consequent; + yield node.alternate; + break; + } + case 'Return': { + yield node.value.node; + break; + } + case 'Value': { + yield* [...node.dependencies.keys()]; + if (node.controlDependency != null) { + yield node.controlDependency; + } + break; + } + case 'Scope': { + yield* [...node.dependencies.keys()]; + // yield* [...node.declarations.keys()]; + yield node.body; + break; + } + default: { + assertExhaustive(node, `Unexpected node kind '${(node as any).kind}'`); + } + } +} + +export function* eachNodeReference( + node: ReactiveNode, +): Iterable { + switch (node.kind) { + case 'LoadArgument': + case 'Empty': { + break; + } + case 'Return': { + yield node.value; + break; + } + case 'If': { + yield node.test; + break; + } + case 'Value': { + yield* [...node.dependencies].map(([node, dep]) => ({ + node, + from: dep.from, + as: dep.as, + })); + break; + } + case 'Scope': { + yield* [...node.dependencies].map(([node, dep]) => ({ + node, + from: dep.from, + as: dep.as, + })); + // yield* [...node.declarations].map(([node, dep]) => ({ + // node, + // from: dep.from, + // as: dep.as, + // })); + break; + } + default: { + assertExhaustive(node, `Unexpected node kind '${(node as any).kind}'`); + } + } +} + +function printNodeReference({node, from, as}: NodeReference): string { + return `£${node}.${printPlace(from)} => ${printPlace(as)}`; +} + +export function printNodeDependencies(deps: NodeDependencies): string { + const buffer: Array = []; + for (const [id, dep] of deps) { + buffer.push(printNodeReference({node: id, from: dep.from, as: dep.as})); + } + return buffer.join(', '); +} + +export function printReactiveGraph(graph: ReactiveGraph): string { + const buffer: Array = []; + buffer.push( + `${graph.fnType} ${graph.id ?? ''}(` + + graph.params + .map(param => { + if (param.kind === 'Identifier') { + return printPlace(param); + } else { + return `...${printPlace(param.place)}`; + } + }) + .join(', ') + + ')', + ); + writeReactiveNodes(buffer, graph.nodes); + buffer.push(`Exit £${graph.exit}`); + return buffer.join('\n'); +} + +export function printReactiveNodes( + nodes: Map, +): string { + const buffer: Array = []; + writeReactiveNodes(buffer, nodes); + return buffer.join('\n'); +} + +function writeReactiveNodes( + buffer: Array, + nodes: Map, +): void { + for (const [id, node] of nodes) { + const deps = [...eachNodeReference(node)] + .map(id => printNodeReference(id)) + .join(' '); + switch (node.kind) { + case 'LoadArgument': { + buffer.push(`£${id} LoadArgument ${printPlace(node.place)}`); + break; + } + case 'Empty': { + buffer.push(`£${id} Empty deps=[${deps}]`); + break; + } + case 'Return': { + buffer.push(`£${id} Return ${printNodeReference(node.value)}`); + break; + } + case 'If': { + buffer.push( + `£${id} If test=${printNodeReference(node.test)} consequent=£${node.consequent} alternate=£${node.alternate}`, + ); + break; + } + case 'Value': { + buffer.push(`£${id} Intermediate deps=[${deps}]`); + buffer.push(' ' + printInstruction(node.value)); + break; + } + case 'Scope': { + buffer.push( + // `£${id} Scope @${node.scope.id} deps=[${printNodeDependencies(node.dependencies)}] declarations=[${printNodeDependencies(node.declarations)}]`, + `£${id} Scope @${node.scope.id} deps=[${printNodeDependencies(node.dependencies)}] body=£${node.body}`, + ); + break; + } + default: { + assertExhaustive(node, `Unexpected node kind ${(node as any).kind}`); + } + } + } +} diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-graph.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-graph.js new file mode 100644 index 0000000000000..ac76010e16d75 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reactive-graph.js @@ -0,0 +1,8 @@ +// @enableReactiveGraph +function Component(props) { + const elements = []; + if (props.value) { + elements.push(
{props.value}
); + } + return elements; +}