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-source-id: aea6dab469acb1f20058b85cb6f9aafab5d167cd Pull Request resolved: https://github.com/facebook/react/pull/29781
Joe Savona committed
Jun 7, 2024 at 12:10 UTC
aa0930452b06cee00a03544f1fbe1a6888578801
12 files changed
+113
-6
compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
+8
-1
@@ -65,12 +65,19 @@ export type ReactiveScopeBlock = {
65
instructions: ReactiveBlock;
66
};
67
68
+export type PrunedReactiveScopeBlock = {
69
+ kind: "pruned-scope";
70
+ scope: ReactiveScope;
71
+ instructions: ReactiveBlock;
72
+};
73
+
74
export type ReactiveBlock = Array<ReactiveStatement>;
75
76
export type ReactiveStatement =
77
| ReactiveInstructionStatement
78
| ReactiveTerminalStatement
73
- | ReactiveScopeBlock;
79
+ | ReactiveScopeBlock
80
+ | PrunedReactiveScopeBlock;
81
82
export type ReactiveInstructionStatement = {
83
kind: "instruction";
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/BuildReactiveBlocks.ts
+1
@@ -183,6 +183,7 @@ function visitBlock(context: Context, block: ReactiveBlock): void {
183
context.append(stmt, stmt.label);
184
break;
185
}
186
+ case "pruned-scope":
187
case "scope": {
188
CompilerError.invariant(false, {
189
reason: "Expected the function to not have scopes already assigned",
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts
+5
@@ -400,6 +400,11 @@ function codegenBlockNoReset(
400
}
401
break;
402
}
403
+ case "pruned-scope": {
404
+ const scopeBlock = codegenBlockNoReset(cx, item.instructions);
405
+ statements.push(...scopeBlock.body);
406
+ break;
407
+ }
408
case "scope": {
409
const temp = new Map(cx.temp);
410
codegenReactiveScope(cx, statements, item.scope, item.instructions);
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenReactiveLoops.ts
+8
-1
@@ -34,7 +34,14 @@ class Transform extends ReactiveFunctionTransform<boolean> {
34
): Transformed<ReactiveStatement> {
35
this.visitScope(scope, isWithinLoop);
36
if (isWithinLoop) {
37
- return { kind: "replace-many", value: scope.instructions };
37
+ return {
38
+ kind: "replace",
39
+ value: {
40
+ kind: "pruned-scope",
41
+ scope: scope.scope,
42
+ instructions: scope.instructions,
43
+ },
44
+ };
45
} else {
46
return { kind: "keep" };
47
}
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/FlattenScopesWithHooksOrUse.ts
+8
-1
@@ -66,7 +66,14 @@ class Transform extends ReactiveFunctionTransform<State> {
66
this.visitScope(scope, innerState);
67
outerState.hasHook ||= innerState.hasHook;
68
if (innerState.hasHook) {
69
- return { kind: "replace-many", value: scope.instructions };
69
+ return {
70
+ kind: "replace",
71
+ value: {
72
+ kind: "pruned-scope",
73
+ scope: scope.scope,
74
+ instructions: scope.instructions,
75
+ },
76
+ };
77
} else {
78
return { kind: "keep" };
79
}
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts
+10
@@ -174,6 +174,16 @@ class Transform extends ReactiveFunctionTransform<ReactiveScopeDependencies | nu
174
}
175
break;
176
}
177
+ case "pruned-scope": {
178
+ // For now we don't merge across pruned scopes
179
+ if (current !== null) {
180
+ log(
181
+ `Reset scope @${current.block.scope.id} from pruned scope @${instr.scope.id}`
182
+ );
183
+ reset();
184
+ }
185
+ break;
186
+ }
187
case "instruction": {
188
switch (instr.instruction.value.kind) {
189
case "ComputedLoad":
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PrintReactiveFunction.ts
+14
@@ -7,6 +7,7 @@
7
8
import { CompilerError } from "../CompilerError";
9
import {
10
+ PrunedReactiveScopeBlock,
11
ReactiveFunction,
12
ReactiveScope,
13
ReactiveScopeBlock,
@@ -83,6 +84,15 @@ export function writeReactiveBlock(
84
writer.writeLine("}");
85
}
86
87
+export function writePrunedScope(
88
+ writer: Writer,
89
+ block: PrunedReactiveScopeBlock
90
+): void {
91
+ writer.writeLine(`<pruned> ${printReactiveScopeSummary(block.scope)} {`);
92
+ writeReactiveInstructions(writer, block.instructions);
93
+ writer.writeLine("}");
94
+}
95
+
96
export function printDependency(dependency: ReactiveScopeDependency): string {
97
const identifier =
98
printIdentifier(dependency.identifier) +
@@ -133,6 +143,10 @@ function writeReactiveInstruction(
143
writeReactiveBlock(writer, instr);
144
break;
145
}
146
+ case "pruned-scope": {
147
+ writePrunedScope(writer, instr);
148
+ break;
149
+ }
150
case "terminal": {
151
if (instr.label !== null) {
152
writer.write(`bb${instr.label.id}: `);
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneAlwaysInvalidatingScopes.ts
+8
-1
@@ -107,7 +107,14 @@ class Transform extends ReactiveFunctionTransform<boolean> {
107
this.unmemoizedValues.add(identifier);
108
}
109
}
110
- return { kind: "replace-many", value: scopeBlock.instructions };
110
+ return {
111
+ kind: "replace",
112
+ value: {
113
+ kind: "pruned-scope",
114
+ scope: scopeBlock.scope,
115
+ instructions: scopeBlock.instructions,
116
+ },
117
+ };
118
}
119
}
120
return { kind: "keep" };
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts
+8
-1
@@ -951,7 +951,14 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
951
return { kind: "keep" };
952
} else {
953
this.prunedScopes.add(scopeBlock.scope.id);
954
- return { kind: "replace-many", value: scopeBlock.instructions };
954
+ return {
955
+ kind: "replace",
956
+ value: {
957
+ kind: "pruned-scope",
958
+ scope: scopeBlock.scope,
959
+ instructions: scopeBlock.instructions,
960
+ },
961
+ };
962
}
963
}
964
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneUnusedScopes.ts
+8
-1
@@ -51,7 +51,14 @@ class Transform extends ReactiveFunctionTransform<State> {
51
*/
52
!hasOwnDeclaration(scopeBlock))
53
) {
54
- return { kind: "replace-many", value: scopeBlock.instructions };
54
+ return {
55
+ kind: "replace",
56
+ value: {
57
+ kind: "pruned-scope",
58
+ scope: scopeBlock.scope,
59
+ instructions: scopeBlock.instructions,
60
+ },
61
+ };
62
} else {
63
return { kind: "keep" };
64
}
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/RenameVariables.ts
+8
@@ -12,6 +12,7 @@ import {
12
IdentifierName,
13
InstructionId,
14
Place,
15
+ PrunedReactiveScopeBlock,
16
ReactiveBlock,
17
ReactiveFunction,
18
ReactiveScopeBlock,
@@ -84,6 +85,13 @@ class Visitor extends ReactiveFunctionVisitor<Scopes> {
85
});
86
}
87
88
+ override visitPrunedScope(
89
+ scopeBlock: PrunedReactiveScopeBlock,
90
+ state: Scopes
91
+ ): void {
92
+ this.traverseBlock(scopeBlock.instructions, state);
93
+ }
94
+
95
override visitScope(scope: ReactiveScopeBlock, state: Scopes): void {
96
for (const [_, declaration] of scope.scope.declarations) {
97
state.visit(declaration.identifier);
compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/visitors.ts
+27
@@ -9,6 +9,7 @@ import {
9
HIRFunction,
10
InstructionId,
11
Place,
12
+ PrunedReactiveScopeBlock,
13
ReactiveBlock,
14
ReactiveFunction,
15
ReactiveInstruction,
@@ -196,6 +197,16 @@ export class ReactiveFunctionVisitor<TState = void> {
197
this.visitBlock(scope.instructions, state);
198
}
199
200
+ visitPrunedScope(scopeBlock: PrunedReactiveScopeBlock, state: TState): void {
201
+ this.traversePrunedScope(scopeBlock, state);
202
+ }
203
+ traversePrunedScope(
204
+ scopeBlock: PrunedReactiveScopeBlock,
205
+ state: TState
206
+ ): void {
207
+ this.visitBlock(scopeBlock.instructions, state);
208
+ }
209
+
210
visitBlock(block: ReactiveBlock, state: TState): void {
211
this.traverseBlock(block, state);
212
}
@@ -210,6 +221,10 @@ export class ReactiveFunctionVisitor<TState = void> {
221
this.visitScope(instr, state);
222
break;
223
}
224
+ case "pruned-scope": {
225
+ this.visitPrunedScope(instr, state);
226
+ break;
227
+ }
228
case "terminal": {
229
this.visitTerminal(instr, state);
230
break;
@@ -273,6 +288,10 @@ export class ReactiveFunctionTransform<
288
transformed = this.transformScope(instr, state);
289
break;
290
}
291
+ case "pruned-scope": {
292
+ transformed = this.transformPrunedScope(instr, state);
293
+ break;
294
+ }
295
case "terminal": {
296
transformed = this.transformTerminal(instr, state);
297
break;
@@ -339,6 +358,14 @@ export class ReactiveFunctionTransform<
358
return { kind: "keep" };
359
}
360
361
+ transformPrunedScope(
362
+ scope: PrunedReactiveScopeBlock,
363
+ state: TState
364
+ ): Transformed<ReactiveStatement> {
365
+ this.visitPrunedScope(scope, state);
366
+ return { kind: "keep" };
367
+ }
368
+
369
transformValue(
370
id: InstructionId,
371
value: ReactiveValue,