[compiler] Collect temporaries and optional chains from inner functions (#31346)
Recursively collect identifier / property loads and optional chains from inner functions. This PR is in preparation for #31200 Previously, we only did this in `collectHoistablePropertyLoads` to understand hoistable property loads from inner functions. 1. collectTemporariesSidemap 2. collectOptionalChainSidemap 3. collectHoistablePropertyLoads - ^ this recursively calls `collectTemporariesSidemap`, `collectOptionalChainSidemap`, and `collectOptionalChainSidemap` on inner functions 4. collectDependencies Now, we have 1. collectTemporariesSidemap - recursively record identifiers in inner functions. Note that we track all temporaries in the same map as `IdentifierIds` are currently unique across functions 2. collectOptionalChainSidemap - recursively records optional chain sidemaps in inner functions 3. collectHoistablePropertyLoads - (unchanged, except to remove recursive collection of temporaries) 4. collectDependencies - unchanged: to be modified to recursively collect dependencies in next PR ' --- [//]: # (BEGIN SAPLING FOOTER) Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/facebook/react/pull/31346). * #31202 * #31203 * #31201 * #31200 * __->__ #31346 * #31199
mofeiZ committed
Nov 5, 2024 at 19:25 UTC
c3570b158d087eb4e3ee5748c4bd9360045c8a26
4 files changed
+139
-54
compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts
-9
@@ -8,7 +8,6 @@ import {
8
Set_union,
9
getOrInsertDefault,
10
} from '../Utils/utils';
11
-import {collectOptionalChainSidemap} from './CollectOptionalChainDependencies';
11
import {
12
BasicBlock,
13
BlockId,
@@ -22,7 +21,6 @@ import {
21
ReactiveScopeDependency,
22
ScopeId,
23
} from './HIR';
25
-import {collectTemporariesSidemap} from './PropagateScopeDependenciesHIR';
24
25
const DEBUG_PRINT = false;
26
@@ -373,17 +371,10 @@ function collectNonNullsInBlocks(
371
!fn.env.config.enableTreatFunctionDepsAsConditional
372
) {
373
const innerFn = instr.value.loweredFunc;
376
- const innerTemporaries = collectTemporariesSidemap(
377
- innerFn.func,
378
- new Set(),
379
- );
380
- const innerOptionals = collectOptionalChainSidemap(innerFn.func);
374
const innerHoistableMap = collectHoistablePropertyLoadsImpl(
375
innerFn.func,
376
{
377
...context,
385
- temporaries: innerTemporaries, // TODO: remove in later PR
386
- hoistableFromOptionals: innerOptionals.hoistableObjects, // TODO: remove in later PR
378
nestedFnImmutableContext:
379
context.nestedFnImmutableContext ??
380
new Set(
compiler/packages/babel-plugin-react-compiler/src/HIR/CollectOptionalChainDependencies.ts
+43
-23
@@ -3,7 +3,6 @@ import {assertNonNull} from './CollectHoistablePropertyLoads';
3
import {
4
BlockId,
5
BasicBlock,
6
- InstructionId,
6
IdentifierId,
7
ReactiveScopeDependency,
8
BranchTerminal,
@@ -15,6 +14,8 @@ import {
14
OptionalTerminal,
15
HIRFunction,
16
DependencyPathEntry,
17
+ Instruction,
18
+ Terminal,
19
} from './HIR';
20
import {printIdentifier} from './PrintHIR';
21
@@ -22,25 +23,14 @@ export function collectOptionalChainSidemap(
23
fn: HIRFunction,
24
): OptionalChainSidemap {
25
const context: OptionalTraversalContext = {
26
+ currFn: fn,
27
blocks: fn.body.blocks,
28
seenOptionals: new Set(),
29
processedInstrsInOptional: new Set(),
30
temporariesReadInOptional: new Map(),
31
hoistableObjects: new Map(),
32
};
31
- for (const [_, block] of fn.body.blocks) {
32
- if (
33
- block.terminal.kind === 'optional' &&
34
- !context.seenOptionals.has(block.id)
35
- ) {
36
- traverseOptionalBlock(
37
- block as TBasicBlock<OptionalTerminal>,
38
- context,
39
- null,
40
- );
41
- }
42
- }
43
-
33
+ traverseFunction(fn, context);
34
return {
35
temporariesReadInOptional: context.temporariesReadInOptional,
36
processedInstrsInOptional: context.processedInstrsInOptional,
@@ -96,8 +86,10 @@ export type OptionalChainSidemap = {
86
* bb5:
87
* $5 = MethodCall $2.$4() <--- here, we want to take a dep on $2 and $4!
88
* ```
89
+ *
90
+ * Also note that InstructionIds are not unique across inner functions.
91
*/
100
- processedInstrsInOptional: ReadonlySet<InstructionId>;
92
+ processedInstrsInOptional: ReadonlySet<Instruction | Terminal>;
93
/**
94
* Records optional chains for which we can safely evaluate non-optional
95
* PropertyLoads. e.g. given `a?.b.c`, we can evaluate any load from `a?.b` at
@@ -115,16 +107,46 @@ export type OptionalChainSidemap = {
107
};
108
109
type OptionalTraversalContext = {
110
+ currFn: HIRFunction;
111
blocks: ReadonlyMap<BlockId, BasicBlock>;
112
113
// Track optional blocks to avoid outer calls into nested optionals
114
seenOptionals: Set<BlockId>;
115
123
- processedInstrsInOptional: Set<InstructionId>;
116
+ processedInstrsInOptional: Set<Instruction | Terminal>;
117
temporariesReadInOptional: Map<IdentifierId, ReactiveScopeDependency>;
118
hoistableObjects: Map<BlockId, ReactiveScopeDependency>;
119
};
120
121
+function traverseFunction(
122
+ fn: HIRFunction,
123
+ context: OptionalTraversalContext,
124
+): void {
125
+ for (const [_, block] of fn.body.blocks) {
126
+ for (const instr of block.instructions) {
127
+ if (
128
+ instr.value.kind === 'FunctionExpression' ||
129
+ instr.value.kind === 'ObjectMethod'
130
+ ) {
131
+ traverseFunction(instr.value.loweredFunc.func, {
132
+ ...context,
133
+ currFn: instr.value.loweredFunc.func,
134
+ blocks: instr.value.loweredFunc.func.body.blocks,
135
+ });
136
+ }
137
+ }
138
+ if (
139
+ block.terminal.kind === 'optional' &&
140
+ !context.seenOptionals.has(block.id)
141
+ ) {
142
+ traverseOptionalBlock(
143
+ block as TBasicBlock<OptionalTerminal>,
144
+ context,
145
+ null,
146
+ );
147
+ }
148
+ }
149
+}
150
/**
151
* Match the consequent and alternate blocks of an optional.
152
* @returns propertyload computed by the consequent block, or null if the
@@ -137,7 +159,7 @@ function matchOptionalTestBlock(
159
consequentId: IdentifierId;
160
property: string;
161
propertyId: IdentifierId;
140
- storeLocalInstrId: InstructionId;
162
+ storeLocalInstr: Instruction;
163
consequentGoto: BlockId;
164
} | null {
165
const consequentBlock = assertNonNull(blocks.get(terminal.consequent));
@@ -149,7 +171,7 @@ function matchOptionalTestBlock(
171
const propertyLoad: TInstruction<PropertyLoad> = consequentBlock
172
.instructions[0] as TInstruction<PropertyLoad>;
173
const storeLocal: StoreLocal = consequentBlock.instructions[1].value;
152
- const storeLocalInstrId = consequentBlock.instructions[1].id;
174
+ const storeLocalInstr = consequentBlock.instructions[1];
175
CompilerError.invariant(
176
propertyLoad.value.object.identifier.id === terminal.test.identifier.id,
177
{
@@ -189,7 +211,7 @@ function matchOptionalTestBlock(
211
consequentId: storeLocal.lvalue.place.identifier.id,
212
property: propertyLoad.value.property,
213
propertyId: propertyLoad.lvalue.identifier.id,
192
- storeLocalInstrId,
214
+ storeLocalInstr,
215
consequentGoto: consequentBlock.terminal.block,
216
};
217
}
@@ -369,10 +391,8 @@ function traverseOptionalBlock(
391
},
392
],
393
};
372
- context.processedInstrsInOptional.add(
373
- matchConsequentResult.storeLocalInstrId,
374
- );
375
- context.processedInstrsInOptional.add(test.id);
394
+ context.processedInstrsInOptional.add(matchConsequentResult.storeLocalInstr);
395
+ context.processedInstrsInOptional.add(test);
396
context.temporariesReadInOptional.set(
397
matchConsequentResult.consequentId,
398
load,
compiler/packages/babel-plugin-react-compiler/src/HIR/PropagateScopeDependenciesHIR.ts
+88
-22
@@ -16,6 +16,7 @@ import {
16
DeclarationId,
17
areEqualPaths,
18
IdentifierId,
19
+ Terminal,
20
} from './HIR';
21
import {
22
collectHoistablePropertyLoads,
@@ -176,8 +177,10 @@ function findTemporariesUsedOutsideDeclaringScope(
177
* $2 = LoadLocal 'foo'
178
* $3 = CallExpression $2($1)
179
* ```
179
- * Only map LoadLocal and PropertyLoad lvalues to their source if we know that
180
- * reordering the read (from the time-of-load to time-of-use) is valid.
180
+ * @param usedOutsideDeclaringScope is used to check the correctness of
181
+ * reordering LoadLocal / PropertyLoad calls. We only track a LoadLocal /
182
+ * PropertyLoad in the returned temporaries map if reordering the read (from the
183
+ * time-of-load to time-of-use) is valid.
184
*
185
* If a LoadLocal or PropertyLoad instruction is within the reactive scope range
186
* (a proxy for mutable range) of the load source, later instructions may
@@ -215,7 +218,29 @@ export function collectTemporariesSidemap(
218
fn: HIRFunction,
219
usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
220
): ReadonlyMap<IdentifierId, ReactiveScopeDependency> {
218
- const temporaries = new Map<IdentifierId, ReactiveScopeDependency>();
221
+ const temporaries = new Map();
222
+ collectTemporariesSidemapImpl(
223
+ fn,
224
+ usedOutsideDeclaringScope,
225
+ temporaries,
226
+ false,
227
+ );
228
+ return temporaries;
229
+}
230
+
231
+/**
232
+ * Recursive collect a sidemap of all `LoadLocal` and `PropertyLoads` with a
233
+ * function and all nested functions.
234
+ *
235
+ * Note that IdentifierIds are currently unique, so we can use a single
236
+ * Map<IdentifierId, ...> across all nested functions.
237
+ */
238
+function collectTemporariesSidemapImpl(
239
+ fn: HIRFunction,
240
+ usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
241
+ temporaries: Map<IdentifierId, ReactiveScopeDependency>,
242
+ isInnerFn: boolean,
243
+): void {
244
for (const [_, block] of fn.body.blocks) {
245
for (const instr of block.instructions) {
246
const {value, lvalue} = instr;
@@ -224,27 +249,51 @@ export function collectTemporariesSidemap(
249
);
250
251
if (value.kind === 'PropertyLoad' && !usedOutside) {
227
- const property = getProperty(
228
- value.object,
229
- value.property,
230
- false,
231
- temporaries,
232
- );
233
- temporaries.set(lvalue.identifier.id, property);
252
+ if (!isInnerFn || temporaries.has(value.object.identifier.id)) {
253
+ /**
254
+ * All dependencies of a inner / nested function must have a base
255
+ * identifier from the outermost component / hook. This is because the
256
+ * compiler cannot break an inner function into multiple granular
257
+ * scopes.
258
+ */
259
+ const property = getProperty(
260
+ value.object,
261
+ value.property,
262
+ false,
263
+ temporaries,
264
+ );
265
+ temporaries.set(lvalue.identifier.id, property);
266
+ }
267
} else if (
268
value.kind === 'LoadLocal' &&
269
lvalue.identifier.name == null &&
270
value.place.identifier.name !== null &&
271
!usedOutside
272
) {
240
- temporaries.set(lvalue.identifier.id, {
241
- identifier: value.place.identifier,
242
- path: [],
243
- });
273
+ if (
274
+ !isInnerFn ||
275
+ fn.context.some(
276
+ context => context.identifier.id === value.place.identifier.id,
277
+ )
278
+ ) {
279
+ temporaries.set(lvalue.identifier.id, {
280
+ identifier: value.place.identifier,
281
+ path: [],
282
+ });
283
+ }
284
+ } else if (
285
+ value.kind === 'FunctionExpression' ||
286
+ value.kind === 'ObjectMethod'
287
+ ) {
288
+ collectTemporariesSidemapImpl(
289
+ value.loweredFunc.func,
290
+ usedOutsideDeclaringScope,
291
+ temporaries,
292
+ true,
293
+ );
294
}
295
}
296
}
247
- return temporaries;
297
}
298
299
function getProperty(
@@ -310,6 +359,12 @@ class Context {
359
#temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>;
360
#temporariesUsedOutsideScope: ReadonlySet<DeclarationId>;
361
362
+ /**
363
+ * Tracks the traversal state. See Context.declare for explanation of why this
364
+ * is needed.
365
+ */
366
+ inInnerFn: boolean = false;
367
+
368
constructor(
369
temporariesUsedOutsideScope: ReadonlySet<DeclarationId>,
370
temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
@@ -360,12 +415,23 @@ class Context {
415
}
416
417
/*
363
- * Records where a value was declared, and optionally, the scope where the value originated from.
364
- * This is later used to determine if a dependency should be added to a scope; if the current
365
- * scope we are visiting is the same scope where the value originates, it can't be a dependency
366
- * on itself.
418
+ * Records where a value was declared, and optionally, the scope where the
419
+ * value originated from. This is later used to determine if a dependency
420
+ * should be added to a scope; if the current scope we are visiting is the
421
+ * same scope where the value originates, it can't be a dependency on itself.
422
+ *
423
+ * Note that we do not track declarations or reassignments within inner
424
+ * functions for the following reasons:
425
+ * - inner functions cannot be split by scope boundaries and are guaranteed
426
+ * to consume their own declarations
427
+ * - reassignments within inner functions are tracked as context variables,
428
+ * which already have extended mutable ranges to account for reassignments
429
+ * - *most importantly* it's currently simply incorrect to compare inner
430
+ * function instruction ids (tracked by `decl`) with outer ones (as stored
431
+ * by root identifier mutable ranges).
432
*/
433
declare(identifier: Identifier, decl: Decl): void {
434
+ if (this.inInnerFn) return;
435
if (!this.#declarations.has(identifier.declarationId)) {
436
this.#declarations.set(identifier.declarationId, decl);
437
}
@@ -575,7 +641,7 @@ function collectDependencies(
641
fn: HIRFunction,
642
usedOutsideDeclaringScope: ReadonlySet<DeclarationId>,
643
temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
578
- processedInstrsInOptional: ReadonlySet<InstructionId>,
644
+ processedInstrsInOptional: ReadonlySet<Instruction | Terminal>,
645
): Map<ReactiveScope, Array<ReactiveScopeDependency>> {
646
const context = new Context(usedOutsideDeclaringScope, temporaries);
647
@@ -614,12 +680,12 @@ function collectDependencies(
680
}
681
}
682
for (const instr of block.instructions) {
617
- if (!processedInstrsInOptional.has(instr.id)) {
683
+ if (!processedInstrsInOptional.has(instr)) {
684
handleInstruction(instr, context);
685
}
686
}
687
622
- if (!processedInstrsInOptional.has(block.terminal.id)) {
688
+ if (!processedInstrsInOptional.has(block.terminal)) {
689
for (const place of eachTerminalOperand(block.terminal)) {
690
context.visitOperand(place);
691
}
compiler/packages/babel-plugin-react-compiler/src/HIR/visitors.ts
+8
@@ -1215,9 +1215,17 @@ export class ScopeBlockTraversal {
1215
}
1216
}
1217
1218
+ /**
1219
+ * @returns if the given scope is currently 'active', i.e. if the scope start
1220
+ * block but not the scope fallthrough has been recorded.
1221
+ */
1222
isScopeActive(scopeId: ScopeId): boolean {
1223
return this.#activeScopes.indexOf(scopeId) !== -1;
1224
}
1225
+
1226
+ /**
1227
+ * The current, innermost active scope.
1228
+ */
1229
get currentScope(): ScopeId | null {
1230
return this.#activeScopes.at(-1) ?? null;
1231
}