From 4fb7a93124090ece27002cb5c473fcd30a69f11b Mon Sep 17 00:00:00 2001 From: Joe Savona Date: Thu, 6 Jun 2024 09:31:30 -0700 Subject: [PATCH] compiler: Represent pruned scopes instead of inlining There are a few places where we want to check whether a value actually got memoized, and we currently have to infer this based on values that "should" have a scope and whether a corresponding scope actually exists. This PR adds a new ReactiveStatement variant to model a reactive scope block that was pruned for some reason, and updates all the passes that prune scopes to instead produce this new variant. [ghstack-poisoned] --- .../src/HIR/HIR.ts | 9 ++++++- .../src/ReactiveScopes/BuildReactiveBlocks.ts | 1 + .../ReactiveScopes/CodegenReactiveFunction.ts | 5 ++++ .../ReactiveScopes/FlattenReactiveLoops.ts | 9 ++++++- .../FlattenScopesWithHooksOrUse.ts | 9 ++++++- ...rgeReactiveScopesThatInvalidateTogether.ts | 10 +++++++ .../ReactiveScopes/PrintReactiveFunction.ts | 14 ++++++++++ .../PruneAlwaysInvalidatingScopes.ts | 9 ++++++- .../ReactiveScopes/PruneNonEscapingScopes.ts | 9 ++++++- .../src/ReactiveScopes/PruneUnusedScopes.ts | 9 ++++++- .../src/ReactiveScopes/RenameVariables.ts | 8 ++++++ .../src/ReactiveScopes/visitors.ts | 27 +++++++++++++++++++ 12 files changed, 113 insertions(+), 6 deletions(-) diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts index f9dfea52f363e..ad009194558be 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts @@ -65,12 +65,19 @@ export type ReactiveScopeBlock = { instructions: ReactiveBlock; }; +export type PrunedReactiveScopeBlock = { + kind: "pruned-scope"; + scope: ReactiveScope; + instructions: ReactiveBlock; +}; + export type ReactiveBlock = Array; export type ReactiveStatement = | ReactiveInstructionStatement | ReactiveTerminalStatement - | ReactiveScopeBlock; + | ReactiveScopeBlock + | PrunedReactiveScopeBlock; export type ReactiveInstructionStatement = { kind: "instruction"; diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveBlocks.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveBlocks.ts index fe59caeb0bfa6..3dee97030a899 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveBlocks.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveBlocks.ts @@ -183,6 +183,7 @@ function visitBlock(context: Context, block: ReactiveBlock): void { context.append(stmt, stmt.label); break; } + case "pruned-scope": case "scope": { CompilerError.invariant(false, { reason: "Expected the function to not have scopes already assigned", diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts index f43cae083176a..e0f5f899e8cbf 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts @@ -400,6 +400,11 @@ function codegenBlockNoReset( } break; } + case "pruned-scope": { + const scopeBlock = codegenBlockNoReset(cx, item.instructions); + statements.push(...scopeBlock.body); + break; + } case "scope": { const temp = new Map(cx.temp); codegenReactiveScope(cx, statements, item.scope, item.instructions); diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenReactiveLoops.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenReactiveLoops.ts index f544043c9cbbe..dd6e7b1cae3d2 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenReactiveLoops.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenReactiveLoops.ts @@ -34,7 +34,14 @@ class Transform extends ReactiveFunctionTransform { ): Transformed { this.visitScope(scope, isWithinLoop); if (isWithinLoop) { - return { kind: "replace-many", value: scope.instructions }; + return { + kind: "replace", + value: { + kind: "pruned-scope", + scope: scope.scope, + instructions: scope.instructions, + }, + }; } else { return { kind: "keep" }; } diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUse.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUse.ts index d9af65d5ec26c..06282baa0e362 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUse.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUse.ts @@ -66,7 +66,14 @@ class Transform extends ReactiveFunctionTransform { this.visitScope(scope, innerState); outerState.hasHook ||= innerState.hasHook; if (innerState.hasHook) { - return { kind: "replace-many", value: scope.instructions }; + return { + kind: "replace", + value: { + kind: "pruned-scope", + scope: scope.scope, + instructions: scope.instructions, + }, + }; } else { return { kind: "keep" }; } diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts index 149ba34bdf651..0509a3c65bead 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts @@ -174,6 +174,16 @@ class Transform extends ReactiveFunctionTransform ${printReactiveScopeSummary(block.scope)} {`); + writeReactiveInstructions(writer, block.instructions); + writer.writeLine("}"); +} + export function printDependency(dependency: ReactiveScopeDependency): string { const identifier = printIdentifier(dependency.identifier) + @@ -133,6 +143,10 @@ function writeReactiveInstruction( writeReactiveBlock(writer, instr); break; } + case "pruned-scope": { + writePrunedScope(writer, instr); + break; + } case "terminal": { if (instr.label !== null) { writer.write(`bb${instr.label.id}: `); diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts index 1227acca86bbb..cfbc934d6287c 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts @@ -107,7 +107,14 @@ class Transform extends ReactiveFunctionTransform { this.unmemoizedValues.add(identifier); } } - return { kind: "replace-many", value: scopeBlock.instructions }; + return { + kind: "replace", + value: { + kind: "pruned-scope", + scope: scopeBlock.scope, + instructions: scopeBlock.instructions, + }, + }; } } return { kind: "keep" }; diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts index cfa3700f81ea3..205eb07c7dd7d 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts @@ -950,7 +950,14 @@ class PruneScopesTransform extends ReactiveFunctionTransform< return { kind: "keep" }; } else { this.prunedScopes.add(scopeBlock.scope.id); - return { kind: "replace-many", value: scopeBlock.instructions }; + return { + kind: "replace", + value: { + kind: "pruned-scope", + scope: scopeBlock.scope, + instructions: scopeBlock.instructions, + }, + }; } } diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts index 0a3b712768873..11cc928874e3c 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts @@ -51,7 +51,14 @@ class Transform extends ReactiveFunctionTransform { */ !hasOwnDeclaration(scopeBlock)) ) { - return { kind: "replace-many", value: scopeBlock.instructions }; + return { + kind: "replace", + value: { + kind: "pruned-scope", + scope: scopeBlock.scope, + instructions: scopeBlock.instructions, + }, + }; } else { return { kind: "keep" }; } diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/RenameVariables.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/RenameVariables.ts index 6f5100bb7b23c..73994a052ba4c 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/RenameVariables.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/RenameVariables.ts @@ -12,6 +12,7 @@ import { IdentifierName, InstructionId, Place, + PrunedReactiveScopeBlock, ReactiveBlock, ReactiveFunction, ReactiveScopeBlock, @@ -84,6 +85,13 @@ class Visitor extends ReactiveFunctionVisitor { }); } + override visitPrunedScope( + scopeBlock: PrunedReactiveScopeBlock, + state: Scopes + ): void { + this.traverseBlock(scopeBlock.instructions, state); + } + override visitScope(scope: ReactiveScopeBlock, state: Scopes): void { for (const [_, declaration] of scope.scope.declarations) { state.visit(declaration.identifier); diff --git a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/visitors.ts b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/visitors.ts index 4d40f6537b496..f7b8f9cc62c05 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/visitors.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/visitors.ts @@ -9,6 +9,7 @@ import { HIRFunction, InstructionId, Place, + PrunedReactiveScopeBlock, ReactiveBlock, ReactiveFunction, ReactiveInstruction, @@ -196,6 +197,16 @@ export class ReactiveFunctionVisitor { this.visitBlock(scope.instructions, state); } + visitPrunedScope(scopeBlock: PrunedReactiveScopeBlock, state: TState): void { + this.traversePrunedScope(scopeBlock, state); + } + traversePrunedScope( + scopeBlock: PrunedReactiveScopeBlock, + state: TState + ): void { + this.visitBlock(scopeBlock.instructions, state); + } + visitBlock(block: ReactiveBlock, state: TState): void { this.traverseBlock(block, state); } @@ -210,6 +221,10 @@ export class ReactiveFunctionVisitor { this.visitScope(instr, state); break; } + case "pruned-scope": { + this.visitPrunedScope(instr, state); + break; + } case "terminal": { this.visitTerminal(instr, state); break; @@ -273,6 +288,10 @@ export class ReactiveFunctionTransform< transformed = this.transformScope(instr, state); break; } + case "pruned-scope": { + transformed = this.transformPrunedScope(instr, state); + break; + } case "terminal": { transformed = this.transformTerminal(instr, state); break; @@ -339,6 +358,14 @@ export class ReactiveFunctionTransform< return { kind: "keep" }; } + transformPrunedScope( + scope: PrunedReactiveScopeBlock, + state: TState + ): Transformed { + this.visitPrunedScope(scope, state); + return { kind: "keep" }; + } + transformValue( id: InstructionId, value: ReactiveValue,