[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: f6797d7652
Pull Request resolved: https://github.com/facebook/react/pull/30077

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts

@@ -280,7 +280,13 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
   scopeMapping = new Map();
   temporaries: Map<IdentifierId, ManualMemoDependency> = new Map();
 
-  collectMaybeMemoDependencies(
+  /**
+   * Recursively visit values and instructions to collect declarations
+   * and property loads.
+   * @returns a @{ManualMemoDependency} representing the variable +
+   * property reads represented by @value
+   */
+  recordDepsInValue(
     value: ReactiveValue,
     state: VisitorState
   ): ManualMemoDependency | null {
@@ -289,16 +295,28 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
         for (const instr of value.instructions) {
           this.visitInstruction(instr, state);
         }
-        const result = this.collectMaybeMemoDependencies(value.value, state);
-
+        const result = this.recordDepsInValue(value.value, state);
         return result;
       }
       case "OptionalExpression": {
-        return this.collectMaybeMemoDependencies(value.value, state);
+        return this.recordDepsInValue(value.value, state);
+      }
+      case "ReactiveFunctionValue": {
+        CompilerError.throwTodo({
+          reason:
+            "Handle ReactiveFunctionValue in ValidatePreserveManualMemoization",
+          loc: value.loc,
+        });
+      }
+      case "ConditionalExpression": {
+        this.recordDepsInValue(value.test, state);
+        this.recordDepsInValue(value.consequent, state);
+        this.recordDepsInValue(value.alternate, state);
+        return null;
       }
-      case "ReactiveFunctionValue":
-      case "ConditionalExpression":
       case "LogicalExpression": {
+        this.recordDepsInValue(value.left, state);
+        this.recordDepsInValue(value.right, state);
         return null;
       }
       default: {
@@ -336,7 +354,7 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
       state.manualMemoState.decls.add(lvalId);
     }
 
-    const maybeDep = this.collectMaybeMemoDependencies(value, state);
+    const maybeDep = this.recordDepsInValue(value, state);
     if (lvalId != null) {
       if (maybeDep != null) {
         temporaries.set(lvalId, maybeDep);
@@ -400,7 +418,10 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
     instruction: ReactiveInstruction,
     state: VisitorState
   ): void {
-    this.traverseInstruction(instruction, state);
+    /**
+     * We don't invoke traverseInstructions because `recordDepsInValue`
+     * recursively visits ReactiveValues and instructions
+     */
     this.recordTemporaries(instruction, state);
     if (instruction.value.kind === "StartMemoize") {
       let depsFromSource: Array<ManualMemoDependency> | null = null;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-slow-validate-preserve-memo.expect.md

@@ -0,0 +1,68 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import { Builder } from "shared-runtime";
+function useTest({ isNull, data }: { isNull: boolean; data: string }) {
+  const result = Builder.makeBuilder(isNull, "hello world")
+    ?.push("1", 2)
+    ?.push(3, {
+      a: 4,
+      b: 5,
+      c: data,
+    })
+    ?.push(6, data)
+    ?.push(7, "8")
+    ?.push("8", Builder.makeBuilder(!isNull)?.push(9).vals)?.vals;
+  return result;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [{ isNull: false, data: "param" }],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validatePreserveExistingMemoizationGuarantees
+
+import { Builder } from "shared-runtime";
+function useTest(t0) {
+  const $ = _c(3);
+  const { isNull, data } = t0;
+  let t1;
+  if ($[0] !== isNull || $[1] !== data) {
+    t1 = Builder.makeBuilder(isNull, "hello world")
+      ?.push("1", 2)
+      ?.push(3, { a: 4, b: 5, c: data })
+      ?.push(
+        6,
+
+        data,
+      )
+      ?.push(7, "8")
+      ?.push("8", Builder.makeBuilder(!isNull)?.push(9).vals)?.vals;
+    $[0] = isNull;
+    $[1] = data;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  const result = t1;
+  return result;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [{ isNull: false, data: "param" }],
+};
+
+```
+      
+### Eval output
+(kind: ok) ["hello world","1",2,3,{"a":4,"b":5,"c":"param"},6,"param",7,"8","8",null]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-slow-validate-preserve-memo.ts

@@ -0,0 +1,21 @@
+// @validatePreserveExistingMemoizationGuarantees
+
+import { Builder } from "shared-runtime";
+function useTest({ isNull, data }: { isNull: boolean; data: string }) {
+  const result = Builder.makeBuilder(isNull, "hello world")
+    ?.push("1", 2)
+    ?.push(3, {
+      a: 4,
+      b: 5,
+      c: data,
+    })
+    ?.push(6, data)
+    ?.push(7, "8")
+    ?.push("8", Builder.makeBuilder(!isNull)?.push(9).vals)?.vals;
+  return result;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useTest,
+  params: [{ isNull: false, data: "param" }],
+};

compiler/packages/snap/src/sprout/shared-runtime.ts

@@ -312,6 +312,22 @@ export function toJSON(value: any, invokeFns: boolean = false): string {
     return val;
   });
 }
+export class Builder {
+  vals: Array<any> = [];
+  static makeBuilder(isNull: boolean, ...args: Array<any>): Builder | null {
+    if (isNull) {
+      return null;
+    } else {
+      const builder = new Builder();
+      builder.push(...args);
+      return builder;
+    }
+  }
+  push(...args: Array<any>): Builder {
+    this.vals.push(...args);
+    return this;
+  }
+}
 
 export const ObjectWithHooks = {
   useFoo(): number {