Update on "[compiler] When preserving user-level useMemos, don't memoize arguments"

Summary:
This work enhances the don't-drop-memoization mode to not memoize the arguments to useMemo and useCallback (or technically, to memoize them in the same block as the hook call -- but since hooks can't be memoized, this is equivalent to never memoizing them). The advantage here is that we preserve the source-level code structure around memoization callbacks more, and avoid the compiler bug of having function expression dependencies be unconditional when their actual semantics are conditional.

This conceptually replaces some of the work of #30177: some of that was motivated by avoiding the deps bug, and the rest of the motivation was to better exercise the compiler's semantics, which was really happening as much as I'd hoped, as mofeiZ explained to me in that PR and offline!

[ghstack-poisoned]

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -155,11 +155,7 @@ function* runWithEnvironment(
   validateContextVariableLValues(hir);
   validateUseMemo(hir);
 
-  if (
-    !env.config.enablePreserveExistingManualUseMemo &&
-    !env.config.disableMemoizationForDebugging &&
-    !env.config.enableChangeDetectionForDebugging
-  ) {
+  if (!env.preserveManualMemo()) {
     dropManualMemoization(hir);
     yield log({kind: 'hir', name: 'DropManualMemoization', value: hir});
   }

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

@@ -765,6 +765,14 @@ export class Environment {
       return DefaultMutatingHook;
     }
   }
+
+  preserveManualMemo(): boolean {
+    return (
+      this.config.enablePreserveExistingManualUseMemo ||
+      this.config.disableMemoizationForDebugging ||
+      this.config.enableChangeDetectionForDebugging != null
+    );
+  }
 }
 
 // From https://github.com/facebook/react/blob/main/packages/eslint-plugin-react-hooks/src/RulesOfHooks.js#LL18C1-L23C2

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

@@ -45,6 +45,7 @@ export function memoizeFbtAndMacroOperandsInSameScope(
   const fbtMacroTags = new Set([
     ...FBT_TAGS,
     ...(fn.env.config.customMacros ?? []),
+    ...(fn.env.preserveManualMemo() ? ["useMemo", "useCallback"] : []),
   ]);
   const fbtValues: Set<IdentifierId> = new Set();
   while (true) {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved-nomemo.expect.md

@@ -25,33 +25,18 @@ import { c as _c } from "react/compiler-runtime"; // @disableMemoizationForDebug
 import { useMemo } from "react";
 
 function Component(t0) {
-  const $ = _c(5);
+  const $ = _c(2);
   const { a } = t0;
+  const x = useMemo(() => [a], []);
   let t1;
-  if ($[0] !== a || true) {
-    t1 = () => [a];
-    $[0] = a;
+  if ($[0] !== x || true) {
+    t1 = <div>{x}</div>;
+    $[0] = x;
     $[1] = t1;
   } else {
     t1 = $[1];
   }
-  let t2;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel") || true) {
-    t2 = [];
-    $[2] = t2;
-  } else {
-    t2 = $[2];
-  }
-  const x = useMemo(t1, t2);
-  let t3;
-  if ($[3] !== x || true) {
-    t3 = <div>{x}</div>;
-    $[3] = x;
-    $[4] = t3;
-  } else {
-    t3 = $[4];
-  }
-  return t3;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md

@@ -25,33 +25,18 @@ import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManu
 import { useMemo } from "react";
 
 function Component(t0) {
-  const $ = _c(5);
+  const $ = _c(2);
   const { a } = t0;
+  const x = useMemo(() => [a], []);
   let t1;
-  if ($[0] !== a) {
-    t1 = () => [a];
-    $[0] = a;
+  if ($[0] !== x) {
+    t1 = <div>{x}</div>;
+    $[0] = x;
     $[1] = t1;
   } else {
     t1 = $[1];
   }
-  let t2;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
-    t2 = [];
-    $[2] = t2;
-  } else {
-    t2 = $[2];
-  }
-  const x = useMemo(t1, t2);
-  let t3;
-  if ($[3] !== x) {
-    t3 = <div>{x}</div>;
-    $[3] = x;
-    $[4] = t3;
-  } else {
-    t3 = $[4];
-  }
-  return t3;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {