[compiler] Patch ValidatePreserveMemo to bailout correctly for refs

ghstack-source-id: 73fc47f23488b392644ffa27f124c309dc82b69c
Pull Request resolved: #30603

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PruneNonEscapingScopes.ts

@@ -9,6 +9,7 @@ import {CompilerError} from '../CompilerError';
 import {
   DeclarationId,
   Environment,
+  Identifier,
   InstructionId,
   Pattern,
   Place,
@@ -24,7 +25,7 @@ import {
   isMutableEffect,
 } from '../HIR';
 import {getFunctionCallSignature} from '../Inference/InferReferenceEffects';
-import {assertExhaustive} from '../Utils/utils';
+import {assertExhaustive, getOrInsertDefault} from '../Utils/utils';
 import {getPlaceScope} from './BuildReactiveBlocks';
 import {
   ReactiveFunctionTransform,
@@ -935,6 +936,11 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
   Set<DeclarationId>
 > {
   prunedScopes: Set<ScopeId> = new Set();
+  /**
+   * Track reassignments so we can correctly set `pruned` flags for
+   * inlined useMemos.
+   */
+  reassignments: Map<DeclarationId, Set<Identifier>> = new Map();
 
   override transformScope(
     scopeBlock: ReactiveScopeBlock,
@@ -977,24 +983,45 @@ class PruneScopesTransform extends ReactiveFunctionTransform<
     }
   }
 
+  /**
+   * If we pruned the scope for a non-escaping value, we know it doesn't
+   * need to be memoized. Remove associated `Memoize` instructions so that
+   * we don't report false positives on "missing" memoization of these values.
+   */
   override transformInstruction(
     instruction: ReactiveInstruction,
     state: Set<DeclarationId>,
   ): Transformed<ReactiveStatement> {
     this.traverseInstruction(instruction, state);
 
-    /**
-     * If we pruned the scope for a non-escaping value, we know it doesn't
-     * need to be memoized. Remove associated `Memoize` instructions so that
-     * we don't report false positives on "missing" memoization of these values.
-     */
-    if (instruction.value.kind === 'FinishMemoize') {
-      const identifier = instruction.value.decl.identifier;
+    const value = instruction.value;
+    if (value.kind === 'StoreLocal' && value.lvalue.kind === 'Reassign') {
+      const ids = getOrInsertDefault(
+        this.reassignments,
+        value.lvalue.place.identifier.declarationId,
+        new Set(),
+      );
+      ids.add(value.value.identifier);
+    } else if (value.kind === 'FinishMemoize') {
+      let decls;
+      if (value.decl.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.
+         */
+        decls = this.reassignments.get(value.decl.identifier.declarationId) ?? [
+          value.decl.identifier,
+        ];
+      } else {
+        decls = [value.decl.identifier];
+      }
+
       if (
-        identifier.scope !== null &&
-        this.prunedScopes.has(identifier.scope.id)
+        [...decls].every(
+          decl => decl.scope == null || this.prunedScopes.has(decl.scope.id),
+        )
       ) {
-        instruction.value.pruned = true;
+        value.pruned = true;
       }
     }
 

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;
 
@@ -434,6 +445,18 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
      */
     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 (value.kind === 'StartMemoize') {
       let depsFromSource: Array<ManualMemoDependency> | null = null;
       if (value.deps != null) {
@@ -451,6 +474,7 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
         decls: new Set(),
         depsFromSource,
         manualMemoId: value.manualMemoId,
+        reassignments: new Map(),
       };
 
       for (const {identifier, loc} of eachInstructionValueOperand(
@@ -483,20 +507,34 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
           suggestions: null,
         },
       );
+      const reassignments = state.manualMemoState.reassignments;
       state.manualMemoState = null;
       if (!value.pruned) {
         for (const {identifier, loc} of eachInstructionValueOperand(
           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 decls;
+          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.
+             */
+            decls = reassignments.get(identifier.declarationId) ?? [identifier];
+          } else {
+            decls = [identifier];
+          }
+
+          for (const identifier of decls) {
+            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/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.expect.md

@@ -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/useMemo-dropped-infer-always-invalidating.ts → compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/error.false-positive-useMemo-dropped-infer-always-invalidating.ts

@@ -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/error.maybe-mutable-ref-not-preserved.expect.md

@@ -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 = {
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.expect.md

@@ -0,0 +1,52 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {identity} from 'shared-runtime';
+
+function useFoo(cond) {
+  useMemo(() => {
+    if (cond) {
+      return 2;
+    } else {
+      return identity(5);
+    }
+  }, [cond]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [true],
+};
+
+```
+
+## Code
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import { useMemo } from "react";
+import { identity } from "shared-runtime";
+
+function useFoo(cond) {
+  let t0;
+  if (cond) {
+    t0 = 2;
+  } else {
+    t0 = identity(5);
+  }
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [true],
+};
+
+```
+      
+### Eval output
+(kind: ok) 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/prune-nonescaping-useMemo-mult-returns-primitive.ts

@@ -0,0 +1,19 @@
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {identity} from 'shared-runtime';
+
+function useFoo(cond) {
+  useMemo(() => {
+    if (cond) {
+      return 2;
+    } else {
+      return identity(5);
+    }
+  }, [cond]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [true],
+};