Update base for Update on "[compiler] Add wrapper functions to wrap change-detection storage and loading from the memo cache"

Summary: We may wish to perform some additional computation on values when they enter or exit the memo cache in change detection mode (e.g. make a deep copy, restore the original value). This builds support for doing so.

In addition, it drops the "ForDebugging" part of the flag name and makes it compatible with "disableMemoization": if memoization is disabled, we implement that by not restoring the old version of the value unless we're in a source-level memo block.

[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/HIR/HIR.ts

@@ -1168,7 +1168,7 @@ export type NonLocalBinding =
       imported: string;
     }
   // let, const, function, etc declared in the module but outside the current component/hook
-  | {kind: 'ModuleLocal'; name: string}
+  | {kind: 'ModuleLocal'; name: string, immutable: boolean}
   // an unresolved binding
   | {kind: 'Global'; name: string};
 

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

@@ -271,9 +271,14 @@ export default class HIRBuilder {
           module: importDeclaration.node.source.value,
         };
       } else {
+        const immutable =
+          (path.isVariableDeclaration() && path.node.kind === "const") ||
+          path.isClassDeclaration() ||
+          path.isClassExpression();
         return {
           kind: 'ModuleLocal',
           name: originalName,
+          immutable,
         };
       }
     }

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

@@ -6,7 +6,7 @@ import {
   makeInstructionId,
 } from '.';
 import {getPlaceScope} from '../ReactiveScopes/BuildReactiveBlocks';
