1
+/**
2
+ * Copyright (c) Facebook, Inc. and its 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 invariant from "invariant";
9
+import prettyFormat from "pretty-format";
10
+import { CompilerError } from "../CompilerError";
11
+import {
12
+ Effect,
13
+ IdentifierId,
14
+ InstructionId,
15
+ Place,
16
+ ReactiveFunction,
17
+ ReactiveInstruction,
18
+ ReactiveScopeBlock,
19
+ ReactiveStatement,
20
+ ReactiveTerminal,
21
+ ReactiveTerminalStatement,
22
+ ReactiveValue,
23
+ ScopeId,
24
+} from "../HIR";
25
+import { eachInstructionLValue } from "../HIR/visitors";
26
+import { log } from "../Utils/logger";
27
+import { assertExhaustive } from "../Utils/utils";
28
+import { getPlaceScope } from "./BuildReactiveBlocks";
29
+import { printReactiveFunction } from "./PrintReactiveFunction";
30
+import {
31
+ eachReactiveValueOperand,
32
+ ReactiveFunctionTransform,
33
+ ReactiveFunctionVisitor,
34
+ Transformed,
35
+ visitReactiveFunction,
36
+} from "./visitors";
37
+
38
+/**
39
+ * This pass prunes reactive scopes that are not necessary to bound downstream computation.
40
+ * Specifically, the pass identifies the set of identifiers which are directly returned by
41
+ * the function and/or transitively aliased by a return value - ie, values that "escape".
42
+ *
43
+ * Example to build intuition:
44
+ *
45
+ * ```javascript
46
+ * function Component(props) {
47
+ * const a = {}; // not aliased or returned: *not* memoized
48
+ * const b = {}; // aliased by c, which is returned: memoized
49
+ * const c = [b]; // directly returned: memoized
50
+ * return c;
51
+ * }
52
+ * ```
53
+ *
54
+ * However, this logic alone is insufficient for two reasons:
55
+ * - Statically memoizing JSX elements *may* be inefficient compared to using dynamic
56
+ * memoization with `React.memo()`. Static memoization may be JIT'd and can look at
57
+ * the precise props w/o dynamic iteration, but incurs potentially large code-size
58
+ * overhead. Dynamic memoization with `React.memo()` incurs potentially increased
59
+ * runtime overhead for smaller code size. We plan to experiment with both variants
60
+ * for JSX.
61
+ * - Because we merge values whose mutations _interleave_ into a single scope, there
62
+ * can be cases where a non-escaping value needs to be memoized anyway to avoid breaking
63
+ * a memoization input. As a rule, for any scope that has a memoized output, all of that
64
+ * scope's transitive dependencies must also be memoized _even if they don't escape_.
65
+ * Failing to memoize them would cause the scope to invalidate more often than necessary
66
+ * and break downstream memoization.
67
+ *
68
+ * Example of this second case:
69
+ *
70
+ * ```javascript
71
+ * function Component(props) {
72
+ * // a can be independently memoized but it doesn't escape, so naively we may think its
73
+ * // safe to not memoize. but not memoizing would break caching of b, which does
74
+ * // escape.
75
+ * const a = [props.a];
76
+ *
77
+ * // b and c are interleaved and grouped into a single scope,
78
+ * // but they are independent values. c does not escape, but
79
+ * // we need to ensure that a is memoized or else b will invalidate
80
+ * // on every render since a is a dependency.
81
+ * const b = [];
82
+ * const c = {};
83
+ * c.a = a;
84
+ * b.push(props.b);
85
+ *
86
+ * return b;
87
+ * }
88
+ * ```
89
+ *
90
+ * ## Algorithm
91
+ *
92
+ * 1. First we build up a graph, a mapping of IdentifierId to a node describing all the
93
+ * scopes and inputs involved in creating that identifier. Individual nodes are marked
94
+ * as definitely aliased, conditionally aliased, or unaliased:
95
+ * a. Arrays, objects, function calls all produce a new value and are always marked as aliased
96
+ * b. Conditional and logical expressions (and a few others) are conditinally aliased,
97
+ * depending on whether their result value is aliased.
98
+ * c. JSX is always unaliased (though its props children may be)
99
+ * 2. The same pass which builds the graph also stores the set of returned identifiers.
100
+ * 3. We traverse the graph starting from the returned identifiers and mark reachable dependencies
101
+ * as escaping, based on the combination of the parent node's type and its children (eg a
102
+ * conditional node with an aliased dep promotes to aliased).
103
+ * 4. Finally we prune scopes whose outputs weren't marked.
104
+ */
105
+export function pruneNonEscapingScopes(fn: ReactiveFunction): void {
106
+ // First build up a map of which instructions are involved in creating which values,
107
+ // and which values are returned.
108
+ const state = new State();
109
+ if (fn.id !== null) {
110
+ state.declare(fn.id.id);
111
+ }
112
+ for (const param of fn.params) {
113
+ state.declare(param.identifier.id);
114
+ }
115
+ visitReactiveFunction(fn, new CollectDependenciesVisitor(), state);
116
+
117
+ log(() => prettyFormat(state));
118
+
119
+ // Then walk outward from the returned values and find all captured operands.
120
+ // This forms the set of identifiers which should be memoized.
121
+ const memoized = computeMemoizedIdentifiers(state);
122
+
123
+ log(() => prettyFormat(memoized));
124
+
125
+ log(() => printReactiveFunction(fn));
126
+
127
+ // Prune scopes that do not declare/reassign any escaping values
128
+ visitReactiveFunction(fn, new PruneScopesTransform(), memoized);
129
+}
130
+
131
+// Describes how to determine whether a value should be memoized, relative to dependees and dependencies
132
+enum MemoizationLevel {
133
+ // The value should be memoized if it escapes
134
+ Memoized = "Memoized",
135
+ // Values that are memoized if their dependencies are memoized (used for logical/ternary and
136
+ // other expressions that propagate dependencies wo changing them)
137
+ Conditional = "Conditional",
138
+ // Values that cannot be compared with Object.is, but which by default don't need to be memoized
139
+ // unless forced
140
+ Unmemoized = "Unmemoized",
141
+ // The value will never be memoized: used for values that can be cheaply compared w Object.is
142
+ Never = "Never",
143
+}
144
+
145
+// Given an identifier that appears as an lvalue multiple times with different memoization levels,
146
+// determines the final memoization level.
147
+function joinAliases(
148
+ kind1: MemoizationLevel,
149
+ kind2: MemoizationLevel
150
+): MemoizationLevel {
151
+ if (
152
+ kind1 === MemoizationLevel.Memoized ||
153
+ kind2 === MemoizationLevel.Memoized
154
+ ) {
155
+ return MemoizationLevel.Memoized;
156
+ } else if (
157
+ kind1 === MemoizationLevel.Conditional ||
158
+ kind2 === MemoizationLevel.Conditional
159
+ ) {
160
+ return MemoizationLevel.Conditional;
161
+ } else if (
162
+ kind1 === MemoizationLevel.Unmemoized ||
163
+ kind2 === MemoizationLevel.Unmemoized
164
+ ) {
165
+ return MemoizationLevel.Unmemoized;
166
+ } else {
167
+ return MemoizationLevel.Never;
168
+ }
169
+}
170
+
171
+// A node in the graph describing the memoization level of a given identifier as well as its dependencies and scopes.
172
+type IdentifierNode = {
173
+ level: MemoizationLevel;
174
+ memoized: boolean;
175
+ dependencies: Set<IdentifierId>;
176
+ scopes: Set<ScopeId>;
177
+ seen: boolean;
178
+};
179
+
180
+// A scope node describing its dependencies
181
+type ScopeNode = {
182
+ dependencies: Array<IdentifierId>;
183
+ seen: boolean;
184
+};
185
+
186
+// Stores the identifier and scope graphs, set of returned identifiers, etc
187
+class State {
188
+ // Maps lvalues for LoadLocal to the identifier being loaded, to resolve indirections
189
+ // in subsequent lvalues/rvalues
190
+ definitions: Map<IdentifierId, IdentifierId> = new Map();
191
+
192
+ identifiers: Map<IdentifierId, IdentifierNode> = new Map();
193
+ scopes: Map<ScopeId, ScopeNode> = new Map();
194
+ returned: Set<IdentifierId> = new Set();
195
+
196
+ /**
197
+ * Declare a new identifier, used for function id and params
198
+ */
199
+ declare(id: IdentifierId): void {
200
+ this.identifiers.set(id, {
201
+ level: MemoizationLevel.Never,
202
+ memoized: false,
203
+ dependencies: new Set(),
204
+ scopes: new Set(),
205
+ seen: false,
206
+ });
207
+ }
208
+
209
+ /**
210
+ * Associates the identifier with its scope, if there is one and it is active for the given instruction id:
211
+ * - Records the scope and its dependencies
212
+ * - Associates the identifier with this scope
213
+ */
214
+ visitOperand(
215
+ id: InstructionId,
216
+ place: Place,
217
+ identifier: IdentifierId
218
+ ): void {
219
+ const scope = getPlaceScope(id, place);
220
+ if (scope !== null) {
221
+ let node = this.scopes.get(scope.id);
222
+ if (node === undefined) {
223
+ node = {
224
+ dependencies: [...scope.dependencies].map((dep) => dep.identifier.id),
225
+ seen: false,
226
+ };
227
+ this.scopes.set(scope.id, node);
228
+ }
229
+ const identifierNode = this.identifiers.get(identifier);
230
+ invariant(
231
+ identifierNode !== undefined,
232
+ "Expected identifier to be initialized"
233
+ );
234
+ identifierNode.scopes.add(scope.id);
235
+ }
236
+ }
237
+}
238
+
239
+/**
240
+ * Given a state derived from visiting the function, walks the graph from the returned nodes
241
+ * to determine which other values should be memoized. Returns a set of all identifiers
242
+ * that should be memoized.
243
+ */
244
+function computeMemoizedIdentifiers(state: State): Set<IdentifierId> {
245
+ const memoized = new Set<IdentifierId>();
246
+
247
+ // Visit an identifier, optionally forcing it to be memoized
248
+ function visit(id: IdentifierId, forceMemoize: boolean = false): boolean {
249
+ const node = state.identifiers.get(id);
250
+ invariant(node !== undefined, "Expected a node for all identifiers");
251
+ if (node.seen) {
252
+ return node.memoized;
253
+ }
254
+ node.seen = true;
255
+
256
+ // Note: in case of cycles we temporarily mark the identifier as non-memoized,
257
+ // this is reset later after processing dependencies
258
+ node.memoized = false;
259
+
260
+ // Visit dependencies, determine if any of them are memoized
261
+ let hasMemoizedDependency = false;
262
+ for (const dep of node.dependencies) {
263
+ const isDepMemoized = visit(dep);
264
+ hasMemoizedDependency ||= isDepMemoized;
265
+ }
266
+
267
+ if (
268
+ node.level === MemoizationLevel.Memoized ||
269
+ (node.level === MemoizationLevel.Conditional &&
270
+ (hasMemoizedDependency || forceMemoize)) ||
271
+ (node.level === MemoizationLevel.Unmemoized && forceMemoize)
272
+ ) {
273
+ node.memoized = true;
274
+ memoized.add(id);
275
+ for (const scope of node.scopes) {
276
+ forceMemoizeScopeDependencies(scope);
277
+ }
278
+ }
279
+ return node.memoized;
280
+ }
281
+
282
+ // Force all the scope's optionally-memoizeable dependencies (not "Never") to be memoized
283
+ function forceMemoizeScopeDependencies(id: ScopeId): void {
284
+ const node = state.scopes.get(id);
285
+ invariant(node !== undefined, "Expected a node for all scopes");
286
+ if (node.seen) {
287
+ return;
288
+ }
289
+ node.seen = true;
290
+
291
+ for (const dep of node.dependencies) {
292
+ visit(dep, true);
293
+ }
294
+ return;
295
+ }
296
+
297
+ // Walk from the "roots" aka returned identifiers.
298
+ for (const returned of state.returned) {
299
+ visit(returned);
300
+ }
301
+
302
+ return memoized;
303
+}
304
+
305
+/**
306
+ * Given a value, returns a description of how it should be memoized:
307
+ * - lvalues: optional extra places that are lvalue-like in the sense of
308
+ * aliasing the rvalues
309
+ * - rvalues: places that are aliased by the instruction's lvalues.
310
+ * - level: the level of memoization to apply to this value
311
+ */
312
+function computeMemoizationInputs(value: ReactiveValue): {
313
+ // can optionally return a custom set of lvalues per instruction
314
+ lvalues: Array<Place> | null;
315
+ rvalues: Array<Place>;
316
+ level: MemoizationLevel;
317
+} {
318
+ switch (value.kind) {
319
+ case "ConditionalExpression": {
320
+ return {
321
+ lvalues: null,
322
+ rvalues: [
323
+ // Conditionals do not alias their test value.
324
+ ...computeMemoizationInputs(value.consequent).rvalues,
325
+ ...computeMemoizationInputs(value.alternate).rvalues,
326
+ ],
327
+ // Only need to memoize if the rvalues are memoized
328
+ level: MemoizationLevel.Conditional,
329
+ };
330
+ }
331
+ case "LogicalExpression": {
332
+ return {
333
+ lvalues: null,
334
+ rvalues: [
335
+ ...computeMemoizationInputs(value.left).rvalues,
336
+ ...computeMemoizationInputs(value.right).rvalues,
337
+ ],
338
+ // Only need to memoize if the rvalues are memoized
339
+ level: MemoizationLevel.Conditional,
340
+ };
341
+ }
342
+ case "SequenceExpression": {
343
+ return {
344
+ lvalues: null,
345
+ // Only the final value of the sequence is a true rvalue:
346
+ // values from the sequence's instructions are evaluated
347
+ // as separate nodes
348
+ rvalues: computeMemoizationInputs(value.value).rvalues,
349
+ // Only memoize if the final value was memoized
350
+ level: MemoizationLevel.Conditional,
351
+ };
352
+ }
353
+ case "JsxExpression": {
354
+ const operands: Array<Place> = [];
355
+ operands.push(value.tag);
356
+ for (const prop of value.props) {
357
+ if (prop.kind === "JsxAttribute") {
358
+ operands.push(prop.place);
359
+ } else {
360
+ operands.push(prop.argument);
361
+ }
362
+ }
363
+ if (value.children !== null) {
364
+ for (const child of value.children) {
365
+ operands.push(child);
366
+ }
367
+ }
368
+ return {
369
+ lvalues: null,
370
+ rvalues: operands,
371
+ // JSX elements themselves are not memoized unless forced to
372
+ // avoid breaking downstream memoization
373
+ level: MemoizationLevel.Unmemoized,
374
+ };
375
+ }
376
+ case "JsxFragment": {
377
+ return {
378
+ lvalues: null,
379
+ rvalues: value.children,
380
+ // JSX elements themselves are not memoized unless forced to
381
+ // avoid breaking downstream memoization
382
+ level: MemoizationLevel.Unmemoized,
383
+ };
384
+ }
385
+ case "ComputedDelete":
386
+ case "PropertyDelete":
387
+ case "LoadGlobal":
388
+ case "TemplateLiteral":
389
+ case "Primitive":
390
+ case "JSXText":
391
+ case "BinaryExpression":
392
+ case "UnaryExpression": {
393
+ return {
394
+ lvalues: null,
395
+ rvalues: [],
396
+ // All of these instructions return a primitive value and never need to be memoized
397
+ level: MemoizationLevel.Never,
398
+ };
399
+ }
400
+ case "TypeCastExpression": {
401
+ return {
402
+ lvalues: null,
403
+ // Indirection for the inner value, memoized if the value is
404
+ rvalues: [value.value],
405
+ level: MemoizationLevel.Conditional,
406
+ };
407
+ }
408
+ case "LoadLocal": {
409
+ return {
410
+ lvalues: null,
411
+ // Indirection for the inner value, memoized if the value is
412
+ rvalues: [value.place],
413
+ level: MemoizationLevel.Conditional,
414
+ };
415
+ }
416
+ case "Destructure":
417
+ case "StoreLocal": {
418
+ return {
419
+ lvalues: null,
420
+ // Indirection for the inner value, memoized if the value is
421
+ rvalues: [value.value],
422
+ level: MemoizationLevel.Conditional,
423
+ };
424
+ }
425
+ case "ComputedLoad":
426
+ case "PropertyLoad": {
427
+ return {
428
+ lvalues: null,
429
+ // Only the object is aliased to the result, and the result only needs to be
430
+ // memoized if the object is
431
+ rvalues: [value.object],
432
+ level: MemoizationLevel.Conditional,
433
+ };
434
+ }
435
+ case "ComputedStore": {
436
+ // The object being stored to acts as an lvalue (it aliases the value), but
437
+ // the computed key is not aliased
438
+ return {
439
+ lvalues: [value.object],
440
+ rvalues: [value.value],
441
+ level: MemoizationLevel.Conditional,
442
+ };
443
+ }
444
+ case "FunctionExpression":
445
+ case "TaggedTemplateExpression":
446
+ case "CallExpression":
447
+ case "ArrayExpression":
448
+ case "NewExpression":
449
+ case "ObjectExpression":
450
+ case "ComputedCall":
451
+ case "PropertyCall":
452
+ case "PropertyStore": {
453
+ // All of these instructions may produce new values which must be memoized if
454
+ // reachable from a return value. Any mutable rvalue may alias any other rvalue
455
+ const operands = [...eachReactiveValueOperand(value)];
456
+ return {
457
+ lvalues: operands.filter((operand) => isMutableEffect(operand.effect)),
458
+ rvalues: operands,
459
+ level: MemoizationLevel.Memoized,
460
+ };
461
+ }
462
+ case "UnsupportedNode": {
463
+ CompilerError.invariant(`Unexpected unsupported node`, value.loc);
464
+ }
465
+ default: {
466
+ assertExhaustive(value, `Unexpected value kind '${(value as any).kind}'`);
467
+ }
468
+ }
469
+}
470
+
471
+/**
472
+ * Populates the input state with the set of returned identifiers and information about each
473
+ * identifier's and scope's dependencies.
474
+ */
475
+class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
476
+ override visitInstruction(
477
+ instruction: ReactiveInstruction,
478
+ state: State
479
+ ): void {
480
+ this.traverseInstruction(instruction, state);
481
+
482
+ // Determe the level of memoization for this value and the lvalues/rvalues
483
+ const aliasing = computeMemoizationInputs(instruction.value);
484
+
485
+ // Associate all the rvalues with the instruction's scope if it has one
486
+ for (const operand of aliasing.rvalues) {
487
+ const operandId =
488
+ state.definitions.get(operand.identifier.id) ?? operand.identifier.id;
489
+ state.visitOperand(instruction.id, operand, operandId);
490
+ }
491
+
492
+ // Add the operands as dependencies of all lvalues.
493
+ const lvalues =
494
+ aliasing.lvalues !== null
495
+ ? [...eachInstructionLValue(instruction), ...aliasing.lvalues]
496
+ : [...eachInstructionLValue(instruction)];
497
+ for (const lvalue of lvalues) {
498
+ const lvalueId =
499
+ state.definitions.get(lvalue.identifier.id) ?? lvalue.identifier.id;
500
+ let node = state.identifiers.get(lvalueId);
501
+ if (node === undefined) {
502
+ node = {
503
+ level: MemoizationLevel.Never,
504
+ memoized: false,
505
+ dependencies: new Set(),
506
+ scopes: new Set(),
507
+ seen: false,
508
+ };
509
+ state.identifiers.set(lvalueId, node);
510
+ }
511
+ node.level = joinAliases(node.level, aliasing.level);
512
+ // This looks like NxM iterations but in practice all instructions with multiple
513
+ // lvalues have only a single rvalue
514
+ for (const operand of aliasing.rvalues) {
515
+ const operandId =
516
+ state.definitions.get(operand.identifier.id) ?? operand.identifier.id;
517
+ if (operandId === lvalueId) {
518
+ continue;
519
+ }
520
+ node.dependencies.add(operandId);
521
+ }
522
+
523
+ state.visitOperand(instruction.id, lvalue, lvalueId);
524
+ }
525
+
526
+ if (instruction.value.kind === "LoadLocal" && instruction.lvalue !== null) {
527
+ state.definitions.set(
528
+ instruction.lvalue.identifier.id,
529
+ instruction.value.place.identifier.id
530
+ );
531
+ }
532
+ }
533
+
534
+ override visitTerminal(
535
+ stmt: ReactiveTerminalStatement<ReactiveTerminal>,
536
+ state: State
537
+ ): void {
538
+ this.traverseTerminal(stmt, state);
539
+
540
+ if (stmt.terminal.kind === "return" && stmt.terminal.value !== null) {
541
+ state.returned.add(stmt.terminal.value.identifier.id);
542
+ }
543
+ }
544
+}
545
+
546
+/**
547
+ * Prune reactive scopes that do not have any memoized outputs
548
+ */
549
+class PruneScopesTransform extends ReactiveFunctionTransform<
550
+ Set<IdentifierId>
551
+> {
552
+ override transformScope(
553
+ scope: ReactiveScopeBlock,
554
+ state: Set<IdentifierId>
555
+ ): Transformed<ReactiveStatement> {
556
+ this.visitScope(scope, state);
557
+ const hasMemoizedOutput =
558
+ Array.from(scope.scope.declarations.keys()).some((id) => state.has(id)) ||
559
+ Array.from(scope.scope.reassignments).some((identifier) =>
560
+ state.has(identifier.id)
561
+ );
562
+ if (hasMemoizedOutput) {
563
+ return { kind: "keep" };
564
+ } else {
565
+ return { kind: "replace-many", value: scope.instructions };
566
+ }
567
+ }
568
+}
569
+
570
+function isMutableEffect(effect: Effect): boolean {
571
+ switch (effect) {
572
+ case Effect.Capture:
573
+ case Effect.Mutate:
574
+ case Effect.Store: {
575
+ return true;
576
+ }
577
+ default: {
578
+ return false;
579
+ }
580
+ }
581
+}