[compiler][patch] Patch O(n^2) traversal in validatePreserveMemo
Double checked by syncing internally and verifying the # of `visitInstruction` calls with unique `InstructionId`s. This is a bit of an awkward pattern though. A cleaner alternative might be to override `visitValue` and store its results in a sidemap (instead of returning) ghstack-source-id: f6797d765224fb49c7d26cd377319662830d7348 Pull Request resolved: https://github.com/facebook/react/pull/30077
Mofei Zhang committed
Jun 25, 2024 at 16:06 UTC
4bfab0783204810cb51b9dda24464bb57777eb97
4 files changed
+134
-8
compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts
+29
-8
@@ -280,7 +280,13 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
280
scopeMapping = new Map();
281
temporaries: Map<IdentifierId, ManualMemoDependency> = new Map();
282
283
- collectMaybeMemoDependencies(
283
+ /**
284
+ * Recursively visit values and instructions to collect declarations
285
+ * and property loads.
286
+ * @returns a @{ManualMemoDependency} representing the variable +
287
+ * property reads represented by @value
288
+ */
289
+ recordDepsInValue(
290
value: ReactiveValue,
291
state: VisitorState
292
): ManualMemoDependency | null {
@@ -289,16 +295,28 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
295
for (const instr of value.instructions) {
296
this.visitInstruction(instr, state);
297
}
292
- const result = this.collectMaybeMemoDependencies(value.value, state);
293
-
298
+ const result = this.recordDepsInValue(value.value, state);
299
return result;
300
}
301
case "OptionalExpression": {
297
- return this.collectMaybeMemoDependencies(value.value, state);
302
+ return this.recordDepsInValue(value.value, state);
303
+ }
304
+ case "ReactiveFunctionValue": {
305
+ CompilerError.throwTodo({
306
+ reason:
307
+ "Handle ReactiveFunctionValue in ValidatePreserveManualMemoization",
308
+ loc: value.loc,
309
+ });
310
+ }
311
+ case "ConditionalExpression": {
312
+ this.recordDepsInValue(value.test, state);
313
+ this.recordDepsInValue(value.consequent, state);
314
+ this.recordDepsInValue(value.alternate, state);
315
+ return null;
316
}
299
- case "ReactiveFunctionValue":
300
- case "ConditionalExpression":
317
case "LogicalExpression": {
318
+ this.recordDepsInValue(value.left, state);
319
+ this.recordDepsInValue(value.right, state);
320
return null;
321
}
322
default: {
@@ -336,7 +354,7 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
354
state.manualMemoState.decls.add(lvalId);
355
}
356
339
- const maybeDep = this.collectMaybeMemoDependencies(value, state);
357
+ const maybeDep = this.recordDepsInValue(value, state);
358
if (lvalId != null) {
359
if (maybeDep != null) {
360
temporaries.set(lvalId, maybeDep);
@@ -400,7 +418,10 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
418
instruction: ReactiveInstruction,
419
state: VisitorState
420
): void {
403
- this.traverseInstruction(instruction, state);
421
+ /**
422
+ * We don't invoke traverseInstructions because `recordDepsInValue`
423
+ * recursively visits ReactiveValues and instructions
424
+ */
425
this.recordTemporaries(instruction, state);
426
if (instruction.value.kind === "StartMemoize") {
427
let depsFromSource: Array<ManualMemoDependency> | null = null;
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-slow-validate-preserve-memo.expect.md
new
+68
@@ -0,0 +1,68 @@
1
+
2
+## Input
3
+
4
+```javascript
5
+// @validatePreserveExistingMemoizationGuarantees
6
+
7
+import { Builder } from "shared-runtime";
8
+function useTest({ isNull, data }: { isNull: boolean; data: string }) {
9
+ const result = Builder.makeBuilder(isNull, "hello world")
10
+ ?.push("1", 2)
11
+ ?.push(3, {
12
+ a: 4,
13
+ b: 5,
14
+ c: data,
15
+ })
16
+ ?.push(6, data)
17
+ ?.push(7, "8")
18
+ ?.push("8", Builder.makeBuilder(!isNull)?.push(9).vals)?.vals;
19
+ return result;
20
+}
21
+
22
+export const FIXTURE_ENTRYPOINT = {
23
+ fn: useTest,
24
+ params: [{ isNull: false, data: "param" }],
25
+};
26
+
27
+```
28
+
29
+## Code
30
+
31
+```javascript
32
+import { c as _c } from "react/compiler-runtime"; // @validatePreserveExistingMemoizationGuarantees
33
+
34
+import { Builder } from "shared-runtime";
35
+function useTest(t0) {
36
+ const $ = _c(3);
37
+ const { isNull, data } = t0;
38
+ let t1;
39
+ if ($[0] !== isNull || $[1] !== data) {
40
+ t1 = Builder.makeBuilder(isNull, "hello world")
41
+ ?.push("1", 2)
42
+ ?.push(3, { a: 4, b: 5, c: data })
43
+ ?.push(
44
+ 6,
45
+
46
+ data,
47
+ )
48
+ ?.push(7, "8")
49
+ ?.push("8", Builder.makeBuilder(!isNull)?.push(9).vals)?.vals;
50
+ $[0] = isNull;
51
+ $[1] = data;
52
+ $[2] = t1;
53
+ } else {
54
+ t1 = $[2];
55
+ }
56
+ const result = t1;
57
+ return result;
58
+}
59
+
60
+export const FIXTURE_ENTRYPOINT = {
61
+ fn: useTest,
62
+ params: [{ isNull: false, data: "param" }],
63
+};
64
+
65
+```
66
+
67
+### Eval output
68
+(kind: ok) ["hello world","1",2,3,{"a":4,"b":5,"c":"param"},6,"param",7,"8","8",null]
\ No newline at end of file
compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-slow-validate-preserve-memo.ts
new
+21
@@ -0,0 +1,21 @@
1
+// @validatePreserveExistingMemoizationGuarantees
2
+
3
+import { Builder } from "shared-runtime";
4
+function useTest({ isNull, data }: { isNull: boolean; data: string }) {
5
+ const result = Builder.makeBuilder(isNull, "hello world")
6
+ ?.push("1", 2)
7
+ ?.push(3, {
8
+ a: 4,
9
+ b: 5,
10
+ c: data,
11
+ })
12
+ ?.push(6, data)
13
+ ?.push(7, "8")
14
+ ?.push("8", Builder.makeBuilder(!isNull)?.push(9).vals)?.vals;
15
+ return result;
16
+}
17
+
18
+export const FIXTURE_ENTRYPOINT = {
19
+ fn: useTest,
20
+ params: [{ isNull: false, data: "param" }],
21
+};
compiler/packages/snap/src/sprout/shared-runtime.ts
+16
@@ -312,6 +312,22 @@ export function toJSON(value: any, invokeFns: boolean = false): string {
312
return val;
313
});
314
}
315
+export class Builder {
316
+ vals: Array<any> = [];
317
+ static makeBuilder(isNull: boolean, ...args: Array<any>): Builder | null {
318
+ if (isNull) {
319
+ return null;
320
+ } else {
321
+ const builder = new Builder();
322
+ builder.push(...args);
323
+ return builder;
324
+ }
325
+ }
326
+ push(...args: Array<any>): Builder {
327
+ this.vals.push(...args);
328
+ return this;
329
+ }
330
+}
331
332
export const ObjectWithHooks = {
333
useFoo(): number {