[compiler] Add wrapper functions to wrap change-detection storage and loading from the memo cache

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

@@ -411,7 +411,7 @@ function* runWithEnvironment(
     value: reactiveFunction,
   });
 
-  if (env.config.enableChangeDetectionForDebugging != null) {
+  if (env.config.enableChangeDetection != null) {
     pruneInitializationDependencies(reactiveFunction);
     yield log({
       kind: 'reactive',

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

@@ -538,11 +538,27 @@ export function compileProgram(
       externalFunctions.push(enableEmitHookGuards);
     }
 
-    if (pass.opts.environment?.enableChangeDetectionForDebugging != null) {
-      const enableChangeDetectionForDebugging = tryParseExternalFunction(
-        pass.opts.environment.enableChangeDetectionForDebugging,
-      );
-      externalFunctions.push(enableChangeDetectionForDebugging);
+    if (pass.opts.environment?.enableChangeDetection != null) {
+      const enableChangeDetection = tryParseExternalFunction({
+        importSpecifierName:
+          pass.opts.environment.enableChangeDetection.structuralCheck,
+        source: pass.opts.environment.enableChangeDetection.source,
+      });
+      externalFunctions.push(enableChangeDetection);
+      if (pass.opts.environment.enableChangeDetection.wrappers != null) {
+        const store = tryParseExternalFunction({
+          importSpecifierName:
+            pass.opts.environment.enableChangeDetection.wrappers.store,
+          source: pass.opts.environment.enableChangeDetection.source,
+        });
+        const restore = tryParseExternalFunction({
+          importSpecifierName:
+            pass.opts.environment.enableChangeDetection.wrappers.restore,
+          source: pass.opts.environment.enableChangeDetection.source,
+        });
+        externalFunctions.push(store);
+        externalFunctions.push(restore);
+      }
     }
   } catch (err) {
     handleError(err, pass, null);

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

@@ -393,7 +393,13 @@ const EnvironmentConfigSchema = z.object({
    * computed one. This detects cases where rules of react violations may cause the
    * compiled code to behave differently than the original.
    */
-  enableChangeDetectionForDebugging: ExternalFunctionSchema.nullish(),
+  enableChangeDetection: z
+    .object({
+      source: z.string(),
+      structuralCheck: z.string(),
+      wrappers: z.object({ store: z.string(), restore: z.string() }).nullish(),
+    })
+    .nullish(),
 
   /**
    * The react native re-animated library uses custom Babel transforms that
@@ -456,12 +462,27 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig {
     }
 
     if (
-      key === 'enableChangeDetectionForDebugging' &&
-      (val === undefined || val === 'true')
+      key === "enableChangeDetection" &&
+      (val === undefined || val === "true")
     ) {
       maybeConfig[key] = {
         source: 'react-compiler-runtime',
-        importSpecifierName: '$structuralCheck',
+        structuralCheck: "$structuralCheck",
+      };
+      continue;
+    }
+
+    if (
+      key === "enableChangeDetectionWrappers" &&
+      (val === undefined || val === "true")
+    ) {
+      maybeConfig["enableChangeDetection"] = {
+        source: "react-compiler-runtime",
+        structuralCheck: '$structuralCheck',
+        wrappers: {
+          store: "$store",
+          restore: "$restore",
+        },
       };
       continue;
     }
@@ -549,18 +570,6 @@ export class Environment {
     this.#shapes = new Map(DEFAULT_SHAPES);
     this.#globals = new Map(DEFAULT_GLOBALS);
 
-    if (
-      config.disableMemoizationForDebugging &&
-      config.enableChangeDetectionForDebugging != null
-    ) {
-      CompilerError.throwInvalidConfig({
-        reason: `Invalid environment config: the 'disableMemoizationForDebugging' and 'enableChangeDetectionForDebugging' options cannot be used together`,
-        description: null,
-        loc: null,
-        suggestions: null,
-      });
-    }
-
     for (const [hookName, hook] of this.config.customHooks) {
       CompilerError.invariant(!this.#globals.has(hookName), {
         reason: `[Globals] Found existing definition in global registry for custom hook ${hookName}`,
@@ -770,7 +779,7 @@ export class Environment {
     return (
       this.config.enablePreserveExistingManualUseMemo ||
       this.config.disableMemoizationForDebugging ||
-      this.config.enableChangeDetectionForDebugging != null
+      this.config.enableChangeDetection != null
     );
   }
 }

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

@@ -670,14 +670,6 @@ function codegenReactiveScope(
   }
 
   if (cx.env.config.disableMemoizationForDebugging) {
-    CompilerError.invariant(
-      cx.env.config.enableChangeDetectionForDebugging == null,
-      {
-        reason: `Expected to not have both change detection enabled and memoization disabled`,
-        description: `Incompatible config options`,
-        loc: null,
-      },
-    );
     testCondition = t.logicalExpression(
       '||',
       testCondition,
@@ -687,38 +679,76 @@ function codegenReactiveScope(
   let computationBlock = codegenBlock(cx, block);
 
   let memoStatement;
-  if (cx.env.config.enableChangeDetectionForDebugging != null) {
+  if (cx.env.config.enableChangeDetection != null) {
     const loc =
       typeof scope.loc === 'symbol'
         ? 'unknown location'
         : `(${scope.loc.start.line}:${scope.loc.end.line})`;
     const detectionFunction =
-      cx.env.config.enableChangeDetectionForDebugging.importSpecifierName;
+      cx.env.config.enableChangeDetection.structuralCheck;
     const cacheLoadOldValueStatements: Array<t.Statement> = [];
+    const restoreOldValueStatements: Array<t.Statement> = [];
     const changeDetectionStatements: Array<t.Statement> = [];
     const idempotenceDetectionStatements: Array<t.Statement> = [];
 
-    for (const {name, index, value} of cacheLoads) {
-      const loadName = cx.synthesizeName(`old$${name.name}`);
-      const slot = t.memberExpression(
-        t.identifier(cx.synthesizeName('$')),
+    for (const {
+      name: { name: nameStr },
+      index,
+      value,
+    } of cacheLoads) {
+      const baseSlot = t.memberExpression(
+        t.identifier(cx.synthesizeName("$")),
         t.numericLiteral(index),
         true,
       );
+
+      const genSlot = (): t.MemberExpression => t.cloneNode(baseSlot, true);
+
+      const loadNameStr = cx.synthesizeName(`old$${nameStr}`);
+
+      let storedValue, restoredValue, restoredRecomputed;
+      if (cx.env.config.enableChangeDetection.wrappers != null) {
+        storedValue = t.callExpression(
+          t.identifier(cx.env.config.enableChangeDetection.wrappers.store),
+          [value]
+        );
+        restoredValue = t.callExpression(
+          t.identifier(cx.env.config.enableChangeDetection.wrappers.restore),
+          [t.identifier(loadNameStr)]
+        );
+        restoredRecomputed = t.callExpression(
+          t.identifier(cx.env.config.enableChangeDetection.wrappers.restore),
+          [genSlot()]
+        );
+      } else {
+        storedValue = value;
+        restoredValue = t.identifier(loadNameStr);
+        restoredRecomputed = genSlot();
+      }
+
       cacheStoreStatements.push(
-        t.expressionStatement(t.assignmentExpression('=', slot, value)),
+        t.expressionStatement(
+          t.assignmentExpression("=", genSlot(), storedValue)
+        )
       );
       cacheLoadOldValueStatements.push(
         t.variableDeclaration('let', [
-          t.variableDeclarator(t.identifier(loadName), slot),
+          t.variableDeclarator(t.identifier(loadNameStr), genSlot()),
         ]),
       );
+      if (!cx.env.config.disableMemoizationForDebugging) {
+        restoreOldValueStatements.push(
+          t.expressionStatement(
+            t.assignmentExpression("=", t.identifier(nameStr), restoredValue)
+          )
+        );
+      }
       changeDetectionStatements.push(
         t.expressionStatement(
           t.callExpression(t.identifier(detectionFunction), [
-            t.identifier(loadName),
-            t.cloneNode(name, true),
-            t.stringLiteral(name.name),
+            t.identifier(loadNameStr),
+            t.identifier(nameStr),
+            t.stringLiteral(nameStr),
             t.stringLiteral(cx.fnName),
             t.stringLiteral('cached'),
             t.stringLiteral(loc),
@@ -728,17 +758,19 @@ function codegenReactiveScope(
       idempotenceDetectionStatements.push(
         t.expressionStatement(
           t.callExpression(t.identifier(detectionFunction), [
-            t.cloneNode(slot, true),
-            t.cloneNode(name, true),
-            t.stringLiteral(name.name),
+            genSlot(),
+            t.identifier(nameStr),
+            t.stringLiteral(nameStr),
             t.stringLiteral(cx.fnName),
             t.stringLiteral('recomputed'),
             t.stringLiteral(loc),
           ]),
         ),
       );
       idempotenceDetectionStatements.push(
-        t.expressionStatement(t.assignmentExpression('=', name, slot)),
+        t.expressionStatement(
+          t.assignmentExpression("=", t.identifier(nameStr), restoredRecomputed)
+        )
       );
     }
     const condition = cx.synthesizeName('condition');

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

@@ -264,7 +264,7 @@ function mayAllocate(env: Environment, instruction: Instruction): boolean {
  * mutated elsewhere.
  */
 function mayHaveChanged(env: Environment, instruction: Instruction): boolean {
-  if (env.config.enableChangeDetectionForDebugging == null) {
+  if (env.config.enableChangeDetection == null) {
     return false;
   }
   switch (instruction.value.kind) {

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

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

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

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableChangeDetectionForDebugging
+// @enableChangeDetection
 function Component(props) {
   let x = null;
   if (props.cond) {
@@ -18,7 +18,7 @@ function Component(props) {
 
 ```javascript
 import { $structuralCheck } from "react-compiler-runtime";
-import { c as _c } from "react/compiler-runtime"; // @enableChangeDetectionForDebugging
+import { c as _c } from "react/compiler-runtime"; // @enableChangeDetection
 function Component(props) {
   const $ = _c(2);
   let x = null;

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

@@ -1,4 +1,4 @@
-// @enableChangeDetectionForDebugging
+// @enableChangeDetection
 function Component(props) {
   let x = null;
   if (props.cond) {

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

@@ -0,0 +1,48 @@
+
+## Input
+
+```javascript
+// @enableChangeDetectionWrappers
+function Component(props) {
+  let x = null;
+  if (props.cond) {
+    x = [];
+    x.push(props.value);
+  }
+  return x;
+}
+
+```
+
+## Code
+
+```javascript
+import { $structuralCheck, $store, $restore } from "react-compiler-runtime";
+import { c as _c } from "react/compiler-runtime"; // @enableChangeDetectionWrappers
+function Component(props) {
+  const $ = _c(2);
+  let x = null;
+  if (props.cond) {
+    {
+      x = [];
+      x.push(props.value);
+      let condition = $[0] !== props.value;
+      if (!condition) {
+        let old$x = $[1];
+        $structuralCheck(old$x, x, "x", "Component", "cached", "(3:6)");
+      }
+      $[0] = props.value;
+      $[1] = $store(x);
+      if (condition) {
+        x = [];
+        x.push(props.value);
+        $structuralCheck($[1], x, "x", "Component", "recomputed", "(3:6)");
+        x = $[1];
+      }
+    }
+  }
+  return x;
+}
+
+```
+

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

@@ -0,0 +1,9 @@
+// @enableChangeDetectionWrappers
+function Component(props) {
+  let x = null;
+  if (props.cond) {
+    x = [];
+    x.push(props.value);
+  }
+  return x;
+}

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

@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enableChangeDetectionForDebugging
+// @enableChangeDetection
 let glob = 1;
 
 function Component(props) {
@@ -25,7 +25,7 @@ function Component(props) {
 
 ```javascript
 import { $structuralCheck } from "react-compiler-runtime";
-import { c as _c } from "react/compiler-runtime"; // @enableChangeDetectionForDebugging
+import { c as _c } from "react/compiler-runtime"; // @enableChangeDetection
 let glob = 1;
 
 function Component(props) {

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

@@ -1,4 +1,4 @@
-// @enableChangeDetectionForDebugging
+// @enableChangeDetection
 let glob = 1;
 
 function Component(props) {