[compiler] Patch ValidatePreserveMemo to bailout correctly for refs

ghstack-source-id: 3f5d5c1
Pull Request resolved: #30603

Author: mofeiZ

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

@@ -30,6 +30,7 @@ import {
   ReactiveFunctionVisitor,
   visitReactiveFunction,
 } from '../ReactiveScopes/visitors';
+import {getOrInsertDefault} from '../Utils/utils';
 
 /**
  * Validates that all explicit manual memoization (useMemo/useCallback) was accurately
@@ -52,6 +53,16 @@ export function validatePreservedManualMemoization(fn: ReactiveFunction): void {
 const DEBUG = false;
 
 type ManualMemoBlockState = {
+  /**
+   * Tracks reassigned temporaries.
+   * This is necessary because useMemo calls are usually inlined.
+   * Inlining produces a `let` declaration, followed by reassignments
+   * to the newly declared variable (one per return statement).
+   * Since InferReactiveScopes does not merge scopes across reassigned
+   * variables (except in the case of a mutate-after-phi), we need to
+   * track reassignments to validate we're retaining manual memo.
+   */
+  reassignments: Map<DeclarationId, Set<Identifier>>;
   // The source of the original memoization, used when reporting errors
   loc: SourceLocation;
 
@@ -433,6 +444,19 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
      * recursively visits ReactiveValues and instructions
      */
     this.recordTemporaries(instruction, state);
+    const value = instruction.value;
+    if (
+      value.kind === 'StoreLocal' &&
+      value.lvalue.kind === 'Reassign' &&
+      state.manualMemoState != null
+    ) {
+      const ids = getOrInsertDefault(
+        state.manualMemoState.reassignments,
+        value.lvalue.place.identifier.declarationId,
+        new Set(),
+      );
+      ids.add(value.value.identifier);
+    }
     if (instruction.value.kind === 'StartMemoize') {
       let depsFromSource: Array<ManualMemoDependency> | null = null;
       if (instruction.value.deps != null) {
@@ -450,6 +474,7 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
         decls: new Set(),
         depsFromSource,
         manualMemoId: instruction.value.manualMemoId,
+        reassignments: new Map(),
       };
 
       for (const {identifier, loc} of eachInstructionValueOperand(
@@ -482,20 +507,34 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
           suggestions: null,
         },
       );
+      const reassignments = state.manualMemoState.reassignments;
       state.manualMemoState = null;
       if (!instruction.value.pruned) {
         for (const {identifier, loc} of eachInstructionValueOperand(
           instruction.value as InstructionValue,
         )) {
-          if (isUnmemoized(identifier, this.scopes)) {
-            state.errors.push({
-              reason:
-                'React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output.',
-              description: null,
-              severity: ErrorSeverity.CannotPreserveMemoization,
-              loc,
-              suggestions: null,
-            });
+          let manualMemoRvals;
+          if (identifier.scope == null) {
+            // If the manual memo was a useMemo that got inlined, iterate through
+            // all reassignments to the iife temporary to ensure they're memoized.
+            manualMemoRvals = reassignments.get(identifier.declarationId) ?? [
+              identifier,
+            ];
+          } else {
+            manualMemoRvals = [identifier];
+          }
+
+          for (const identifier of manualMemoRvals) {
+            if (isUnmemoized(identifier, this.scopes)) {
+              state.errors.push({
+                reason:
+                  'React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output.',
+                description: null,
+                severity: ErrorSeverity.CannotPreserveMemoization,
+                loc,
+                suggestions: null,
+              });
+            }
           }
         }
       }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-maybe-mutable-ref-memo-not-preserved.expect.md

@@ -0,0 +1,38 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {identity} from 'shared-runtime';
+
+/**
+ * This is technically a false positive, although it makes sense
+ * to bailout as source code might be doing something sketchy.
+ */
+function useFoo(x) {
+  useMemo(() => identity(x), [x]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [2],
+};
+
+```
+
+
+## Error
+
+```
+   9 |  */
+  10 | function useFoo(x) {
+> 11 |   useMemo(() => identity(x), [x]);
+     |   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (11:11)
+  12 | }
+  13 |
+  14 | export const FIXTURE_ENTRYPOINT = {
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-prune-nonescaping.expect.md

@@ -0,0 +1,17 @@
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {identity} from 'shared-runtime';
+
+/**
+ * This is technically a false positive, although it makes sense
+ * to bailout as source code might be doing something sketchy.
+ */
+function useFoo(x) {
+  useMemo(() => identity(x), [x]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [2],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-prune-nonescaping.ts

@@ -0,0 +1,42 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {useHook} from 'shared-runtime';
+
+// useMemo values may not be memoized in Forget output if we
+// infer that their deps always invalidate.
+// This is technically a false positive as the useMemo in source
+// was effectively a no-op
+function useFoo(props) {
+  const x = [];
+  useHook();
+  x.push(props);
+
+  return useMemo(() => [x], [x]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{}],
+};
+
+```
+
+
+## Error
+
+```
+  13 |   x.push(props);
+  14 |
+> 15 |   return useMemo(() => [x], [x]);
+     |          ^^^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (15:15)
+  16 | }
+  17 |
+  18 | export const FIXTURE_ENTRYPOINT = {
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.expect.md

@@ -5,8 +5,8 @@ import {useHook} from 'shared-runtime';
 
 // useMemo values may not be memoized in Forget output if we
 // infer that their deps always invalidate.
-// This is still correct as the useMemo in source was effectively
-// a no-op already.
+// This is technically a false positive as the useMemo in source
+// was effectively a no-op
 function useFoo(props) {
   const x = [];
   useHook();

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/useMemo-dropped-infer-always-invalidating.ts renamed to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.ts

@@ -0,0 +1,35 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees:true
+
+import {useRef, useMemo} from 'react';
+import {makeArray} from 'shared-runtime';
+
+function useFoo() {
+  const r = useRef();
+  return useMemo(() => makeArray(r), []);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [],
+};
+
+```
+
+
+## Error
+
+```
+   6 | function useFoo() {
+   7 |   const r = useRef();
+>  8 |   return useMemo(() => makeArray(r), []);
+     |          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (8:8)
+   9 | }
+  10 |
+  11 | export const FIXTURE_ENTRYPOINT = {
+```
+          
+