-import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables';
+import {isMutableAtInstruction} from '../ReactiveScopes/InferReactiveScopeVariables';
 import DisjointSet from '../Utils/DisjointSet';
 import {getOrInsertDefault} from '../Utils/utils';
 import {
@@ -254,7 +254,7 @@ function visitPlace(
    * of the stack to the mutated outer scope.
    */
   const placeScope = getPlaceScope(id, place);
-  if (placeScope != null && isMutable({id} as any, place)) {
+  if (placeScope != null && isMutableAtInstruction({id} as any, place)) {
     const placeScopeIdx = activeScopes.indexOf(placeScope);
     if (placeScopeIdx !== -1 && placeScopeIdx !== activeScopes.length - 1) {
       joined.union([placeScope, ...activeScopes.slice(placeScopeIdx + 1)]);

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

@@ -26,7 +26,7 @@ import {
 } from '../HIR/visitors';
 import {
   findDisjointMutableValues,
-  isMutable,
+  isMutableAtInstruction,
 } from '../ReactiveScopes/InferReactiveScopeVariables';
 import DisjointSet from '../Utils/DisjointSet';
 import {assertExhaustive} from '../Utils/utils';
@@ -232,7 +232,7 @@ export function inferReactivePlaces(fn: HIRFunction): void {
               case Effect.Store:
               case Effect.ConditionallyMutate:
               case Effect.Mutate: {
-                if (isMutable(instruction, operand)) {
+                if (isMutableAtInstruction(instruction, operand)) {
                   const resolvedId = identifierMapping.get(operand.identifier);
                   if (resolvedId !== undefined) {
                     reactiveIdentifiers.markReactiveIdentifier(resolvedId);

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

@@ -687,10 +687,7 @@ function codegenReactiveScope(
   let computationBlock = codegenBlock(cx, block);
 
   let memoStatement;
-  if (
-    cx.env.config.enableChangeDetectionForDebugging != null &&
-    changeExpressions.length > 0
-  ) {
+  if (cx.env.config.enableChangeDetectionForDebugging != null) {
     const loc =
       typeof scope.loc === 'symbol'
         ? 'unknown location'

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

@@ -185,7 +185,10 @@ function mergeLocation(l: SourceLocation, r: SourceLocation): SourceLocation {
 }
 
 // Is the operand mutable at this given instruction
-export function isMutable({id}: Instruction, place: Place): boolean {
+export function isMutableAtInstruction(
+  {id}: Instruction,
+  place: Place
+): boolean {
   const range = place.identifier.mutableRange;
   return id >= range.start && id < range.end;
 }
@@ -253,6 +256,89 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
   }
 }
 
+/*
+ * These instructions may pick up external changes due to rules of react violations.
+ * Instructions should be included here if they may change without their inputs changing.
+ * For example, PostfixUpdate is not included because it only has a changed lval if
+ * it has a changed argument, but LoadProperty is included because the argument can be
+ * mutated elsewhere.
+ */
+function mayHaveChanged(env: Environment, instruction: Instruction): boolean {
+  if (env.config.enableChangeDetectionForDebugging == null) {
+    return false;
+  }
+  switch (instruction.value.kind) {
+    case "Await":
+    case "ComputedLoad":
+    case "Destructure":
+    case "GetIterator":
+    case "IteratorNext":
+    case "NextPropertyOf":
+    case "CallExpression":
+    case "MethodCall":
+    case "NewExpression": {
+      return true;
+    }
+    case "PropertyLoad": {
+      return instruction.value.property !== "current"
+    }
+    case "LoadGlobal": {
+      return (
+        instruction.value.binding.kind === "ModuleLocal" &&
+        !instruction.value.binding.immutable
+      );
+    }
+    case "PostfixUpdate":
+    case "PrefixUpdate":
+    case "DeclareLocal":
+    case "DeclareContext":
+    case "StoreLocal":
+    case "MetaProperty":
+    case "TypeCastExpression":
+    case "LoadLocal":
+    case "LoadContext":
+    case "StoreContext":
+    case "PropertyDelete":
+    case "ComputedDelete":
+    case "JSXText":
+    case "TemplateLiteral":
+    case "Primitive":
+    case "Debugger":
+    case "StartMemoize":
+    case "FinishMemoize":
+    case "UnaryExpression":
+    case "BinaryExpression":
+    case "StoreGlobal":
+    case "RegExpLiteral":
+    case "PropertyStore":
+    case "ComputedStore":
+    case "ArrayExpression":
+    case "JsxExpression":
+    case "JsxFragment":
+    case "ObjectExpression":
+    case "UnsupportedNode":
+    case "ObjectMethod":
+    case "FunctionExpression":
+    case "TaggedTemplateExpression": {
+      return false;
+    }
+    default: {
+      assertExhaustive(
+        instruction.value,
+        `Unexpected value kind \`${(instruction.value as any).kind}\``
+      );
+    }
+  }
+}
+
+function isIdentifierMutable(id: Identifier): boolean {
+  return id.mutableRange.end > id.mutableRange.start + 1;
+}
+
+function identifierHasMutableRange(id: Identifier): boolean {
+  return id.mutableRange.start > 0;
+}
+
 export function findDisjointMutableValues(
   fn: HIRFunction,
 ): DisjointSet<Identifier> {
@@ -265,7 +351,7 @@ export function findDisjointMutableValues(
     for (const phi of block.phis) {
       if (
         // The phi was reset because it was not mutated after creation
-        phi.id.mutableRange.start + 1 !== phi.id.mutableRange.end &&
+        isIdentifierMutable(phi.id) &&
         phi.id.mutableRange.end >
           (block.instructions.at(0)?.id ?? block.terminal.id)
       ) {
@@ -281,50 +367,52 @@ export function findDisjointMutableValues(
 
     for (const instr of block.instructions) {
       const operands: Array<Identifier> = [];
-      const range = instr.lvalue.identifier.mutableRange;
-      if (range.end > range.start + 1 || mayAllocate(fn.env, instr)) {
+      if (
+        isIdentifierMutable(instr.lvalue.identifier) ||
+        mayAllocate(fn.env, instr) ||
+        mayHaveChanged(fn.env, instr)
+      ) {
         operands.push(instr.lvalue!.identifier);
       }
       if (
         instr.value.kind === 'StoreLocal' ||
         instr.value.kind === 'StoreContext'
       ) {
-        if (
-          instr.value.lvalue.place.identifier.mutableRange.end >
-          instr.value.lvalue.place.identifier.mutableRange.start + 1
-        ) {
+        if (isIdentifierMutable(instr.value.lvalue.place.identifier)) {
           operands.push(instr.value.lvalue.place.identifier);
         }
         if (
-          isMutable(instr, instr.value.value) &&
-          instr.value.value.identifier.mutableRange.start > 0
+          isMutableAtInstruction(instr, instr.value.value) &&
+          identifierHasMutableRange(instr.value.value.identifier)
         ) {
           operands.push(instr.value.value.identifier);
         }
       } else if (instr.value.kind === 'Destructure') {
         for (const place of eachPatternOperand(instr.value.lvalue.pattern)) {
           if (
-            place.identifier.mutableRange.end >
-            place.identifier.mutableRange.start + 1
+            isIdentifierMutable(place.identifier) ||
+            mayHaveChanged(fn.env, instr)
           ) {
             operands.push(place.identifier);
           }
         }
         if (
-          isMutable(instr, instr.value.value) &&
-          instr.value.value.identifier.mutableRange.start > 0
+          (isMutableAtInstruction(instr, instr.value.value) &&
+            identifierHasMutableRange(instr.value.value.identifier)) ||
+          mayHaveChanged(fn.env, instr)
         ) {
           operands.push(instr.value.value.identifier);
         }
       } else if (instr.value.kind === 'MethodCall') {
         for (const operand of eachInstructionOperand(instr)) {
           if (
-            isMutable(instr, operand) &&
-            /*
-             * exclude global variables from being added to scopes, we can't recreate them!
-             * TODO: improve handling of module-scoped variables and globals
-             */
-            operand.identifier.mutableRange.start > 0
+            (isMutableAtInstruction(instr, operand) &&
+              /*
+               * exclude global variables from being added to scopes, we can't recreate them!
+               * TODO: improve handling of module-scoped variables and globals
+               */
+              identifierHasMutableRange(operand.identifier)) ||
+            mayHaveChanged(fn.env, instr)
           ) {
             operands.push(operand.identifier);
           }
@@ -337,12 +425,13 @@ export function findDisjointMutableValues(
       } else {
         for (const operand of eachInstructionOperand(instr)) {
           if (
-            isMutable(instr, operand) &&
-            /*
-             * exclude global variables from being added to scopes, we can't recreate them!
-             * TODO: improve handling of module-scoped variables and globals
-             */
-            operand.identifier.mutableRange.start > 0
+            (isMutableAtInstruction(instr, operand) &&
+              /*
+               * exclude global variables from being added to scopes, we can't recreate them!
+               * TODO: improve handling of module-scoped variables and globals
+               */
+              identifierHasMutableRange(operand.identifier)) ||
+            mayHaveChanged(fn.env, instr)
           ) {
             operands.push(operand.identifier);
           }

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/ReactiveScopes/PruneNonEscapingScopes.ts

@@ -811,7 +811,9 @@ class CollectDependenciesVisitor extends ReactiveFunctionVisitor<State> {
     this.env = env;
     this.options = {
       memoizeJsxElements: !this.env.config.enableForest,
-      forceMemoizePrimitives: this.env.config.enableForest,
+      forceMemoizePrimitives:
+        this.env.config.enableForest ||
+        this.env.config.enableChangeDetectionForDebugging != null,
     };
   }
 

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

@@ -17,7 +17,7 @@ import {
   isUseInsertionEffectHookType,
   isUseLayoutEffectHookType,
 } from '../HIR';
-import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables';
+import {isMutableAtInstruction} from '../ReactiveScopes/InferReactiveScopeVariables';
 import {
   ReactiveFunctionVisitor,
   visitReactiveFunction,
@@ -103,7 +103,7 @@ class Visitor extends ReactiveFunctionVisitor<CompilerError> {
          * TODO: isMutable is not safe to call here as it relies on identifier mutableRange which is no longer valid at this point
          * in the pipeline
          */
-        (isMutable(instruction as Instruction, deps) ||
+        (isMutableAtInstruction(instruction as Instruction, deps) ||
           isUnmemoized(deps.identifier, this.scopes))
       ) {
         state.push({

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/change-detect.expect.md

@@ -0,0 +1,139 @@
+
+## Input
+
+```javascript
+// @enableChangeDetectionForDebugging
+let glob = 1;
+
+function Component(props) {
+  const a = props.x;
+  const { b, ...c } = props.y;
+  const d = glob;
+  return (
+    <div>
+      {a}
+      {b}
+      {c.c}
+      {d}
+    </div>
+  );
+}
+
+```
+
+## Code
+
+```javascript
+import { $structuralCheck } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @enableChangeDetectionForDebugging
+let glob = 1;
+
+function Component(props) {
+  const $ = _c(12);
+  let t0;
+  {
+    t0 = props.x;
+    let condition = $[0] !== props.x;
+    if (!condition) {
+      let old$t0 = $[1];
+      $structuralCheck(old$t0, t0, "t0", "Component", "cached", "(5:5)");
+    }
+    $[0] = props.x;
+    $[1] = t0;
+    if (condition) {
+      t0 = props.x;
+      $structuralCheck($[1], t0, "t0", "Component", "recomputed", "(5:5)");
+      t0 = $[1];
+    }
+  }
+  const a = t0;
+  let c;
+  let b;
+  {
+    ({ b, ...c } = props.y);
+    let condition = $[2] !== props.y;
+    if (!condition) {
+      let old$c = $[3];
+      let old$b = $[4];
+      $structuralCheck(old$c, c, "c", "Component", "cached", "(6:6)");
+      $structuralCheck(old$b, b, "b", "Component", "cached", "(6:6)");
+    }
+    $[2] = props.y;
+    $[3] = c;
+    $[4] = b;
+    if (condition) {
+      ({ b, ...c } = props.y);
+      $structuralCheck($[3], c, "c", "Component", "recomputed", "(6:6)");
+      c = $[3];
+      $structuralCheck($[4], b, "b", "Component", "recomputed", "(6:6)");
+      b = $[4];
+    }
+  }
+  let t1;
+  {
+    t1 = c.c;
+    let condition = $[5] !== c.c;
+    if (!condition) {
+      let old$t1 = $[6];
+      $structuralCheck(old$t1, t1, "t1", "Component", "cached", "(12:12)");
+    }
+    $[5] = c.c;
+    $[6] = t1;
+    if (condition) {
+      t1 = c.c;
+      $structuralCheck($[6], t1, "t1", "Component", "recomputed", "(12:12)");
+      t1 = $[6];
+    }
+  }
+  let t2;
+  {
+    t2 = glob;
+    let condition = $[7] === Symbol.for("react.memo_cache_sentinel");
+    if (!condition) {
+      let old$t2 = $[7];
+      $structuralCheck(old$t2, t2, "t2", "Component", "cached", "(13:13)");
+    }
+    $[7] = t2;
+    if (condition) {
+      t2 = glob;
+      $structuralCheck($[7], t2, "t2", "Component", "recomputed", "(13:13)");
+      t2 = $[7];
+    }
+  }
+  let t3;
+  {
+    t3 = (
+      <div>
+        {a}
+        {b}
+        {t1}
+        {t2}
+      </div>
+    );
+    let condition = $[8] !== a || $[9] !== b || $[10] !== t1;
+    if (!condition) {
+      let old$t3 = $[11];
+      $structuralCheck(old$t3, t3, "t3", "Component", "cached", "(9:14)");
+    }
+    $[8] = a;
+    $[9] = b;
+    $[10] = t1;
+    $[11] = t3;
+    if (condition) {
+      t3 = (
+        <div>
+          {a}
+          {b}
+          {t1}
+          {t2}
+        </div>
+      );
+      $structuralCheck($[11], t3, "t3", "Component", "recomputed", "(9:14)");
+      t3 = $[11];
+    }
+  }
+  return t3;
+}
+
+```
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/change-detect.js

@@ -0,0 +1,16 @@
+// @enableChangeDetectionForDebugging
+let glob = 1;
+
+function Component(props) {
+  const a = props.x;
+  const { b, ...c } = props.y;
+  const d = glob;
+  return (
+    <div>
+      {a}
+      {b}
+      {c.c}
+      {d}
+    </div>
+  );
+}

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 = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useState-pruned-dependency-change-detect.expect.md

@@ -20,32 +20,25 @@ import { c as _c } from "react/compiler-runtime"; // @enableChangeDetectionForDe
 import { useState } from "react";
 
 function Component(props) {
-  const $ = _c(3);
+  const $ = _c(2);
+  const [x] = useState(f(props.x));
   let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = f(props.x);
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const [x] = useState(t0);
-  let t1;
   {
-    t1 = <div>{x}</div>;
-    let condition = $[1] !== x;
+    t0 = <div>{x}</div>;
+    let condition = $[0] !== x;
     if (!condition) {
-      let old$t1 = $[2];
-      $structuralCheck(old$t1, t1, "t1", "Component", "cached", "(6:6)");
+      let old$t0 = $[1];
+      $structuralCheck(old$t0, t0, "t0", "Component", "cached", "(6:6)");
     }
-    $[1] = x;
-    $[2] = t1;
+    $[0] = x;
+    $[1] = t0;
     if (condition) {
-      t1 = <div>{x}</div>;
-      $structuralCheck($[2], t1, "t1", "Component", "recomputed", "(6:6)");
-      t1 = $[2];
+      t0 = <div>{x}</div>;
+      $structuralCheck($[1], t0, "t0", "Component", "recomputed", "(6:6)");
+      t0 = $[1];
     }
   }
-  return t1;
+  return t0;
 }
 
 ```

compiler/packages/snap/src/SproutTodoFilter.ts

@@ -499,6 +499,7 @@ const skipFilter = new Set([
   'useState-unpruned-dependency',
   'useState-and-other-hook-unpruned-dependency',
   'change-detect-reassign',
+  "change-detect",
 
   // needs to be executed as a module
   'meta-property',