From 64b555c5fb67c7ec90b969bcae035fcef2b278c8 Mon Sep 17 00:00:00 2001 From: Mofei Zhang Date: Thu, 22 May 2025 15:44:35 -0400 Subject: [PATCH] [compiler] Prepare HIRBuilder to be used by later passes --- .../src/Entrypoint/Pipeline.ts | 1 + .../src/HIR/BuildHIR.ts | 25 ++++++++--------- .../src/HIR/Environment.ts | 5 +++- .../src/HIR/HIRBuilder.ts | 28 +++++++++++-------- 4 files changed, 34 insertions(+), 25 deletions(-) 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 831d1ca38054e..fe97c8d642f60 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -130,6 +130,7 @@ function run( mode, config, contextIdentifiers, + func, logger, filename, code, diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts index b9f82eea18e9f..cfb15fb595ccc 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/BuildHIR.ts @@ -70,12 +70,14 @@ import {BuiltInArrayId} from './ObjectShape'; export function lower( func: NodePath, env: Environment, + // Bindings captured from the outer function, in case lower() is called recursively (for lambdas) bindings: Bindings | null = null, capturedRefs: Array = [], - // the outermost function being compiled, in case lower() is called recursively (for lambdas) - parent: NodePath | null = null, ): Result { - const builder = new HIRBuilder(env, parent ?? func, bindings, capturedRefs); + const builder = new HIRBuilder(env, { + bindings, + context: capturedRefs, + }); const context: HIRFunction['context'] = []; for (const ref of capturedRefs ?? []) { @@ -215,7 +217,7 @@ export function lower( return Ok({ id, params, - fnType: parent == null ? env.fnType : 'Other', + fnType: bindings == null ? env.fnType : 'Other', returnTypeAnnotation: null, // TODO: extract the actual return type node if present returnType: makeType(), body: builder.build(), @@ -3417,7 +3419,7 @@ function lowerFunction( | t.ObjectMethod >, ): LoweredFunction | null { - const componentScope: Scope = builder.parentFunction.scope; + const componentScope: Scope = builder.environment.parentFunction.scope; const capturedContext = gatherCapturedContext(expr, componentScope); /* @@ -3428,13 +3430,10 @@ function lowerFunction( * This isn't a problem in practice because use Babel's scope analysis to * identify the correct references. */ - const lowering = lower( - expr, - builder.environment, - builder.bindings, - [...builder.context, ...capturedContext], - builder.parentFunction, - ); + const lowering = lower(expr, builder.environment, builder.bindings, [ + ...builder.context, + ...capturedContext, + ]); let loweredFunc: HIRFunction; if (lowering.isErr()) { lowering @@ -3456,7 +3455,7 @@ function lowerExpressionToTemporary( return lowerValueToTemporary(builder, value); } -function lowerValueToTemporary( +export function lowerValueToTemporary( builder: HIRBuilder, value: InstructionValue, ): Place { 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 6e6643cd1d68f..27b578b3c7834 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -47,7 +47,7 @@ import { ShapeRegistry, addHook, } from './ObjectShape'; -import {Scope as BabelScope} from '@babel/traverse'; +import {Scope as BabelScope, NodePath} from '@babel/traverse'; import {TypeSchema} from './TypeSchema'; export const ReactElementSymbolSchema = z.object({ @@ -675,6 +675,7 @@ export class Environment { #contextIdentifiers: Set; #hoistedIdentifiers: Set; + parentFunction: NodePath; constructor( scope: BabelScope, @@ -682,6 +683,7 @@ export class Environment { compilerMode: CompilerMode, config: EnvironmentConfig, contextIdentifiers: Set, + parentFunction: NodePath, // the outermost function being compiled logger: Logger | null, filename: string | null, code: string | null, @@ -740,6 +742,7 @@ export class Environment { this.#moduleTypes.set(REANIMATED_MODULE_NAME, reanimatedModuleType); } + this.parentFunction = parentFunction; this.#contextIdentifiers = contextIdentifiers; this.#hoistedIdentifiers = new Set(); } diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts index 44dd34b7d6cae..9ed37bb2fc85f 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts @@ -110,7 +110,6 @@ export default class HIRBuilder { #bindings: Bindings; #env: Environment; #exceptionHandlerStack: Array = []; - parentFunction: NodePath; errors: CompilerError = new CompilerError(); /** * Traversal context: counts the number of `fbt` tag parents @@ -136,16 +135,17 @@ export default class HIRBuilder { constructor( env: Environment, - parentFunction: NodePath, // the outermost function being compiled - bindings: Bindings | null = null, - context: Array | null = null, + options?: { + bindings?: Bindings | null; + context?: Array; + entryBlockKind?: BlockKind; + }, ) { this.#env = env; - this.#bindings = bindings ?? new Map(); - this.parentFunction = parentFunction; - this.#context = context ?? []; + this.#bindings = options?.bindings ?? new Map(); + this.#context = options?.context ?? []; this.#entry = makeBlockId(env.nextBlockId); - this.#current = newBlock(this.#entry, 'block'); + this.#current = newBlock(this.#entry, options?.entryBlockKind ?? 'block'); } currentBlockKind(): BlockKind { @@ -239,7 +239,7 @@ export default class HIRBuilder { // Check if the binding is from module scope const outerBinding = - this.parentFunction.scope.parent.getBinding(originalName); + this.#env.parentFunction.scope.parent.getBinding(originalName); if (babelBinding === outerBinding) { const path = babelBinding.path; if (path.isImportDefaultSpecifier()) { @@ -293,7 +293,7 @@ export default class HIRBuilder { const binding = this.#resolveBabelBinding(path); if (binding) { // Check if the binding is from module scope, if so return null - const outerBinding = this.parentFunction.scope.parent.getBinding( + const outerBinding = this.#env.parentFunction.scope.parent.getBinding( path.node.name, ); if (binding === outerBinding) { @@ -376,7 +376,7 @@ export default class HIRBuilder { } // Terminate the current block w the given terminal, and start a new block - terminate(terminal: Terminal, nextBlockKind: BlockKind | null): void { + terminate(terminal: Terminal, nextBlockKind: BlockKind | null): BlockId { const {id: blockId, kind, instructions} = this.#current; this.#completed.set(blockId, { kind, @@ -390,6 +390,7 @@ export default class HIRBuilder { const nextId = this.#env.nextBlockId; this.#current = newBlock(nextId, nextBlockKind); } + return blockId; } /* @@ -746,6 +747,11 @@ function getReversePostorderedBlocks(func: HIR): HIR['blocks'] { * (eg bb2 then bb1), we ensure that they get reversed back to the correct order. */ const block = func.blocks.get(blockId)!; + CompilerError.invariant(block != null, { + reason: '[HIRBuilder] Unexpected null block', + description: `expected block ${blockId} to exist`, + loc: GeneratedSource, + }); const successors = [...eachTerminalSuccessor(block.terminal)].reverse(); const fallthrough = terminalFallthrough(block.terminal);