[compiler][wip] Fix for ValidatePreserveMemo edge case w multiple return in useMemo

Potential fix for #34750. With a useMemo with multiple returns we can sometimes end up thinking the FinishMemoize value wasn't memoized, even though it definitely is. But the same case works if you rewrite to have a let binding, assign to the let instead of returning, and then have a single return with the let — which is the exact shape we convert to when we inline the useMemo! The only difference is that when we inline, we end up with an extra LoadLocal. Adding an extra LoadLocal "fixes" the issue but this is definitely a hack.

However, along the way this helped me find that we also have a bug in PruneAlwaysInvalidatingScopes, where we were forgetting to set FinishMemoize instructions to pruned if their declaration always invalidates.

Putting up for feedback, not necessarily expecting to land as-is.

Author: josephsavona

compiler/packages/babel-plugin-react-compiler/src/HIR/HIRBuilder.ts

@@ -988,7 +988,7 @@ export function createTemporaryPlace(
     identifier: makeTemporaryIdentifier(env.nextIdentifierId, loc),
     reactive: false,
     effect: Effect.Unknown,
-    loc: GeneratedSource,
+    loc,
   };
 }
 

compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts

@@ -11,7 +11,6 @@ import {
   CallExpression,
   Effect,
   Environment,
-  FinishMemoize,
   FunctionExpression,
   HIRFunction,
   IdentifierId,
@@ -25,7 +24,6 @@ import {
   Place,
   PropertyLoad,
   SpreadPattern,
-  StartMemoize,
   TInstruction,
   getHookKindForType,
   makeInstructionId,
@@ -184,36 +182,52 @@ function makeManualMemoizationMarkers(
   depsList: Array<ManualMemoDependency> | null,
   memoDecl: Place,
   manualMemoId: number,
-): [TInstruction<StartMemoize>, TInstruction<FinishMemoize>] {
+): [Array<Instruction>, Array<Instruction>] {
+  const temp = createTemporaryPlace(env, memoDecl.loc);
   return [
-    {
-      id: makeInstructionId(0),
-      lvalue: createTemporaryPlace(env, fnExpr.loc),
-      value: {
-        kind: 'StartMemoize',
-        manualMemoId,
-        /*
-         * Use deps list from source instead of inferred deps
-         * as dependencies
-         */
-        deps: depsList,
+    [
+      {
+        id: makeInstructionId(0),
+        lvalue: createTemporaryPlace(env, fnExpr.loc),
+        value: {
+          kind: 'StartMemoize',
+          manualMemoId,
+          /*
+           * Use deps list from source instead of inferred deps
+           * as dependencies
+           */
+          deps: depsList,
+          loc: fnExpr.loc,
+        },
+        effects: null,
         loc: fnExpr.loc,
       },
-      effects: null,
-      loc: fnExpr.loc,
-    },
-    {
-      id: makeInstructionId(0),
-      lvalue: createTemporaryPlace(env, fnExpr.loc),
-      value: {
-        kind: 'FinishMemoize',
-        manualMemoId,
-        decl: {...memoDecl},
-        loc: fnExpr.loc,
+    ],
+    [
+      {
+        id: makeInstructionId(0),
+        lvalue: {...temp},
+        value: {
+          kind: 'LoadLocal',
+          place: {...memoDecl},
+          loc: memoDecl.loc,
+        },
+        effects: null,
+        loc: memoDecl.loc,
+      },
+      {
+        id: makeInstructionId(0),
+        lvalue: createTemporaryPlace(env, memoDecl.loc),
+        value: {
+          kind: 'FinishMemoize',
+          manualMemoId,
+          decl: {...temp},
+          loc: memoDecl.loc,
+        },
+        effects: null,
+        loc: memoDecl.loc,
       },
-      effects: null,
-      loc: fnExpr.loc,
-    },
+    ],
   ];
 }
 
@@ -409,10 +423,7 @@ export function dropManualMemoization(
    *   LoadLocal fnArg
    * - (if validation is enabled) collect manual memoization markers
    */
-  const queuedInserts: Map<
-    InstructionId,
-    TInstruction<StartMemoize> | TInstruction<FinishMemoize>
-  > = new Map();
+  const queuedInserts: Map<InstructionId, Array<Instruction>> = new Map();
   for (const [_, block] of func.body.blocks) {
     for (let i = 0; i < block.instructions.length; i++) {
       const instr = block.instructions[i]!;
@@ -557,11 +568,11 @@ export function dropManualMemoization(
       let nextInstructions: Array<Instruction> | null = null;
       for (let i = 0; i < block.instructions.length; i++) {
         const instr = block.instructions[i];
-        const insertInstr = queuedInserts.get(instr.id);
-        if (insertInstr != null) {
+        const insertInstructions = queuedInserts.get(instr.id);
+        if (insertInstructions != null) {
           nextInstructions = nextInstructions ?? block.instructions.slice(0, i);
           nextInstructions.push(instr);
-          nextInstructions.push(insertInstr);
+          nextInstructions.push(...insertInstructions);
         } else if (nextInstructions != null) {
           nextInstructions.push(instr);
         }

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

@@ -77,6 +77,15 @@ class Transform extends ReactiveFunctionTransform<boolean> {
         }
         break;
       }
+      case 'FinishMemoize': {
+        if (
+          !withinScope &&
+          this.alwaysInvalidatingValues.has(value.decl.identifier)
+        ) {
+          value.pruned = true;
+        }
+        break;
+      }
     }
     return {kind: 'keep'};
   }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.bug-invariant-expected-consistent-destructuring.expect.md

@@ -29,7 +29,7 @@ Found 1 error:
 
 Invariant: Expected consistent kind for destructuring
 
-Other places were `Reassign` but 'mutate? #t8$46[7:9]{reactive}' is const.
+Other places were `Reassign` but 'mutate? #t8$47[7:9]{reactive}' is const.
 
 error.bug-invariant-expected-consistent-destructuring.ts:9:9
    7 |

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,56 @@
+
+## Input
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {useHook} from 'shared-runtime';
+
+// If we can prove that a useMemo was ineffective because it would always invalidate,
+// then we shouldn't throw a "couldn't preserve existing memoization" error
+// TODO: consider reporting a separate error to the user for this case, if you're going
+// to memoize manually, then you probably want to know that it's a no-op
+function useFoo(props) {
+  const x = [];
+  useHook();
+  x.push(props);
+
+  return useMemo(() => [x], [x]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+// @validatePreserveExistingMemoizationGuarantees
+
+import { useMemo } from "react";
+import { useHook } from "shared-runtime";
+
+// If we can prove that a useMemo was ineffective because it would always invalidate,
+// then we shouldn't throw a "couldn't preserve existing memoization" error
+// TODO: consider reporting a separate error to the user for this case, if you're going
+// to memoize manually, then you probably want to know that it's a no-op
+function useFoo(props) {
+  const x = [];
+  useHook();
+  x.push(props);
+  return [x];
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) [[{}]]

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,21 @@
+// @validatePreserveExistingMemoizationGuarantees
+
+import {useMemo} from 'react';
+import {useHook} from 'shared-runtime';
+
+// If we can prove that a useMemo was ineffective because it would always invalidate,
+// then we shouldn't throw a "couldn't preserve existing memoization" error
+// TODO: consider reporting a separate error to the user for this case, if you're going
+// to memoize manually, then you probably want to know that it's a no-op
+function useFoo(props) {
+  const x = [];
+  useHook();
+  x.push(props);
+
+  return useMemo(() => [x], [x]);
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{}],
+};