Improve MergeConsecutiveScopes
Rewrites the core logic of MergeConsecutiveScopes to be easier to follow and fix bugs. We now do a two-pass approach: * First we iterate block instructions to identify scopes which can be merged, without actually merging the instructions themselves. * Then we iterate again, copying instructions from the block either into the new output block, or into their merged scope, as appropriate. I think the simplicity here is worth the performance cost, and we can always revisit later as necessary.
Joe Savona committed
Oct 9, 2023 at 16:15 UTC
a1461016df10f2a5f3e0e6fb60b1820c55c1f5be
3 files changed
+250
-181
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/MergeConsecutiveScopes.ts
+198
-110
@@ -5,6 +5,7 @@
5
* LICENSE file in the root directory of this source tree.
6
*/
7
8
+import { CompilerError } from "..";
9
import {
10
IdentifierId,
11
InstructionId,
@@ -17,6 +18,8 @@ import {
18
ReactiveScopeDependency,
19
makeInstructionId,
20
} from "../HIR";
21
+import { assertExhaustive } from "../Utils/utils";
22
+import { printReactiveScopeSummary } from "./PrintReactiveFunction";
23
import {
24
ReactiveFunctionTransform,
25
ReactiveFunctionVisitor,
@@ -62,6 +65,13 @@ export function mergeConsecutiveScopes(fn: ReactiveFunction): void {
65
);
66
}
67
68
+const DEBUG: boolean = false;
69
+function log(msg: string): void {
70
+ if (DEBUG) {
71
+ console.log(msg);
72
+ }
73
+}
74
+
75
class FindLastUsageVisitor extends ReactiveFunctionVisitor<void> {
76
lastUsage: Map<IdentifierId, InstructionId> = new Map();
77
@@ -84,137 +94,194 @@ class Transform extends ReactiveFunctionTransform<void> {
94
}
95
96
override visitBlock(block: ReactiveBlock, state: void): void {
97
+ // Pass 1: visit nested blocks to potentially merge their scopes
98
this.traverseBlock(block, state);
99
89
- // The current reactive scope which is a candidate for subsequent scopes
90
- // to be merged into
91
- let currentScope: {
92
- // The scope itself
100
+ // Pass 2: identify scopes for merging
101
+ type MergedScope = {
102
scope: ReactiveScopeBlock;
94
- // the starting index within `block` of this scope (inclusive)
103
from: number;
96
- // the index within `block` of instructions which are merged into this
97
- // scope (exclusive)
104
to: number;
99
- // Whether this block has been emitted yet onto `nextInstructions`
100
- merged: boolean;
101
- } | null = null;
102
-
103
- // Tracks the lvalues of instructions which occur between reactive scopes
104
- // We can't merge two scopes if their intervening instructions are needed
105
- // by subsequent code
106
- const lvalues: Set<IdentifierId> = new Set();
107
-
108
- // The updated set of instructions for the block. Stays null until
109
- // we make changes (ie merge scopes)
110
- let nextInstructions: ReactiveBlock | null = null;
111
- // The maximum index within the original instructions that we have reached.
112
- // Used to avoid emitting duplicate instructions
113
- let maxIndex: number = 0;
114
-
115
- // Called when we find some instruction that cannot be merged into a
116
- // preceding scope, or we otherwise need to reset and not consider
117
- // the previous candidate scope to be mergeable anymore.
118
- function resetCurrentScope(index: number): void {
119
- if (nextInstructions !== null) {
120
- if (currentScope !== null && !currentScope.merged) {
121
- currentScope.merged = true;
122
- nextInstructions.push(block[currentScope.from]!);
123
- }
124
- if (currentScope !== null) {
125
- nextInstructions.push(...block.slice(currentScope.to, index));
126
- }
127
- // We can sometimes call resetCurrentScope twice for the same index,
128
- // such as when an instruction resets and then a subsequent scope also resets.
129
- // This is the only case in which we push instructions w/o gating on
130
- // `currentScope != null`, so we avoid duplicates by checking the max index
131
- // already emitted.
132
- if (index < block.length && index > maxIndex) {
133
- nextInstructions.push(block[index]!);
134
- maxIndex = index;
135
- }
105
+ lvalues: Set<IdentifierId>;
106
+ };
107
+ let current: MergedScope | null = null;
108
+ const merged: Array<MergedScope> = [];
109
+ function reset(): void {
110
+ CompilerError.invariant(current !== null, {
111
+ loc: null,
112
+ reason:
113
+ "MergeConsecutiveScopes: expected current scope to be non-null if reset()",
114
+ suggestions: null,
115
+ description: null,
116
+ });
117
+ if (current.to > current.from + 1) {
118
+ merged.push(current);
119
}
137
- currentScope = null;
120
+ current = null;
121
}
139
-
122
for (let i = 0; i < block.length; i++) {
123
const instr = block[i]!;
142
- if (instr.kind === "terminal") {
143
- // Don't merge scopes with terminals in between.
144
- // In theory we could allow certain types of terminals,
145
- // such as loops, but for simplicity we just skip all
146
- // cases with terminals
147
- resetCurrentScope(i);
148
- } else if (instr.kind === "instruction") {
149
- switch (instr.instruction.value.kind) {
150
- case "JSXText":
151
- case "Primitive":
152
- case "LoadLocal":
153
- case "PropertyLoad":
154
- case "ComputedLoad": {
155
- // Allow simple instructions between scopes
156
- if (currentScope === null && nextInstructions !== null) {
157
- nextInstructions.push(instr);
158
- } else if (
159
- currentScope !== null &&
160
- instr.instruction.lvalue !== null
161
- ) {
162
- lvalues.add(instr.instruction.lvalue.identifier.id);
124
+ switch (instr.kind) {
125
+ case "terminal": {
126
+ // For now we don't merge across terminals
127
+ if (current !== null) {
128
+ log(
129
+ `Reset scope @${current.scope.scope.id} from terminal [${instr.terminal.id}]`
130
+ );
131
+ reset();
132
+ }
133
+ break;
134
+ }
135
+ case "instruction": {
136
+ switch (instr.instruction.value.kind) {
137
+ case "ComputedLoad":
138
+ case "JSXText":
139
+ case "LoadLocal":
140
+ case "Primitive":
141
+ case "PropertyLoad": {
142
+ // We can merge two scopes if there are intervening instructions, but:
143
+ // - Only if the instructions are simple and it's okay to make them
144
+ // execute conditionally (hence allowing a conservative subset of value kinds)
145
+ // - The values produced are used at or before the next scope. If they are used
146
+ // later and we move them into the scope, then they wouldn't be accessible to
147
+ // subsequent code wo expanding the set of declarations, which we want to avoid
148
+ if (current !== null && instr.instruction.lvalue !== null) {
149
+ current.lvalues.add(instr.instruction.lvalue.identifier.id);
150
+ }
151
+ break;
152
+ }
153
+ default: {
154
+ // Other instructions are known to prevent merging, so we reset the scope if present
155
+ if (current !== null) {
156
+ log(
157
+ `Reset scope @${current.scope.scope.id} from instruction [${instr.instruction.id}]`
158
+ );
159
+ reset();
160
+ }
161
}
164
- break;
162
}
166
- default: {
167
- // skip merging if there are complex intermediate instructions
168
- resetCurrentScope(i);
163
+ break;
164
+ }
165
+ case "scope": {
166
+ if (
167
+ current !== null &&
168
+ canMergeScopes(current.scope, instr) &&
169
+ areLValuesLastUsedByScope(
170
+ instr.scope,
171
+ current.lvalues,
172
+ this.lastUsage
173
+ )
174
+ ) {
175
+ // The current and next scopes can merge!
176
+ log(
177
+ `Can merge scope @${current.scope.scope.id} with @${instr.scope.id}`
178
+ );
179
+ // Update the merged scope's range
180
+ current.scope.scope.range.end = makeInstructionId(
181
+ Math.max(current.scope.scope.range.end, instr.scope.range.end)
182
+ );
183
+ // Add declarations
184
+ for (const [key, value] of instr.scope.declarations) {
185
+ current.scope.scope.declarations.set(key, value);
186
+ }
187
+ // Then prune declarations - this removes declarations from the earlier
188
+ // scope that are last-used at or before the newly merged subsequent scope
189
+ updateScopeDeclarations(current.scope.scope, this.lastUsage);
190
+ current.to = i + 1;
191
+ // We already checked that intermediate values were used at-or-before the merged
192
+ // scoped, so we can reset
193
+ current.lvalues.clear();
194
+
195
+ if (!scopeAlwaysInvalidatesOnDependencyChanges(instr)) {
196
+ // The subsequent scope that we just merged isn't guaranteed to invalidate if its
197
+ // inputs change, so it is not a candidate for future merging
198
+ log(
199
+ ` but scope @${instr.scope.id} doesnt guaranteed invalidate so it cannot merge further`
200
+ );
201
+ reset();
202
+ }
203
+ } else {
204
+ // No previous scope, or the scope cannot merge
205
+ if (current !== null) {
206
+ // Reset if necessary
207
+ log(
208
+ `Reset scope @${current.scope.scope.id}, not mergeable with subsequent scope @${instr.scope.id}`
209
+ );
210
+ reset();
211
+ }
212
+ // Only set a new merge candidate if the scope is guaranteed to invalidate on changes
213
+ if (scopeAlwaysInvalidatesOnDependencyChanges(instr)) {
214
+ current = {
215
+ scope: instr,
216
+ from: i,
217
+ to: i + 1,
218
+ lvalues: new Set(),
219
+ };
220
+ } else {
221
+ log(
222
+ `scope @${instr.scope.id} doesnt guaranteed invalidate so it cannot merge further`
223
+ );
224
+ }
225
}
226
+ break;
227
}
171
- } else {
172
- if (
173
- currentScope !== null &&
174
- canMergeScopes(currentScope.scope, instr) &&
175
- // If there are intermediate instructions, we can only merge the scopes
176
- // if those intermediate instructions are all used by the second scope.
177
- // if not, merging them would make those values unavailable to subsequent
178
- // code by moving them inside a different block scope in the output.
179
- areLValuesLastUsedByScope(instr.scope, lvalues, this.lastUsage)
180
- ) {
181
- const intermediateInstructions = block.slice(currentScope.to, i);
182
- currentScope.scope.scope.range.end = makeInstructionId(
183
- Math.max(currentScope.scope.scope.range.end, instr.scope.range.end)
228
+ default: {
229
+ assertExhaustive(
230
+ instr,
231
+ `Unexpected instruction kind '${(instr as any).kind}'`
232
);
185
- currentScope.scope.instructions.push(...intermediateInstructions);
186
- currentScope.scope.instructions.push(...instr.instructions);
187
- for (const [key, value] of instr.scope.declarations) {
188
- currentScope.scope.scope.declarations.set(key, value);
189
- }
190
- if (nextInstructions === null) {
191
- nextInstructions = block.slice(0, currentScope.from);
192
- nextInstructions.push(currentScope.scope);
193
- currentScope.merged = true;
194
- }
195
- currentScope.to = i + 1;
196
- lvalues.clear();
197
- } else {
198
- resetCurrentScope(i - 1); // don't include the current scope
199
- currentScope = { scope: instr, from: i, to: i + 1, merged: false };
200
- lvalues.clear();
233
}
234
}
235
}
204
- if (currentScope !== null && nextInstructions !== null) {
205
- nextInstructions.push(...block.slice(currentScope.to, block.length));
236
+ if (current !== null) {
237
+ reset();
238
+ }
239
+ if (merged.length) {
240
+ log(`merged ${merged.length} scopes:`);
241
+ for (const entry of merged) {
242
+ log(
243
+ printReactiveScopeSummary(entry.scope.scope) +
244
+ ` from=${entry.from} to=${entry.to}`
245
+ );
246
+ }
247
}
248
208
- if (nextInstructions !== null) {
209
- for (const instr of nextInstructions) {
249
+ // Pass 3: optional: if scopes can be merged, merge them and update the block
250
+ if (merged.length === 0) {
251
+ // Nothing merged, nothing to do!
252
+ return;
253
+ }
254
+ const nextInstructions = [];
255
+ let index = 0;
256
+ for (const entry of merged) {
257
+ if (index < entry.from) {
258
+ nextInstructions.push(...block.slice(index, entry.from));
259
+ index = entry.from;
260
+ }
261
+ const mergedScope = block[entry.from]!;
262
+ CompilerError.invariant(mergedScope.kind === "scope", {
263
+ loc: null,
264
+ reason:
265
+ "MergeConsecutiveScopes: Expected scope starting index to be a scope",
266
+ description: null,
267
+ suggestions: null,
268
+ });
269
+ nextInstructions.push(mergedScope);
270
+ index++;
271
+ while (index < entry.to) {
272
+ const instr = block[index++]!;
273
if (instr.kind === "scope") {
211
- updateScopeDeclarations(instr.scope, this.lastUsage);
274
+ mergedScope.instructions.push(...instr.instructions);
275
+ } else {
276
+ mergedScope.instructions.push(instr);
277
}
278
}
214
-
215
- block.length = 0;
216
- block.push(...nextInstructions);
279
}
280
+ while (index < block.length) {
281
+ nextInstructions.push(block[index++]!);
282
+ }
283
+ block.length = 0;
284
+ block.push(...nextInstructions);
285
}
286
}
287
@@ -247,6 +314,7 @@ function areLValuesLastUsedByScope(
314
for (const lvalue of lvalues) {
315
const lastUsedAt = lastUsage.get(lvalue)!;
316
if (lastUsedAt >= scope.range.end) {
317
+ log(` lvalue ${lvalue} used after scope @${scope.id}, cannot merge`);
318
return false;
319
}
320
}
@@ -256,10 +324,12 @@ function areLValuesLastUsedByScope(
324
function canMergeScopes(a: ReactiveScopeBlock, b: ReactiveScopeBlock): boolean {
325
// Don't merge scopes with reassignments
326
if (a.scope.reassignments.size !== 0 || b.scope.reassignments.size !== 0) {
327
+ log(` cannot merge, has reassignments`);
328
return false;
329
}
330
// Merge scopes whose dependencies are identical
331
if (areEqualDependencies(a.scope.dependencies, b.scope.dependencies)) {
332
+ log(` canMergeScopes: dependencies are equal`);
333
return true;
334
}
335
// Merge scopes where the outputs of the previous scope are the inputs
@@ -278,11 +348,14 @@ function canMergeScopes(a: ReactiveScopeBlock, b: ReactiveScopeBlock): boolean {
348
}))
349
),
350
b.scope.dependencies
281
- ) &&
282
- scopeAlwaysInvalidatesOnDependencyChanges(a)
351
+ )
352
) {
353
+ log(` outputs of prev are input to current`);
354
return true;
355
}
356
+ log(` cannot merge scopes:`);
357
+ log(` ${printReactiveScopeSummary(a.scope)}`);
358
+ log(` ${printReactiveScopeSummary(b.scope)}`);
359
return false;
360
}
361
@@ -332,6 +405,13 @@ class DeclarationTypeVisitor extends ReactiveFunctionVisitor<void> {
405
this.scope = scope;
406
}
407
408
+ override visitScope(scope: ReactiveScopeBlock, state: void): void {
409
+ if (scope.scope.id !== this.scope.id) {
410
+ return;
411
+ }
412
+ this.traverseScope(scope, state);
413
+ }
414
+
415
override visitInstruction(
416
instruction: ReactiveInstruction,
417
state: void
@@ -343,6 +423,14 @@ class DeclarationTypeVisitor extends ReactiveFunctionVisitor<void> {
423
) {
424
// no lvalue or this instruction isn't directly constructing a
425
// scope output value, skip
426
+ log(
427
+ ` skip instruction lvalue=${
428
+ instruction.lvalue?.identifier.id
429
+ } declaration?=${
430
+ instruction.lvalue != null &&
431
+ this.scope.declarations.has(instruction.lvalue.identifier.id)
432
+ } scope=${printReactiveScopeSummary(this.scope)}`
433
+ );
434
return;
435
}
436
switch (instruction.value.kind) {
compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/PrintReactiveFunction.ts
+24
-13
@@ -8,6 +8,7 @@
8
import { CompilerError } from "../CompilerError";
9
import {
10
ReactiveFunction,
11
+ ReactiveScope,
12
ReactiveScopeBlock,
13
ReactiveScopeDependency,
14
ReactiveStatement,
@@ -40,23 +41,25 @@ export function printReactiveFunction(fn: ReactiveFunction): string {
41
return writer.complete();
42
}
43
44
+export function printReactiveScopeSummary(scope: ReactiveScope): string {
45
+ return `scope @${scope.id} [${scope.range.start}:${
46
+ scope.range.end
47
+ }] dependencies=[${Array.from(scope.dependencies)
48
+ .map((dep) => printDependency(dep))
49
+ .join(", ")}] declarations=[${Array.from(scope.declarations)
50
+ .map(([, decl]) =>
51
+ printIdentifier({ ...decl.identifier, scope: decl.scope })
52
+ )
53
+ .join(", ")}] reassignments=[${Array.from(scope.reassignments).map(
54
+ (reassign) => printIdentifier(reassign)
55
+ )}]`;
56
+}
57
+
58
export function writeReactiveBlock(
59
writer: Writer,
60
block: ReactiveScopeBlock
61
): void {
47
- writer.writeLine(
48
- `scope @${block.scope.id} [${block.scope.range.start}:${
49
- block.scope.range.end
50
- }] dependencies=[${Array.from(block.scope.dependencies)
51
- .map((dep) => printDependency(dep))
52
- .join(", ")}] declarations=[${Array.from(block.scope.declarations)
53
- .map(([, decl]) =>
54
- printIdentifier({ ...decl.identifier, scope: decl.scope })
55
- )
56
- .join(", ")}] reassignments=[${Array.from(block.scope.reassignments).map(
57
- (reassign) => printIdentifier(reassign)
58
- )}] {`
59
- );
62
+ writer.writeLine(`${printReactiveScopeSummary(block.scope)} {`);
63
writeReactiveInstructions(writer, block.instructions);
64
writer.writeLine("}");
65
}
@@ -68,6 +71,14 @@ function printDependency(dependency: ReactiveScopeDependency): string {
71
return `${identifier}${dependency.path.map((prop) => `.${prop}`).join("")}`;
72
}
73
74
+export function printReactiveInstructions(
75
+ instructions: Array<ReactiveStatement>
76
+): string {
77
+ const writer = new Writer();
78
+ writeReactiveInstructions(writer, instructions);
79
+ return writer.complete();
80
+}
81
+
82
export function writeReactiveInstructions(
83
writer: Writer,
84
instructions: Array<ReactiveStatement>
compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/merge-consecutive-scopes-objects.expect.md
+28
-58
@@ -39,7 +39,7 @@ import { unstable_useMemoCache as useMemoCache } from "react"; // @enableMergeCo
39
// prevent scome scopes from merging, which concealed a bug with the merging logic.
40
// By avoiding JSX we eliminate extraneous instructions and more accurately test the merging.
41
function Component(props) {
42
- const $ = useMemoCache(18);
42
+ const $ = useMemoCache(11);
43
const [state, setState] = useState(0);
44
let t0;
45
if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
@@ -51,82 +51,52 @@ function Component(props) {
51
const c_1 = $[1] !== state;
52
let t1;
53
if (c_1) {
54
- const t2 = [state];
55
- t1 = { children: t2 };
54
+ t1 = { component: "span", props: { children: [state] } };
55
$[1] = state;
56
$[2] = t1;
57
} else {
58
t1 = $[2];
59
}
61
- const c_3 = $[3] !== t2;
60
+ const c_3 = $[3] !== state;
61
+ let t2;
62
if (c_3) {
63
- t1 = { children: t2 };
64
- $[3] = t2;
65
- $[4] = t1;
63
+ t2 = () => setState(state + 1);
64
+ $[3] = state;
65
+ $[4] = t2;
66
} else {
67
- t1 = $[4];
67
+ t2 = $[4];
68
}
69
- const c_5 = $[5] !== t1;
69
let t3;
71
- if (c_5) {
72
- t3 = { component: "span", props: t1 };
73
- $[5] = t1;
74
- $[6] = t3;
70
+ if ($[5] === Symbol.for("react.memo_cache_sentinel")) {
71
+ t3 = ["increment"];
72
+ $[5] = t3;
73
} else {
76
- t3 = $[6];
74
+ t3 = $[5];
75
}
78
- const c_7 = $[7] !== state;
76
+ const c_6 = $[6] !== t2;
77
let t4;
80
- if (c_7) {
81
- t4 = () => setState(state + 1);
82
- $[7] = state;
83
- $[8] = t4;
78
+ if (c_6) {
79
+ t4 = {
80
+ component: "button",
81
+ props: { "data-testid": "button", onClick: t2, children: t3 },
82
+ };
83
+ $[6] = t2;
84
+ $[7] = t4;
85
} else {
85
- t4 = $[8];
86
+ t4 = $[7];
87
}
88
+ const c_8 = $[8] !== t1;
89
+ const c_9 = $[9] !== t4;
90
let t5;
88
- if ($[9] === Symbol.for("react.memo_cache_sentinel")) {
89
- t5 = ["increment"];
90
- $[9] = t5;
91
- } else {
92
- t5 = $[9];
93
- }
94
- if ($[10] === Symbol.for("react.memo_cache_sentinel")) {
95
- t5 = ["increment"];
91
+ if (c_8 || c_9) {
92
+ t5 = [t0, t1, t4];
93
+ $[8] = t1;
94
+ $[9] = t4;
95
$[10] = t5;
96
} else {
97
t5 = $[10];
98
}
100
- const c_11 = $[11] !== t4;
101
- let t6;
102
- if (c_11) {
103
- const t7 = { "data-testid": "button", onClick: t4, children: t5 };
104
- t6 = { component: "button", props: t7 };
105
- $[11] = t4;
106
- $[12] = t6;
107
- } else {
108
- t6 = $[12];
109
- }
110
- const c_13 = $[13] !== t7;
111
- if (c_13) {
112
- t6 = { component: "button", props: t7 };
113
- $[13] = t7;
114
- $[14] = t6;
115
- } else {
116
- t6 = $[14];
117
- }
118
- const c_15 = $[15] !== t3;
119
- const c_16 = $[16] !== t6;
120
- let t8;
121
- if (c_15 || c_16) {
122
- t8 = [t0, t3, t6];
123
- $[15] = t3;
124
- $[16] = t6;
125
- $[17] = t8;
126
- } else {
127
- t8 = $[17];
128
- }
129
- return t8;
99
+ return t5;
100
}
101
102
export const FIXTURE_ENTRYPOINT = {