1
+/**
2
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
3
+ *
4
+ * This source code is licensed under the MIT license found in the
5
+ * LICENSE file in the root directory of this source tree.
6
+ */
7
+
8
+import {
9
+ IdentifierId,
10
+ InstructionId,
11
+ Place,
12
+ ReactiveBlock,
13
+ ReactiveFunction,
14
+ ReactiveScope,
15
+ ReactiveScopeBlock,
16
+ ReactiveScopeDependency,
17
+} from "../HIR";
18
+import { assertExhaustive } from "../Utils/utils";
19
+import {
20
+ ReactiveFunctionTransform,
21
+ ReactiveFunctionVisitor,
22
+ visitReactiveFunction,
23
+} from "./visitors";
24
+
25
+/**
26
+ * The primary goal of this pass is to reduce memoization overhead, specifically:
27
+ * - Use fewer memo slots
28
+ * - Reduce the number of comparisons and other memoization-related instructions
29
+ *
30
+ * This is achieved by merging consecutive reactive scopes when the two scopes
31
+ * will always invalidate together. The idea is that if two scopes would always
32
+ * invalidate together, it's more efficient to group the scopes together to save
33
+ * on memoization overhead.
34
+ *
35
+ * This optimization is necessarily somewhat limited. First, we only merge
36
+ * scopes that are in the same (reactive) block, ie we don't merge across
37
+ * control-flow or block-scoping boundaries. Second, we can only merge scopes
38
+ * so long as any intermediate instructions are safe to memoize — specifically,
39
+ * as long as the values created by those instructions are only referenced by
40
+ * the second scope and not elsewhere. This is to avoid changing control-flow
41
+ * and to avoid increasing the number of scope outputs (which defeats the optimization).
42
+ *
43
+ * With that in mind we can apply the optimization in two cases. Given a block with
44
+ * scope A, some safe-to-memoize instructions I, and scope B, we can merge scopes when:
45
+ * - A and B have identical dependencies. This means they will invalidate together, so
46
+ * by merging the scopes we can avoid duplicate cache slots and duplicate checks of
47
+ * those dependencies.
48
+ * - The output of A is the input to B. Any invalidation of A will change its output
49
+ * which invalidates B, so we can similarly merge scopes. Note that this optimization
50
+ * may not be beneficial if the outupts of A are not guaranteed to change if its input
51
+ * changes, but in practice this is generally the case.
52
+ *
53
+ */
54
+export function mergeConsecutiveScopes(fn: ReactiveFunction): void {
55
+ visitReactiveFunction(fn, new Transform(), undefined);
56
+}
57
+
58
+class Transform extends ReactiveFunctionTransform<void> {
59
+ override visitBlock(block: ReactiveBlock, state: void): void {
60
+ this.traverseBlock(block, state);
61
+
62
+ // The current reactive scope which is a candidate for subsequent scopes
63
+ // to be merged into
64
+ let currentScope: {
65
+ // The scope itself
66
+ scope: ReactiveScopeBlock;
67
+ // the starting index within `block` of this scope (inclusive)
68
+ from: number;
69
+ // the index within `block` of instructions which are merged into this
70
+ // scope (exclusive)
71
+ to: number;
72
+ // Whether this block has been emitted yet onto `nextInstructions`
73
+ merged: boolean;
74
+ } | null = null;
75
+
76
+ // Tracks the lvalues of instructions which occur between reactive scopes
77
+ // We can't merge two scopes if their intervening instructions are needed
78
+ // by subsequent code
79
+ const lvalues: Set<IdentifierId> = new Set();
80
+
81
+ // The updated set of instructions for the block. Stays null until
82
+ // we make changes (ie merge scopes)
83
+ let nextInstructions: ReactiveBlock | null = null;
84
+
85
+ // Called when we find some instruction that cannot be merged into a
86
+ // preceding scope, or we otherwise need to reset and not consider
87
+ // the previous candidate scope to be mergeable anymore.
88
+ function resetCurrentScope(index: number): void {
89
+ if (nextInstructions !== null) {
90
+ if (currentScope !== null && !currentScope.merged) {
91
+ currentScope.merged = true;
92
+ nextInstructions.push(block[currentScope.from]!);
93
+ }
94
+ if (currentScope !== null) {
95
+ nextInstructions.push(...block.slice(currentScope.to, index));
96
+ }
97
+ if (index < block.length) {
98
+ nextInstructions.push(block[index]!);
99
+ }
100
+ }
101
+ currentScope = null;
102
+ }
103
+
104
+ for (let i = 0; i < block.length; i++) {
105
+ const instr = block[i]!;
106
+ if (instr.kind === "terminal") {
107
+ // Don't merge scopes with terminals in between.
108
+ // In theory we could allow certain types of terminals,
109
+ // such as loops, but for simplicity we just skip all
110
+ // cases with terminals
111
+ resetCurrentScope(i);
112
+ } else if (instr.kind === "instruction") {
113
+ switch (instr.instruction.value.kind) {
114
+ case "JSXText":
115
+ case "Primitive":
116
+ case "LoadLocal":
117
+ case "PropertyLoad":
118
+ case "ComputedLoad": {
119
+ // Allow simple instructions between scopes
120
+ if (currentScope === null && nextInstructions !== null) {
121
+ nextInstructions.push(instr);
122
+ } else if (
123
+ currentScope !== null &&
124
+ instr.instruction.lvalue !== null
125
+ ) {
126
+ lvalues.add(instr.instruction.lvalue.identifier.id);
127
+ }
128
+ break;
129
+ }
130
+ default: {
131
+ // skip merging if there are complex intermediate instructions
132
+ resetCurrentScope(i);
133
+ }
134
+ }
135
+ } else {
136
+ if (
137
+ currentScope !== null &&
138
+ canMergeScopes(currentScope.scope.scope, instr.scope) &&
139
+ // If there are intermediate instructions, we can only merge the scopes
140
+ // if those intermediate instructions are all used by the second scope.
141
+ // if not, merging them would make those values unavailable to subsequent
142
+ // code by moving them inside a different block scope in the output.
143
+ usesAllLvalues(instr.instructions, lvalues)
144
+ ) {
145
+ const intermediateInstructions = block.slice(currentScope.to, i);
146
+ currentScope.scope.instructions.push(...intermediateInstructions);
147
+ currentScope.scope.instructions.push(...instr.instructions);
148
+ for (const [key, value] of instr.scope.declarations) {
149
+ currentScope.scope.scope.declarations.set(key, value);
150
+ }
151
+ if (nextInstructions === null) {
152
+ nextInstructions = block.slice(0, currentScope.from);
153
+ nextInstructions.push(currentScope.scope);
154
+ currentScope.merged = true;
155
+ }
156
+ currentScope.to = i + 1;
157
+ lvalues.clear();
158
+ } else {
159
+ resetCurrentScope(i - 1); // don't include the current scope
160
+ currentScope = { scope: instr, from: i, to: i + 1, merged: false };
161
+ lvalues.clear();
162
+ }
163
+ }
164
+ }
165
+ if (currentScope !== null && nextInstructions !== null) {
166
+ nextInstructions.push(...block.slice(currentScope.to, block.length));
167
+ }
168
+
169
+ if (nextInstructions !== null) {
170
+ block.length = 0;
171
+ block.push(...nextInstructions);
172
+ }
173
+ }
174
+}
175
+
176
+function usesAllLvalues(
177
+ block: ReactiveBlock,
178
+ lvalues: Set<IdentifierId>
179
+): boolean {
180
+ if (lvalues.size === 0) {
181
+ return true;
182
+ }
183
+ const visitor = new OperandVisitor();
184
+ visitor.traverseBlock(block, lvalues);
185
+ return lvalues.size === 0;
186
+}
187
+
188
+class OperandVisitor extends ReactiveFunctionVisitor<Set<IdentifierId>> {
189
+ override visitPlace(
190
+ _id: InstructionId,
191
+ place: Place,
192
+ state: Set<IdentifierId>
193
+ ): void {
194
+ state.delete(place.identifier.id);
195
+ }
196
+
197
+ override traverseBlock(block: ReactiveBlock, state: Set<IdentifierId>): void {
198
+ for (const instr of block) {
199
+ if (state.size === 0) {
200
+ return;
201
+ }
202
+ switch (instr.kind) {
203
+ case "instruction": {
204
+ this.visitInstruction(instr.instruction, state);
205
+ break;
206
+ }
207
+ case "scope": {
208
+ this.visitScope(instr, state);
209
+ break;
210
+ }
211
+ case "terminal": {
212
+ this.visitTerminal(instr, state);
213
+ break;
214
+ }
215
+ default: {
216
+ assertExhaustive(
217
+ instr,
218
+ `Unexpected instruction kind '${(instr as any).kind}'`
219
+ );
220
+ }
221
+ }
222
+ }
223
+ }
224
+}
225
+
226
+function canMergeScopes(a: ReactiveScope, b: ReactiveScope): boolean {
227
+ // Don't merge scopes with reassignments
228
+ if (a.reassignments.size !== 0 || b.reassignments.size !== 0) {
229
+ return false;
230
+ }
231
+ if (areEqualDependencies(a.dependencies, b.dependencies)) {
232
+ return true;
233
+ }
234
+ if (
235
+ areEqualDependencies(
236
+ new Set(
237
+ [...a.declarations.values()].map((declaration) => ({
238
+ identifier: declaration.identifier,
239
+ path: [],
240
+ }))
241
+ ),
242
+ b.dependencies
243
+ )
244
+ ) {
245
+ return true;
246
+ }
247
+ return false;
248
+}
249
+
250
+function areEqualDependencies(
251
+ a: Set<ReactiveScopeDependency>,
252
+ b: Set<ReactiveScopeDependency>
253
+): boolean {
254
+ if (a.size !== b.size) {
255
+ return false;
256
+ }
257
+ for (const aValue of a) {
258
+ let found = false;
259
+ for (const bValue of b) {
260
+ if (
261
+ aValue.identifier === bValue.identifier &&
262
+ areEqualPaths(aValue.path, bValue.path)
263
+ ) {
264
+ found = true;
265
+ break;
266
+ }
267
+ }
268
+ if (!found) {
269
+ return false;
270
+ }
271
+ }
272
+ return true;
273
+}
274
+
275
+function areEqualPaths(a: Array<string>, b: Array<string>): boolean {
276
+ return a.length === b.length && a.every((item, ix) => item === b[ix]);
277
+}