[compiler] New inference repros/fixes

Substantially improves the last major known issue with the new inference model's implementation: inferring effects of function expressions. I knowingly used a really simple (dumb) approach in InferFunctionExpressionAliasingEffects but it worked surprisingly well on a ton of code. However, investigating during the sync I saw that we the algorithm was literally running out of memory, or crashing from arrays that exceeded the maximum capacity. We were accumluating data flow in a way that could lead to lists of data flow captures compounding on themselves and growing very large very quickly. Plus, we were incorrectly recording some data flow, leading to cases where we reported false positive "can't mutate frozen value" for example.

So I went back to the drawing board. InferMutationAliasingRanges already builds up a data flow graph which it uses to figure out what values would be affected by mutations of other values, and update mutable ranges. Well, the key question that we really want to answer for inferring a function expression's aliasing effects is which values alias/capture where. Per the docs I wrote up, we only have to record such aliasing _if they are observable via mutations_. So, lightbulb: simulate mutations of the params, free variables, and return of the function expression and see which params/free-vars would be affected! That's what we do now, giving us precise information about which such values alias/capture where. When the "into" is a param/context-var we use Capture, iwhen the destination is the return we use Alias to be conservative.

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

@@ -1770,6 +1770,10 @@ export function isUseStateType(id: Identifier): boolean {
   return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInUseState';
 }
 
+export function isJsxType(type: Type): boolean {
+  return type.kind === 'Object' && type.shapeId === 'BuiltInJsx';
+}
+
 export function isRefOrRefValue(id: Identifier): boolean {
   return isUseRefType(id) || isRefValueType(id);
 }

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

@@ -22,7 +22,6 @@ import {inferMutableRanges} from './InferMutableRanges';
 import inferReferenceEffects from './InferReferenceEffects';
 import {assertExhaustive} from '../Utils/utils';
 import {inferMutationAliasingEffects} from './InferMutationAliasingEffects';
-import {inferFunctionExpressionAliasingEffectsSignature} from './InferFunctionExpressionAliasingEffectsSignature';
 import {inferMutationAliasingRanges} from './InferMutationAliasingRanges';
 
 export default function analyseFunctions(func: HIRFunction): void {
@@ -68,19 +67,12 @@ function lowerWithMutationAliasing(fn: HIRFunction): void {
   analyseFunctions(fn);
   inferMutationAliasingEffects(fn, {isFunctionExpression: true});
   deadCodeElimination(fn);
-  inferMutationAliasingRanges(fn, {isFunctionExpression: true});
+  const functionEffects = inferMutationAliasingRanges(fn, {
+    isFunctionExpression: true,
+  }).unwrap();
   rewriteInstructionKindsBasedOnReassignment(fn);
   inferReactiveScopeVariables(fn);
-  const effects = inferFunctionExpressionAliasingEffectsSignature(fn);
-  fn.env.logger?.debugLogIRs?.({
-    kind: 'hir',
-    name: 'AnalyseFunction (inner)',
-    value: fn,
-  });
-  if (effects != null) {
-    fn.aliasingEffects ??= [];
-    fn.aliasingEffects?.push(...effects);
-  }
+  fn.aliasingEffects = functionEffects;
 
   /**
    * Phase 2: populate the Effect of each context variable to use in inferring
@@ -88,7 +80,7 @@ function lowerWithMutationAliasing(fn: HIRFunction): void {
    * effects to decide if the function may be mutable or not.
    */
   const capturedOrMutated = new Set<IdentifierId>();
-  for (const effect of effects ?? []) {
+  for (const effect of functionEffects) {
     switch (effect.kind) {
       case 'Assign':
       case 'Alias':
@@ -140,6 +132,12 @@ function lowerWithMutationAliasing(fn: HIRFunction): void {
       operand.effect = Effect.Read;
     }
   }
+
+  fn.env.logger?.debugLogIRs?.({
+    kind: 'hir',
+    name: 'AnalyseFunction (inner)',
+    value: fn,
+  });
 }
 
 function lower(func: HIRFunction): void {

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

@@ -1,206 +0,0 @@
-/**
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- */
-
-import {HIRFunction, IdentifierId, Place, ValueKind, ValueReason} from '../HIR';
-import {getOrInsertDefault} from '../Utils/utils';
-import {AliasingEffect} from './AliasingEffects';
-
-/**
- * This function tracks data flow within an inner function expression in order to
- * compute a set of data-flow aliasing effects describing data flow between the function's
- * params, context variables, and return value.
- *
- * For example, consider the following function expression:
- *
- * ```
- * (x) => { return [x, y] }
- * ```
- *
- * This function captures both param `x` and context variable `y` into the return value.
- * Unlike our previous inference which counted this as a mutation of x and y, we want to
- * build a signature for the function that describes the data flow. We would infer
- * `Capture x -> return, Capture y -> return` effects for this function.
- *
- * This function *also* propagates more ambient-style effects (MutateFrozen, MutateGlobal, Impure, Render)
- * from instructions within the function up to the function itself.
- */
-export function inferFunctionExpressionAliasingEffectsSignature(
-  fn: HIRFunction,
-): Array<AliasingEffect> | null {
-  const effects: Array<AliasingEffect> = [];
-
-  /**
-   * Map used to identify tracked variables: params, context vars, return value
-   * This is used to detect mutation/capturing/aliasing of params/context vars
-   */
-  const tracked = new Map<IdentifierId, Place>();
-  tracked.set(fn.returns.identifier.id, fn.returns);
-  for (const operand of [...fn.context, ...fn.params]) {
-    const place = operand.kind === 'Identifier' ? operand : operand.place;
-    tracked.set(place.identifier.id, place);
-  }
-
-  /**
-   * Track capturing/aliasing of context vars and params into each other and into the return.
-   * We don't need to track locals and intermediate values, since we're only concerned with effects
-   * as they relate to arguments visible outside the function.
-   *
-   * For each aliased identifier we track capture/alias/createfrom and then merge this with how
-   * the value is used. Eg capturing an alias => capture. See joinEffects() helper.
-   */
-  type AliasedIdentifier = {
-    kind: AliasingKind;
-    place: Place;
-  };
-  const dataFlow = new Map<IdentifierId, Array<AliasedIdentifier>>();
-
-  /*
-   * Check for aliasing of tracked values. Also joins the effects of how the value is
-   * used (@param kind) with the aliasing type of each value
-   */
-  function lookup(
-    place: Place,
-    kind: AliasedIdentifier['kind'],
-  ): Array<AliasedIdentifier> | null {
-    if (tracked.has(place.identifier.id)) {
-      return [{kind, place}];
-    }
-    return (
-      dataFlow.get(place.identifier.id)?.map(aliased => ({
-        kind: joinEffects(aliased.kind, kind),
-        place: aliased.place,
-      })) ?? null
-    );
-  }
-
-  // todo: fixpoint
-  for (const block of fn.body.blocks.values()) {
-    for (const phi of block.phis) {
-      const operands: Array<AliasedIdentifier> = [];
-      for (const operand of phi.operands.values()) {
-        const inputs = lookup(operand, 'Alias');
-        if (inputs != null) {
-          operands.push(...inputs);
-        }
-      }
-      if (operands.length !== 0) {
-        dataFlow.set(phi.place.identifier.id, operands);
-      }
-    }
-    for (const instr of block.instructions) {
-      if (instr.effects == null) continue;
-      for (const effect of instr.effects) {
-        if (
-          effect.kind === 'Assign' ||
-          effect.kind === 'Capture' ||
-          effect.kind === 'Alias' ||
-          effect.kind === 'CreateFrom'
-        ) {
-          const from = lookup(effect.from, effect.kind);
-          if (from == null) {
-            continue;
-          }
-          const into = lookup(effect.into, 'Alias');
-          if (into == null) {
-            getOrInsertDefault(dataFlow, effect.into.identifier.id, []).push(
-              ...from,
-            );
-          } else {
-            for (const aliased of into) {
-              getOrInsertDefault(
-                dataFlow,
-                aliased.place.identifier.id,
-                [],
-              ).push(...from);
-            }
-          }
-        } else if (
-          effect.kind === 'Create' ||
-          effect.kind === 'CreateFunction'
-        ) {
-          getOrInsertDefault(dataFlow, effect.into.identifier.id, [
-            {kind: 'Alias', place: effect.into},
-          ]);
-        } else if (
-          effect.kind === 'MutateFrozen' ||
-          effect.kind === 'MutateGlobal' ||
-          effect.kind === 'Impure' ||
-          effect.kind === 'Render'
-        ) {
-          effects.push(effect);
-        }
-      }
-    }
-    if (block.terminal.kind === 'return') {
-      const from = lookup(block.terminal.value, 'Alias');
-      if (from != null) {
-        getOrInsertDefault(dataFlow, fn.returns.identifier.id, []).push(
-          ...from,
-        );
-      }
-    }
-  }
-
-  // Create aliasing effects based on observed data flow
-  let hasReturn = false;
-  for (const [into, from] of dataFlow) {
-    const input = tracked.get(into);
-    if (input == null) {
-      continue;
-    }
-    for (const aliased of from) {
-      if (
-        aliased.place.identifier.id === input.identifier.id ||
-        !tracked.has(aliased.place.identifier.id)
-      ) {
-        continue;
-      }
-      const effect = {kind: aliased.kind, from: aliased.place, into: input};
-      effects.push(effect);
-      if (
-        into === fn.returns.identifier.id &&
-        (aliased.kind === 'Assign' || aliased.kind === 'CreateFrom')
-      ) {
-        hasReturn = true;
-      }
-    }
-  }
-  // TODO: more precise return effect inference
-  if (!hasReturn) {
-    effects.unshift({
-      kind: 'Create',
-      into: fn.returns,
-      value:
-        fn.returnType.kind === 'Primitive'
-          ? ValueKind.Primitive
-          : ValueKind.Mutable,
-      reason: ValueReason.KnownReturnSignature,
-    });
-  }
-
-  return effects;
-}
-
-export enum MutationKind {
-  None = 0,
-  Conditional = 1,
-  Definite = 2,
-}
-
-type AliasingKind = 'Alias' | 'Capture' | 'CreateFrom' | 'Assign';
-function joinEffects(
-  effect1: AliasingKind,
-  effect2: AliasingKind,
-): AliasingKind {
-  if (effect1 === 'Capture' || effect2 === 'Capture') {
-    return 'Capture';
-  } else if (effect1 === 'Assign' || effect2 === 'Assign') {
-    return 'Assign';
-  } else {
-    return 'Alias';
-  }
-}

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

@@ -57,7 +57,6 @@ import {
 import {
   printAliasingEffect,
   printAliasingSignature,
-  printFunction,
   printIdentifier,
   printInstruction,
   printInstructionValue,
@@ -194,19 +193,15 @@ export function inferMutationAliasingEffects(
     hoistedContextDeclarations,
   );
 
-  let count = 0;
+  let iterationCount = 0;
   while (queuedStates.size !== 0) {
-    count++;
-    if (count > 100) {
-      console.log(
-        'oops infinite loop',
-        fn.id,
-        typeof fn.loc !== 'symbol' ? fn.loc?.filename : null,
-      );
-      if (DEBUG) {
-        console.log(printFunction(fn));
-      }
-      throw new Error('infinite loop');
+    iterationCount++;
+    if (iterationCount > 100) {
+      CompilerError.invariant(false, {
+        reason: `[InferMutationAliasingEffects] Potential infinite loop`,
+        description: `A value, temporary place, or effect was not cached properly`,
+        loc: fn.loc,
+      });
     }
     for (const [blockId, block] of fn.body.blocks) {
       const incomingState = queuedStates.get(blockId);
@@ -217,11 +212,6 @@ export function inferMutationAliasingEffects(
 
       statesByBlock.set(blockId, incomingState);
       const state = incomingState.clone();
-      if (DEBUG) {
-        console.log('*************');
-        console.log(`bb${block.id}`);
-        console.log('*************');
-      }
       inferBlock(context, state, block);
 
       for (const nextBlockId of eachTerminalSuccessor(block.terminal)) {
@@ -362,6 +352,11 @@ function inferBlock(
   } else if (terminal.kind === 'maybe-throw') {
     const handlerParam = context.catchHandlers.get(terminal.handler);
     if (handlerParam != null) {
+      CompilerError.invariant(state.kind(handlerParam) != null, {
+        reason:
+          'Expected catch binding to be intialized with a DeclareLocal Catch instruction',
+        loc: terminal.loc,
+      });
       const effects: Array<AliasingEffect> = [];
       for (const instr of block.instructions) {
         if (
@@ -476,14 +471,14 @@ function applySignature(
    * Track which values we've already aliased once, so that we can switch to
    * appendAlias() for subsequent aliases into the same value
    */
-  const aliased = new Set<IdentifierId>();
+  const initialized = new Set<IdentifierId>();
 
   if (DEBUG) {
     console.log(printInstruction(instruction));
   }
 
   for (const effect of signature.effects) {
-    applyEffect(context, state, effect, aliased, effects);
+    applyEffect(context, state, effect, initialized, effects);
   }
   if (DEBUG) {
     console.log(
@@ -508,7 +503,7 @@ function applyEffect(
   context: Context,
   state: InferenceState,
   _effect: AliasingEffect,
-  aliased: Set<IdentifierId>,
+  initialized: Set<IdentifierId>,
   effects: Array<AliasingEffect>,
 ): void {
   const effect = context.internEffect(_effect);
@@ -524,6 +519,13 @@ function applyEffect(
       break;
     }
     case 'Create': {
+      CompilerError.invariant(!initialized.has(effect.into.identifier.id), {
+        reason: `Cannot re-initialize variable within an instruction`,
+        description: `Re-initialized ${printPlace(effect.into)} in ${printAliasingEffect(effect)}`,
+        loc: effect.into.loc,
+      });
+      initialized.add(effect.into.identifier.id);
+
       let value = context.effectInstructionValueCache.get(effect);
       if (value == null) {
         value = {
@@ -538,6 +540,7 @@ function applyEffect(
         reason: new Set([effect.reason]),
       });
       state.define(effect.into, value);
+      effects.push(effect);
       break;
     }
     case 'ImmutableCapture': {
@@ -555,6 +558,13 @@ function applyEffect(
       break;
     }
     case 'CreateFrom': {
+      CompilerError.invariant(!initialized.has(effect.into.identifier.id), {
+        reason: `Cannot re-initialize variable within an instruction`,
+        description: `Re-initialized ${printPlace(effect.into)} in ${printAliasingEffect(effect)}`,
+        loc: effect.into.loc,
+      });
+      initialized.add(effect.into.identifier.id);
+
       const fromValue = state.kind(effect.from);
       let value = context.effectInstructionValueCache.get(effect);
       if (value == null) {
@@ -573,10 +583,21 @@ function applyEffect(
       switch (fromValue.kind) {
         case ValueKind.Primitive:
         case ValueKind.Global: {
-          // no need to track this data flow
+          effects.push({
+            kind: 'Create',
+            value: fromValue.kind,
+            into: effect.into,
+            reason: [...fromValue.reason][0] ?? ValueReason.Other,
+          });
           break;
         }
         case ValueKind.Frozen: {
+          effects.push({
+            kind: 'Create',
+            value: fromValue.kind,
+            into: effect.into,
+            reason: [...fromValue.reason][0] ?? ValueReason.Other,
+          });
           applyEffect(
             context,
             state,
@@ -585,7 +606,7 @@ function applyEffect(
               from: effect.from,
               into: effect.into,
             },
-            aliased,
+            initialized,
             effects,
           );
           break;
@@ -597,6 +618,13 @@ function applyEffect(
       break;
     }
     case 'CreateFunction': {
+      CompilerError.invariant(!initialized.has(effect.into.identifier.id), {
+        reason: `Cannot re-initialize variable within an instruction`,
+        description: `Re-initialized ${printPlace(effect.into)} in ${printAliasingEffect(effect)}`,
+        loc: effect.into.loc,
+      });
+      initialized.add(effect.into.identifier.id);
+
       effects.push(effect);
       /**
        * We consider the function mutable if it has any mutable context variables or
@@ -653,7 +681,7 @@ function applyEffect(
             from: capture,
             into: effect.into,
           },
-          aliased,
+          initialized,
           effects,
         );
       }
@@ -661,6 +689,14 @@ function applyEffect(
     }
     case 'Alias':
     case 'Capture': {
+      CompilerError.invariant(
+        effect.kind === 'Capture' || initialized.has(effect.into.identifier.id),
+        {
+          reason: `Expected destination value to already be initialized within this instruction for Alias effect`,
+          description: `Destination ${printPlace(effect.into)} is not initialized in this instruction`,
+          loc: effect.into.loc,
+        },
+      );
       /*
        * Capture describes potential information flow: storing a pointer to one value
        * within another. If the destination is not mutable, or the source value has
@@ -698,7 +734,7 @@ function applyEffect(
               from: effect.from,
               into: effect.into,
             },
-            aliased,
+            initialized,
             effects,
           );
           break;
@@ -714,6 +750,13 @@ function applyEffect(
       break;
     }
     case 'Assign': {
+      CompilerError.invariant(!initialized.has(effect.into.identifier.id), {
+        reason: `Cannot re-initialize variable within an instruction`,
+        description: `Re-initialized ${printPlace(effect.into)} in ${printAliasingEffect(effect)}`,
+        loc: effect.into.loc,
+      });
+      initialized.add(effect.into.identifier.id);
+
       /*
        * Alias represents potential pointer aliasing. If the type is a global,
        * a primitive (copy-on-write semantics) then we can prune the effect
@@ -730,7 +773,7 @@ function applyEffect(
               from: effect.from,
               into: effect.into,
             },
-            aliased,
+            initialized,
             effects,
           );
           let value = context.effectInstructionValueCache.get(effect);
@@ -768,12 +811,7 @@ function applyEffect(
           break;
         }
         default: {
-          if (aliased.has(effect.into.identifier.id)) {
-            state.appendAlias(effect.into, effect.from);
-          } else {
-            aliased.add(effect.into.identifier.id);
-            state.alias(effect.into, effect.from);
-          }
+          state.assign(effect.into, effect.from);
           effects.push(effect);
           break;
         }
@@ -784,7 +822,8 @@ function applyEffect(
       const functionValues = state.values(effect.function);
       if (
         functionValues.length === 1 &&
-        functionValues[0].kind === 'FunctionExpression'
+        functionValues[0].kind === 'FunctionExpression' &&
+        functionValues[0].loweredFunc.func.aliasingEffects != null
       ) {
         /*
          * We're calling a locally declared function, we already know it's effects!
@@ -819,18 +858,15 @@ function applyEffect(
             ),
         );
         if (signatureEffects != null) {
-          if (DEBUG) {
-            console.log('apply function expression effects');
-          }
           applyEffect(
             context,
             state,
             {kind: 'MutateTransitiveConditionally', value: effect.function},
-            aliased,
+            initialized,
             effects,
           );
           for (const signatureEffect of signatureEffects) {
-            applyEffect(context, state, signatureEffect, aliased, effects);
+            applyEffect(context, state, signatureEffect, initialized, effects);
           }
           break;
         }
@@ -854,16 +890,10 @@ function applyEffect(
         );
       }
       if (signatureEffects != null) {
-        if (DEBUG) {
-          console.log('apply aliasing signature effects');
-        }
         for (const signatureEffect of signatureEffects) {
-          applyEffect(context, state, signatureEffect, aliased, effects);
+          applyEffect(context, state, signatureEffect, initialized, effects);
         }
       } else if (effect.signature != null) {
-        if (DEBUG) {
-          console.log('apply legacy signature effects');
-        }
         const legacyEffects = computeEffectsForLegacySignature(
           state,
           effect.signature,
@@ -873,12 +903,9 @@ function applyEffect(
           effect.loc,
         );
         for (const legacyEffect of legacyEffects) {
-          applyEffect(context, state, legacyEffect, aliased, effects);
+          applyEffect(context, state, legacyEffect, initialized, effects);
         }
       } else {
-        if (DEBUG) {
-          console.log('default effects');
-        }
         applyEffect(
           context,
           state,
@@ -888,7 +915,7 @@ function applyEffect(
             value: ValueKind.Mutable,
             reason: ValueReason.Other,
           },
-          aliased,
+          initialized,
           effects,
         );
         /*
@@ -911,21 +938,21 @@ function applyEffect(
                 kind: 'MutateTransitiveConditionally',
                 value: operand,
               },
-              aliased,
+              initialized,
               effects,
             );
           }
           const mutateIterator =
             arg.kind === 'Spread' ? conditionallyMutateIterator(operand) : null;
           if (mutateIterator) {
-            applyEffect(context, state, mutateIterator, aliased, effects);
+            applyEffect(context, state, mutateIterator, initialized, effects);
           }
           applyEffect(
             context,
             state,
             // OK: recording information flow
             {kind: 'Alias', from: operand, into: effect.into},
-            aliased,
+            initialized,
             effects,
           );
           for (const otherArg of [
@@ -953,7 +980,7 @@ function applyEffect(
                 from: operand,
                 into: other,
               },
-              aliased,
+              initialized,
               effects,
             );
           }
@@ -1009,7 +1036,7 @@ function applyEffect(
                   suggestions: null,
                 },
               },
-              aliased,
+              initialized,
               effects,
             );
           }
@@ -1028,7 +1055,7 @@ function applyEffect(
                 suggestions: null,
               },
             },
-            aliased,
+            initialized,
             effects,
           );
         } else {
@@ -1059,7 +1086,7 @@ function applyEffect(
                 suggestions: null,
               },
             },
-            aliased,
+            initialized,
             effects,
           );
         }
@@ -1166,7 +1193,7 @@ class InferenceState {
   }
 
   // Updates the value at @param place to point to the same value as @param value.
-  alias(place: Place, value: Place): void {
+  assign(place: Place, value: Place): void {
     const values = this.#variables.get(value.identifier.id);
     CompilerError.invariant(values != null, {
       reason: `[InferMutationAliasingEffects] Expected value for identifier to be initialized`,
@@ -1244,9 +1271,6 @@ class InferenceState {
       kind: ValueKind.Frozen,
       reason: new Set([reason]),
     });
-    if (DEBUG) {
-      console.log(`freeze value: ${printInstructionValue(value)} ${reason}`);
-    }
     if (
       value.kind === 'FunctionExpression' &&
       (this.env.config.enablePreserveExistingMemoizationGuarantees ||
@@ -2103,8 +2127,6 @@ function computeEffectsForLegacySignature(
         const mutateIterator = conditionallyMutateIterator(place);
         if (mutateIterator != null) {
           effects.push(mutateIterator);
-          // TODO: should we always push to captures?
-          captures.push(place);
         }
         effects.push({
           kind: 'Capture',
@@ -2286,17 +2308,6 @@ function computeEffectsForSignature(
     // Too many args and there is no rest param to hold them
     (args.length > signature.params.length && signature.rest == null)
   ) {
-    if (DEBUG) {
-      if (signature.params.length > args.length) {
-        console.log(
-          `not enough args: ${args.length} args for ${signature.params.length} params`,
-        );
-      } else {
-        console.log(
-          `too many args: ${args.length} args for ${signature.params.length} params, with no rest param`,
-        );
-      }
-    }
     return null;
   }
   // Build substitutions
@@ -2311,9 +2322,6 @@ function computeEffectsForSignature(
       continue;
     } else if (params == null || i >= params.length || arg.kind === 'Spread') {
       if (signature.rest == null) {
-        if (DEBUG) {
-          console.log(`no rest value to hold param`);
-        }
         return null;
       }
       const place = arg.kind === 'Identifier' ? arg : arg.place;
@@ -2421,23 +2429,14 @@ function computeEffectsForSignature(
       case 'Apply': {
         const applyReceiver = substitutions.get(effect.receiver.identifier.id);
         if (applyReceiver == null || applyReceiver.length !== 1) {
-          if (DEBUG) {
-            console.log(`too many substitutions for receiver`);
-          }
           return null;
         }
         const applyFunction = substitutions.get(effect.function.identifier.id);
         if (applyFunction == null || applyFunction.length !== 1) {
-          if (DEBUG) {
-            console.log(`too many substitutions for function`);
-          }
           return null;
         }
         const applyInto = substitutions.get(effect.into.identifier.id);
         if (applyInto == null || applyInto.length !== 1) {
-          if (DEBUG) {
-            console.log(`too many substitutions for into`);
-          }
           return null;
         }
         const applyArgs: Array<Place | SpreadPattern | Hole> = [];
@@ -2447,18 +2446,12 @@ function computeEffectsForSignature(
           } else if (arg.kind === 'Identifier') {
             const applyArg = substitutions.get(arg.identifier.id);
             if (applyArg == null || applyArg.length !== 1) {
-              if (DEBUG) {
-                console.log(`too many substitutions for arg`);
-              }
               return null;
             }
             applyArgs.push(applyArg[0]);
           } else {
             const applyArg = substitutions.get(arg.place.identifier.id);
             if (applyArg == null || applyArg.length !== 1) {
-              if (DEBUG) {
-                console.log(`too many substitutions for arg`);
-              }
               return null;
             }
             applyArgs.push({kind: 'Spread', place: applyArg[0]});

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

@@ -5,7 +5,6 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import prettyFormat from 'pretty-format';
 import {CompilerError, SourceLocation} from '..';
 import {
   BlockId,
@@ -14,7 +13,10 @@ import {
   Identifier,
   IdentifierId,
   InstructionId,
+  isJsxType,
   makeInstructionId,
+  ValueKind,
+  ValueReason,
   Place,
 } from '../HIR/HIR';
 import {
@@ -23,43 +25,58 @@ import {
   eachTerminalOperand,
 } from '../HIR/visitors';
 import {assertExhaustive, getOrInsertWith} from '../Utils/utils';
-import {printFunction} from '../HIR';
-import {printIdentifier, printPlace} from '../HIR/PrintHIR';
-import {MutationKind} from './InferFunctionExpressionAliasingEffectsSignature';
-import {Result} from '../Utils/Result';
-
-const DEBUG = false;
-const VERBOSE = false;
+import {Err, Ok, Result} from '../Utils/Result';
+import {AliasingEffect} from './AliasingEffects';
 
 /**
- * Infers mutable ranges for all values in the program, using previously inferred
- * mutation/aliasing effects. This pass builds a data flow graph using the effects,
- * tracking an abstract notion of "when" each effect occurs relative to the others.
- * It then walks each mutation effect against the graph, updating the range of each
- * node that would be reachable at the "time" that the effect occurred.
+ * This pass builds an abstract model of the heap and interprets the effects of the
+ * given function in order to determine the following:
+ * - The mutable ranges of all identifiers in the function
+ * - The externally-visible effects of the function, such as mutations of params and
+ *   context-vars, aliasing between params/context-vars/return-value, and impure side
+ *   effects.
+ * - The legacy `Effect` to store on each Place.
+ *
+ * This pass builds a data flow graph using the effects, tracking an abstract notion
+ * of "when" each effect occurs relative to the others. It then walks each mutation
+ * effect against the graph, updating the range of each node that would be reachable
+ * at the "time" that the effect occurred.
  *
  * This pass also validates against invalid effects: any function that is reachable
  * by being called, or via a Render effect, is validated against mutating globals
  * or calling impure code.
  *
  * Note that this function also populates the outer function's aliasing effects with
- * any mutations that apply to its params or context variables. For example, a
- * function expression such as the following:
+ * any mutations that apply to its params or context variables.
+ *
+ * ## Example
+ * A function expression such as the following:
  *
  * ```
  * (x) => { x.y = true }
  * ```
  *
  * Would populate a `Mutate x` aliasing effect on the outer function.
+ *
+ * ## Returned Function Effects
+ *
+ * The function returns (if successful) a list of externally-visible effects.
+ * This is determined by simulating a conditional, transitive mutation against
+ * each param, context variable, and return value in turn, and seeing which other
+ * such values are affected. If they're affected, they must be captured, so we
+ * record a Capture.
+ *
+ * The only tricky bit is the return value, which could _alias_ (or even assign)
+ * one or more of the params/context-vars rather than just capturing. So we have
+ * to do a bit more tracking for returns.
  */
 export function inferMutationAliasingRanges(
   fn: HIRFunction,
   {isFunctionExpression}: {isFunctionExpression: boolean},
-): Result<void, CompilerError> {
-  if (VERBOSE) {
-    console.log();
-    console.log(printFunction(fn));
-  }
+): Result<Array<AliasingEffect>, CompilerError> {
+  // The set of externally-visible effects
+  const functionEffects: Array<AliasingEffect> = [];
+
   /**
    * Part 1: Infer mutable ranges for values. We build an abstract model of
    * values, the alias/capture edges between them, and the set of mutations.
@@ -115,20 +132,6 @@ export function inferMutationAliasingRanges(
     seenBlocks.add(block.id);
 
     for (const instr of block.instructions) {
-      if (
-        instr.value.kind === 'FunctionExpression' ||
-        instr.value.kind === 'ObjectMethod'
-      ) {
-        state.create(instr.lvalue, {
-          kind: 'Function',
-          function: instr.value.loweredFunc.func,
-        });
-      } else {
-        for (const lvalue of eachInstructionLValue(instr)) {
-          state.create(lvalue, {kind: 'Object'});
-        }
-      }
-
       if (instr.effects == null) continue;
       for (const effect of instr.effects) {
         if (effect.kind === 'Create') {
@@ -141,6 +144,15 @@ export function inferMutationAliasingRanges(
         } else if (effect.kind === 'CreateFrom') {
           state.createFrom(index++, effect.from, effect.into);
         } else if (effect.kind === 'Assign') {
+          /**
+           * TODO: Invariant that the node is not initialized yet
+           *
+           * InferFunctionExpressionAliasingEffectSignatures currently infers
+           * Assign effects in some places that should be Alias, leading to
+           * Assign effects that reinitialize a value. The end result appears to
+           * be fine, but we should fix that inference pass so that we add the
+           * invariant here.
+           */
           if (!state.nodes.has(effect.into.identifier)) {
             state.create(effect.into, {kind: 'Object'});
           }
@@ -183,8 +195,10 @@ export function inferMutationAliasingRanges(
           effect.kind === 'Impure'
         ) {
           errors.push(effect.error);
+          functionEffects.push(effect);
         } else if (effect.kind === 'Render') {
           renders.push({index: index++, place: effect.place});
+          functionEffects.push(effect);
         }
       }
     }
@@ -216,10 +230,6 @@ export function inferMutationAliasingRanges(
     }
   }
 
-  if (VERBOSE) {
-    console.log(state.debug());
-    console.log(pretty(mutations));
-  }
   for (const mutation of mutations) {
     state.mutate(
       mutation.index,
@@ -234,10 +244,6 @@ export function inferMutationAliasingRanges(
   for (const render of renders) {
     state.render(render.index, render.place.identifier, errors);
   }
-  if (DEBUG) {
-    console.log(pretty([...state.nodes.keys()]));
-  }
-  fn.aliasingEffects ??= [];
   for (const param of [...fn.context, ...fn.params]) {
     const place = param.kind === 'Identifier' ? param : param.place;
     const node = state.nodes.get(place.identifier);
@@ -248,13 +254,13 @@ export function inferMutationAliasingRanges(
     if (node.local != null) {
       if (node.local.kind === MutationKind.Conditional) {
         mutated = true;
-        fn.aliasingEffects.push({
+        functionEffects.push({
           kind: 'MutateConditionally',
           value: {...place, loc: node.local.loc},
         });
       } else if (node.local.kind === MutationKind.Definite) {
         mutated = true;
-        fn.aliasingEffects.push({
+        functionEffects.push({
           kind: 'Mutate',
           value: {...place, loc: node.local.loc},
         });
@@ -263,13 +269,13 @@ export function inferMutationAliasingRanges(
     if (node.transitive != null) {
       if (node.transitive.kind === MutationKind.Conditional) {
         mutated = true;
-        fn.aliasingEffects.push({
+        functionEffects.push({
           kind: 'MutateTransitiveConditionally',
           value: {...place, loc: node.transitive.loc},
         });
       } else if (node.transitive.kind === MutationKind.Definite) {
         mutated = true;
-        fn.aliasingEffects.push({
+        functionEffects.push({
           kind: 'MutateTransitive',
           value: {...place, loc: node.transitive.loc},
         });
@@ -458,10 +464,82 @@ export function inferMutationAliasingRanges(
     }
   }
 
-  if (VERBOSE) {
-    console.log(printFunction(fn));
+  /**
+   * Part 3
+   * Finish populating the externally visible effects. Above we bubble-up the side effects
+   * (MutateFrozen/MutableGlobal/Impure/Render) as well as mutations of context variables.
+   * Here we populate an effect to create the return value as well as populating alias/capture
+   * effects for how data flows between the params, context vars, and return.
+   */
+  functionEffects.push({
+    kind: 'Create',
+    into: fn.returns,
+    value:
+      fn.returnType.kind === 'Primitive'
+        ? ValueKind.Primitive
+        : isJsxType(fn.returnType)
+          ? ValueKind.Frozen
+          : ValueKind.Mutable,
+    reason: ValueReason.KnownReturnSignature,
+  });
+  /**
+   * Determine precise data-flow effects by simulating transitive mutations of the params/
+   * captures and seeing what other params/context variables are affected. Anything that
+   * would be transitively mutated needs a capture relationship.
+   */
+  const tracked: Array<Place> = [];
+  const ignoredErrors = new CompilerError();
+  for (const param of [...fn.params, ...fn.context, fn.returns]) {
+    const place = param.kind === 'Identifier' ? param : param.place;
+    tracked.push(place);
   }
-  return errors.asResult();
+  for (const into of tracked) {
+    const mutationIndex = index++;
+    state.mutate(
+      mutationIndex,
+      into.identifier,
+      null,
+      true,
+      MutationKind.Conditional,
+      into.loc,
+      ignoredErrors,
+    );
+    for (const from of tracked) {
+      if (
+        from.identifier.id === into.identifier.id ||
+        from.identifier.id === fn.returns.identifier.id
+      ) {
+        continue;
+      }
+      const fromNode = state.nodes.get(from.identifier);
+      CompilerError.invariant(fromNode != null, {
+        reason: `Expected a node to exist for all parameters and context variables`,
+        loc: into.loc,
+      });
+      if (fromNode.lastMutated === mutationIndex) {
+        if (into.identifier.id === fn.returns.identifier.id) {
+          // The return value could be any of the params/context variables
+          functionEffects.push({
+            kind: 'Alias',
+            from,
+            into,
+          });
+        } else {
+          // Otherwise params/context-vars can only capture each other
+          functionEffects.push({
+            kind: 'Capture',
+            from,
+            into,
+          });
+        }
+      }
+    }
+  }
+
+  if (errors.hasErrors() && !isFunctionExpression) {
+    return Err(errors);
+  }
+  return Ok(functionEffects);
 }
 
 function appendFunctionErrors(errors: CompilerError, fn: HIRFunction): void {
@@ -477,6 +555,12 @@ function appendFunctionErrors(errors: CompilerError, fn: HIRFunction): void {
   }
 }
 
+export enum MutationKind {
+  None = 0,
+  Conditional = 1,
+  Definite = 2,
+}
+
 type Node = {
   id: Identifier;
   createdFrom: Map<Identifier, number>;
@@ -485,6 +569,7 @@ type Node = {
   edges: Array<{index: number; node: Identifier; kind: 'capture' | 'alias'}>;
   transitive: {kind: MutationKind; loc: SourceLocation} | null;
   local: {kind: MutationKind; loc: SourceLocation} | null;
+  lastMutated: number;
   value:
     | {kind: 'Object'}
     | {kind: 'Phi'}
@@ -502,6 +587,7 @@ class AliasingState {
       edges: [],
       transitive: null,
       local: null,
+      lastMutated: 0,
       value,
     });
   }
@@ -511,11 +597,6 @@ class AliasingState {
     const fromNode = this.nodes.get(from.identifier);
     const toNode = this.nodes.get(into.identifier);
     if (fromNode == null || toNode == null) {
-      if (VERBOSE) {
-        console.log(
-          `skip: createFrom ${printPlace(from)}${!!fromNode} -> ${printPlace(into)}${!!toNode}`,
-        );
-      }
       return;
     }
     fromNode.edges.push({index, node: into.identifier, kind: 'alias'});
@@ -528,11 +609,6 @@ class AliasingState {
     const fromNode = this.nodes.get(from.identifier);
     const toNode = this.nodes.get(into.identifier);
     if (fromNode == null || toNode == null) {
-      if (VERBOSE) {
-        console.log(
-          `skip: capture ${printPlace(from)}${!!fromNode} -> ${printPlace(into)}${!!toNode}`,
-        );
-      }
       return;
     }
     fromNode.edges.push({index, node: into.identifier, kind: 'capture'});
@@ -545,11 +621,6 @@ class AliasingState {
     const fromNode = this.nodes.get(from.identifier);
     const toNode = this.nodes.get(into.identifier);
     if (fromNode == null || toNode == null) {
-      if (VERBOSE) {
-        console.log(
-          `skip: assign ${printPlace(from)}${!!fromNode} -> ${printPlace(into)}${!!toNode}`,
-        );
-      }
       return;
     }
     fromNode.edges.push({index, node: into.identifier, kind: 'alias'});
@@ -598,17 +669,13 @@ class AliasingState {
   mutate(
     index: number,
     start: Identifier,
-    end: InstructionId,
+    // Null is used for simulated mutations
+    end: InstructionId | null,
     transitive: boolean,
     kind: MutationKind,
     loc: SourceLocation,
     errors: CompilerError,
   ): void {
-    if (DEBUG) {
-      console.log(
-        `mutate ix=${index} start=$${start.id} end=[${end}]${transitive ? ' transitive' : ''} kind=${kind}`,
-      );
-    }
     const seen = new Set<Identifier>();
     const queue: Array<{
       place: Identifier;
@@ -623,21 +690,14 @@ class AliasingState {
       seen.add(current);
       const node = this.nodes.get(current);
       if (node == null) {
-        if (DEBUG) {
-          console.log(
-            `no node! ${printIdentifier(start)} for identifier ${printIdentifier(current)}`,
-          );
-        }
         continue;
       }
-      if (DEBUG) {
-        console.log(
-          `  mutate $${node.id.id} transitive=${transitive} direction=${direction}`,
+      node.lastMutated = Math.max(node.lastMutated, index);
+      if (end != null) {
+        node.id.mutableRange.end = makeInstructionId(
+          Math.max(node.id.mutableRange.end, end),
         );
       }
-      node.id.mutableRange.end = makeInstructionId(
-        Math.max(node.id.mutableRange.end, end),
-      );
       if (
         node.value.kind === 'Function' &&
         node.transitive == null &&
@@ -701,37 +761,5 @@ class AliasingState {
         }
       }
     }
-    if (DEBUG) {
-      const nodes = new Map();
-      for (const id of seen) {
-        const node = this.nodes.get(id);
-        nodes.set(id.id, node);
-      }
-      console.log(pretty(nodes));
-    }
-  }
-
-  debug(): string {
-    return pretty(this.nodes);
   }
 }
-
-export function pretty(v: any): string {
-  return prettyFormat(v, {
-    plugins: [
-      {
-        test: v =>
-          v !== null && typeof v === 'object' && v.kind === 'Identifier',
-        serialize: v => printPlace(v),
-      },
-      {
-        test: v =>
-          v !== null &&
-          typeof v === 'object' &&
-          typeof v.declarationId === 'number',
-        serialize: v =>
-          `${printIdentifier(v)}:${v.mutableRange.start}:${v.mutableRange.end}`,
-      },
-    ],
-  });
-}

compiler/packages/babel-plugin-react-compiler/src/Inference/MUTABILITY_ALIASING_MODEL.md

@@ -514,9 +514,9 @@ Intuition: these effects are inverses of each other (capturing into an object, e
 Capture then CreatFrom is equivalent to Alias: we have to assume that the result _is_ the original value and that a local mutation of the result could mutate the original.
 
 ```js
-const y = [x]; // capture
-const z = y[0]; // createfrom
-mutate(z); // this clearly can mutate x, so the result must be one of Assign/Alias/CreateFrom 
+const b = [a]; // capture
+const c = b[0]; // createfrom
+mutate(c); // this clearly can mutate a, so the result must be one of Assign/Alias/CreateFrom 
 ```
 
 We use Alias as the return type because the mutability kind of the result is not derived from the source value (there's a fresh object in between due to the capture), so the full set of effects in practice would be a Create+Alias.
@@ -528,17 +528,17 @@ CreateFrom c <- b
 Alias c <- a
 ```
 
-Meanwhile the opposite direction preservers the capture, because the result is not the same as the source:
+Meanwhile the opposite direction preserves the capture, because the result is not the same as the source:
 
 ```js
-const y = x[0]; // createfrom
-const z = [y]; // capture
-mutate(z); // does not mutate x, so the result must be Capture
+const b = a[0]; // createfrom
+const c = [b]; // capture
+mutate(c); // does not mutate a, so the result must be Capture
 ```
 
 ```
-Capture b <- a
-CreateFrom c <- b
+CreateFrom b <- a
+Capture c <- b
 =>
-Capture b <- a
+Capture c <- a
 ```

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/capture-in-function-expression-indirect.expect.md

@@ -0,0 +1,81 @@
+
+## Input
+
+```javascript
+import {Stringify, mutate} from 'shared-runtime';
+
+function Component({foo, bar}) {
+  let x = {foo};
+  let y = {bar};
+  const f0 = function () {
+    let a = {y};
+    let b = {x};
+    a.y.x = b;
+  };
+  f0();
+  mutate(y);
+  return <Stringify x={y} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{foo: 2, bar: 3}],
+  sequentialRenders: [
+    {foo: 2, bar: 3},
+    {foo: 2, bar: 3},
+    {foo: 2, bar: 4},
+    {foo: 3, bar: 4},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { Stringify, mutate } from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(3);
+  const { foo, bar } = t0;
+  let t1;
+  if ($[0] !== bar || $[1] !== foo) {
+    const x = { foo };
+    const y = { bar };
+    const f0 = function () {
+      const a = { y };
+      const b = { x };
+      a.y.x = b;
+    };
+
+    f0();
+    mutate(y);
+    t1 = <Stringify x={y} />;
+    $[0] = bar;
+    $[1] = foo;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ foo: 2, bar: 3 }],
+  sequentialRenders: [
+    { foo: 2, bar: 3 },
+    { foo: 2, bar: 3 },
+    { foo: 2, bar: 4 },
+    { foo: 3, bar: 4 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"x":{"bar":3,"x":{"x":{"foo":2}},"wat0":"joe"}}</div>
+<div>{"x":{"bar":3,"x":{"x":{"foo":2}},"wat0":"joe"}}</div>
+<div>{"x":{"bar":4,"x":{"x":{"foo":2}},"wat0":"joe"}}</div>
+<div>{"x":{"bar":4,"x":{"x":{"foo":3}},"wat0":"joe"}}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/capture-in-function-expression-indirect.js

@@ -0,0 +1,25 @@
+import {Stringify, mutate} from 'shared-runtime';
+
+function Component({foo, bar}) {
+  let x = {foo};
+  let y = {bar};
+  const f0 = function () {
+    let a = {y};
+    let b = {x};
+    a.y.x = b;
+  };
+  f0();
+  mutate(y);
+  return <Stringify x={y} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{foo: 2, bar: 3}],
+  sequentialRenders: [
+    {foo: 2, bar: 3},
+    {foo: 2, bar: 3},
+    {foo: 2, bar: 4},
+    {foo: 3, bar: 4},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/mutate-through-identity-function-expression.expect.md

@@ -0,0 +1,97 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {identity, ValidateMemoization} from 'shared-runtime';
+
+function Component({a, b}) {
+  const x = useMemo(() => ({a}), [a, b]);
+  const f = () => {
+    return identity(x);
+  };
+  const x2 = f();
+  x2.b = b;
+
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 0, b: 1},
+    {a: 1, b: 1},
+    {a: 0, b: 0},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+import { identity, ValidateMemoization } from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(10);
+  const { a, b } = t0;
+  let t1;
+  let x;
+  if ($[0] !== a || $[1] !== b) {
+    t1 = { a };
+    x = t1;
+    const f = () => identity(x);
+
+    const x2 = f();
+    x2.b = b;
+    $[0] = a;
+    $[1] = b;
+    $[2] = x;
+    $[3] = t1;
+  } else {
+    x = $[2];
+    t1 = $[3];
+  }
+  let t2;
+  if ($[4] !== a || $[5] !== b) {
+    t2 = [a, b];
+    $[4] = a;
+    $[5] = b;
+    $[6] = t2;
+  } else {
+    t2 = $[6];
+  }
+  let t3;
+  if ($[7] !== t2 || $[8] !== x) {
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
+    $[7] = t2;
+    $[8] = x;
+    $[9] = t3;
+  } else {
+    t3 = $[9];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 0, b: 0 }],
+  sequentialRenders: [
+    { a: 0, b: 0 },
+    { a: 0, b: 1 },
+    { a: 1, b: 1 },
+    { a: 0, b: 0 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[0,0],"output":{"a":0,"b":0}}</div>
+<div>{"inputs":[0,1],"output":{"a":0,"b":1}}</div>
+<div>{"inputs":[1,1],"output":{"a":1,"b":1}}</div>
+<div>{"inputs":[0,0],"output":{"a":0,"b":0}}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/mutate-through-identity-function-expression.js

@@ -0,0 +1,24 @@
+import {useMemo} from 'react';
+import {identity, ValidateMemoization} from 'shared-runtime';
+
+function Component({a, b}) {
+  const x = useMemo(() => ({a}), [a, b]);
+  const f = () => {
+    return identity(x);
+  };
+  const x2 = f();
+  x2.b = b;
+
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 0, b: 1},
+    {a: 1, b: 1},
+    {a: 0, b: 0},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/mutate-through-identity.expect.md

@@ -0,0 +1,92 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {identity, ValidateMemoization} from 'shared-runtime';
+
+function Component({a, b}) {
+  const x = useMemo(() => ({a}), [a, b]);
+  const x2 = identity(x);
+  x2.b = b;
+
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 0, b: 1},
+    {a: 1, b: 1},
+    {a: 0, b: 0},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+import { identity, ValidateMemoization } from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(10);
+  const { a, b } = t0;
+  let t1;
+  let x;
+  if ($[0] !== a || $[1] !== b) {
+    t1 = { a };
+    x = t1;
+    const x2 = identity(x);
+    x2.b = b;
+    $[0] = a;
+    $[1] = b;
+    $[2] = x;
+    $[3] = t1;
+  } else {
+    x = $[2];
+    t1 = $[3];
+  }
+  let t2;
+  if ($[4] !== a || $[5] !== b) {
+    t2 = [a, b];
+    $[4] = a;
+    $[5] = b;
+    $[6] = t2;
+  } else {
+    t2 = $[6];
+  }
+  let t3;
+  if ($[7] !== t2 || $[8] !== x) {
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
+    $[7] = t2;
+    $[8] = x;
+    $[9] = t3;
+  } else {
+    t3 = $[9];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 0, b: 0 }],
+  sequentialRenders: [
+    { a: 0, b: 0 },
+    { a: 0, b: 1 },
+    { a: 1, b: 1 },
+    { a: 0, b: 0 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[0,0],"output":{"a":0,"b":0}}</div>
+<div>{"inputs":[0,1],"output":{"a":0,"b":1}}</div>
+<div>{"inputs":[1,1],"output":{"a":1,"b":1}}</div>
+<div>{"inputs":[0,0],"output":{"a":0,"b":0}}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/mutate-through-identity.js

@@ -0,0 +1,21 @@
+import {useMemo} from 'react';
+import {identity, ValidateMemoization} from 'shared-runtime';
+
+function Component({a, b}) {
+  const x = useMemo(() => ({a}), [a, b]);
+  const x2 = identity(x);
+  x2.b = b;
+
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 0, b: 1},
+    {a: 1, b: 1},
+    {a: 0, b: 0},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-function-expression-effects-stack-overflow.expect.md

@@ -0,0 +1,60 @@
+
+## Input
+
+```javascript
+function Component() {
+  const x = {};
+  const fn = () => {
+    new Object()
+      .build(x)
+      .build({})
+      .build({})
+      .build({})
+      .build({})
+      .build({})
+      .build({});
+  };
+  return <Stringify x={x} fn={fn} />;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component() {
+  const $ = _c(2);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = {};
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const x = t0;
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    const fn = () => {
+      new Object()
+        .build(x)
+        .build({})
+        .build({})
+        .build({})
+        .build({})
+        .build({})
+        .build({});
+    };
+
+    t1 = <Stringify x={x} fn={fn} />;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-function-expression-effects-stack-overflow.js

@@ -0,0 +1,14 @@
+function Component() {
+  const x = {};
+  const fn = () => {
+    new Object()
+      .build(x)
+      .build({})
+      .build({})
+      .build({})
+      .build({})
+      .build({})
+      .build({});
+  };
+  return <Stringify x={x} fn={fn} />;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-invalid-function-expression-effects-phi.expect.md

@@ -0,0 +1,60 @@
+
+## Input
+
+```javascript
+function Component({a, b}) {
+  const y = {a};
+  const x = {b};
+  const f = () => {
+    let z = null;
+    while (z == null) {
+      z = x;
+    }
+    // z is a phi with a backedge, and we don't realize it could be x,
+    // and therefore fail to record a Capture x <- y effect for this
+    // function expression
+    z.y = y;
+  };
+  f();
+  mutate(x);
+  return <div>{x}</div>;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Component(t0) {
+  const $ = _c(3);
+  const { a, b } = t0;
+  let t1;
+  if ($[0] !== a || $[1] !== b) {
+    const y = { a };
+    const x = { b };
+    const f = () => {
+      let z = null;
+      while (z == null) {
+        z = x;
+      }
+
+      z.y = y;
+    };
+
+    f();
+    mutate(x);
+    t1 = <div>{x}</div>;
+    $[0] = a;
+    $[1] = b;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-invalid-function-expression-effects-phi.js

@@ -0,0 +1,17 @@
+function Component({a, b}) {
+  const y = {a};
+  const x = {b};
+  const f = () => {
+    let z = null;
+    while (z == null) {
+      z = x;
+    }
+    // z is a phi with a backedge, and we don't realize it could be x,
+    // and therefore fail to record a Capture x <- y effect for this
+    // function expression
+    z.y = y;
+  };
+  f();
+  mutate(x);
+  return <div>{x}</div>;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-mutate-new-set-of-frozen-items-in-callback.expect.md

@@ -0,0 +1,74 @@
+
+## Input
+
+```javascript
+// @enableNewMutationAliasingModel:true
+
+export const App = () => {
+  const [selected, setSelected] = useState(new Set<string>());
+  const onSelectedChange = (value: string) => {
+    const newSelected = new Set(selected);
+    if (newSelected.has(value)) {
+      // This should not count as a mutation of `selected`
+      newSelected.delete(value);
+    } else {
+      // This should not count as a mutation of `selected`
+      newSelected.add(value);
+    }
+    setSelected(newSelected);
+  };
+
+  return <Stringify selected={selected} onSelectedChange={onSelectedChange} />;
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel:true
+
+export const App = () => {
+  const $ = _c(6);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = new Set();
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const [selected, setSelected] = useState(t0);
+  let t1;
+  if ($[1] !== selected) {
+    t1 = (value) => {
+      const newSelected = new Set(selected);
+      if (newSelected.has(value)) {
+        newSelected.delete(value);
+      } else {
+        newSelected.add(value);
+      }
+
+      setSelected(newSelected);
+    };
+    $[1] = selected;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  const onSelectedChange = t1;
+  let t2;
+  if ($[3] !== onSelectedChange || $[4] !== selected) {
+    t2 = <Stringify selected={selected} onSelectedChange={onSelectedChange} />;
+    $[3] = onSelectedChange;
+    $[4] = selected;
+    $[5] = t2;
+  } else {
+    t2 = $[5];
+  }
+  return t2;
+};
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/repro-mutate-new-set-of-frozen-items-in-callback.js

@@ -0,0 +1,18 @@
+// @enableNewMutationAliasingModel:true
+
+export const App = () => {
+  const [selected, setSelected] = useState(new Set<string>());
+  const onSelectedChange = (value: string) => {
+    const newSelected = new Set(selected);
+    if (newSelected.has(value)) {
+      // This should not count as a mutation of `selected`
+      newSelected.delete(value);
+    } else {
+      // This should not count as a mutation of `selected`
+      newSelected.add(value);
+    }
+    setSelected(newSelected);
+  };
+
+  return <Stringify selected={selected} onSelectedChange={onSelectedChange} />;
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.expect.md

@@ -23,14 +23,9 @@ function Component({a, b, c}: {a: number; b: number; c: number}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, b, c]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a, b, c]} output={x} />;
       {/* TODO: should only depend on c */}
-      <ValidateMemoization
-        inputs={[a, b, c]}
-        output={x[0]}
-        alwaysCheck={true}
-      />
-      ;
+      <ValidateMemoization inputs={[a, b, c]} output={x[0]} />;
     </>
   );
 }
@@ -98,7 +93,7 @@ function Component(t0) {
   }
   let t3;
   if ($[9] !== t2 || $[10] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} alwaysCheck={true} />;
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
     $[9] = t2;
     $[10] = x;
     $[11] = t3;
@@ -117,7 +112,7 @@ function Component(t0) {
   }
   let t5;
   if ($[16] !== t4 || $[17] !== x[0]) {
-    t5 = <ValidateMemoization inputs={t4} output={x[0]} alwaysCheck={true} />;
+    t5 = <ValidateMemoization inputs={t4} output={x[0]} />;
     $[16] = t4;
     $[17] = x[0];
     $[18] = t5;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-control-flow-sensitive-mutation.tsx

@@ -19,14 +19,9 @@ function Component({a, b, c}: {a: number; b: number; c: number}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, b, c]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a, b, c]} output={x} />;
       {/* TODO: should only depend on c */}
-      <ValidateMemoization
-        inputs={[a, b, c]}
-        output={x[0]}
-        alwaysCheck={true}
-      />
-      ;
+      <ValidateMemoization inputs={[a, b, c]} output={x[0]} />;
     </>
   );
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.expect.md

@@ -22,7 +22,7 @@ function Component({a, b}) {
   typedMutate(z, b);
 
   // TODO: this *should* only depend on `a`
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -86,7 +86,7 @@ function Component(t0) {
   }
   let t3;
   if ($[7] !== t2 || $[8] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} alwaysCheck={true} />;
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
     $[7] = t2;
     $[8] = x;
     $[9] = t3;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/todo-transitivity-createfrom-capture-lambda.tsx

@@ -18,7 +18,7 @@ function Component({a, b}) {
   typedMutate(z, b);
 
   // TODO: this *should* only depend on `a`
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.expect.md

@@ -20,8 +20,8 @@ function Component({a, b}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a]} output={o} alwaysCheck={true} />;
-      <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a]} output={o} />;
+      <ValidateMemoization inputs={[a, b]} output={x} />;
     </>
   );
 }
@@ -92,7 +92,7 @@ function Component(t0) {
   }
   let t5;
   if ($[8] !== o || $[9] !== t4) {
-    t5 = <ValidateMemoization inputs={t4} output={o} alwaysCheck={true} />;
+    t5 = <ValidateMemoization inputs={t4} output={o} />;
     $[8] = o;
     $[9] = t4;
     $[10] = t5;
@@ -110,7 +110,7 @@ function Component(t0) {
   }
   let t7;
   if ($[14] !== t6 || $[15] !== x) {
-    t7 = <ValidateMemoization inputs={t6} output={x} alwaysCheck={true} />;
+    t7 = <ValidateMemoization inputs={t6} output={x} />;
     $[14] = t6;
     $[15] = x;
     $[16] = t7;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-add-captured-array-to-itself.tsx

@@ -16,8 +16,8 @@ function Component({a, b}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a]} output={o} alwaysCheck={true} />;
-      <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+      <ValidateMemoization inputs={[a]} output={o} />;
+      <ValidateMemoization inputs={[a, b]} output={x} />;
     </>
   );
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.expect.md

@@ -21,7 +21,7 @@ function Component({a, b}: {a: number; b: number}) {
   // mutates x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -85,7 +85,7 @@ function Component(t0) {
   }
   let t3;
   if ($[7] !== t2 || $[8] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} alwaysCheck={true} />;
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
     $[7] = t2;
     $[8] = x;
     $[9] = t3;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom-lambda.tsx

@@ -17,7 +17,7 @@ function Component({a, b}: {a: number; b: number}) {
   // mutates x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.expect.md

@@ -17,7 +17,7 @@ function Component({a, b}: {a: number; b: number}) {
   // mutates x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -76,7 +76,7 @@ function Component(t0) {
   }
   let t3;
   if ($[7] !== t2 || $[8] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} alwaysCheck={true} />;
+    t3 = <ValidateMemoization inputs={t2} output={x} />;
     $[7] = t2;
     $[8] = x;
     $[9] = t3;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-capture-createfrom.tsx

@@ -13,7 +13,7 @@ function Component({a, b}: {a: number; b: number}) {
   // mutates x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.expect.md

@@ -17,7 +17,7 @@ function Component({a, b}) {
   // does not mutate x, so x should not depend on b
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -73,7 +73,7 @@ function Component(t0) {
   }
   let t4;
   if ($[4] !== t3 || $[5] !== x) {
-    t4 = <ValidateMemoization inputs={t3} output={x} alwaysCheck={true} />;
+    t4 = <ValidateMemoization inputs={t3} output={x} />;
     $[4] = t3;
     $[5] = x;
     $[6] = t4;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-createfrom-capture.tsx

@@ -13,7 +13,7 @@ function Component({a, b}) {
   // does not mutate x, so x should not depend on b
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.expect.md

@@ -21,7 +21,7 @@ function Component({a, b}) {
   // could mutate x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -92,7 +92,7 @@ function Component(t0) {
   }
   let t4;
   if ($[9] !== t3 || $[10] !== x) {
-    t4 = <ValidateMemoization inputs={t3} output={x} alwaysCheck={true} />;
+    t4 = <ValidateMemoization inputs={t3} output={x} />;
     $[9] = t3;
     $[10] = x;
     $[11] = t4;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/transitivity-phi-assign-or-capture.tsx

@@ -17,7 +17,7 @@ function Component({a, b}) {
   // could mutate x
   typedMutate(z, b);
 
-  return <ValidateMemoization inputs={[a, b]} output={x} alwaysCheck={true} />;
+  return <ValidateMemoization inputs={[a, b]} output={x} />;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.expect.md

@@ -4,6 +4,7 @@
 ```javascript
 // @enableNewMutationAliasingModel
 
+import {useMemo} from 'react';
 import {
   identity,
   makeObject_Primitives,
@@ -14,7 +15,7 @@ import {
 
 function Component({a, b}) {
   // create a mutable value with input `a`
-  const x = makeObject_Primitives(a);
+  const x = useMemo(() => makeObject_Primitives(a), [a]);
 
   // freeze the value
   useIdentity(x);
@@ -49,6 +50,7 @@ export const FIXTURE_ENTRYPOINT = {
 ```javascript
 import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel
 
+import { useMemo } from "react";
 import {
   identity,
   makeObject_Primitives,
@@ -61,13 +63,15 @@ function Component(t0) {
   const $ = _c(7);
   const { a, b } = t0;
   let t1;
+  let t2;
   if ($[0] !== a) {
-    t1 = makeObject_Primitives(a);
+    t2 = makeObject_Primitives(a);
     $[0] = a;
-    $[1] = t1;
+    $[1] = t2;
   } else {
-    t1 = $[1];
+    t2 = $[1];
   }
+  t1 = t2;
   const x = t1;
 
   useIdentity(x);
@@ -75,24 +79,24 @@ function Component(t0) {
   const x2 = typedIdentity(x);
 
   identity(x2, b);
-  let t2;
-  if ($[2] !== a) {
-    t2 = [a];
-    $[2] = a;
-    $[3] = t2;
-  } else {
-    t2 = $[3];
-  }
   let t3;
-  if ($[4] !== t2 || $[5] !== x) {
-    t3 = <ValidateMemoization inputs={t2} output={x} />;
-    $[4] = t2;
-    $[5] = x;
-    $[6] = t3;
+  if ($[2] !== a) {
+    t3 = [a];
+    $[2] = a;
+    $[3] = t3;
   } else {
-    t3 = $[6];
+    t3 = $[3];
   }
-  return t3;
+  let t4;
+  if ($[4] !== t3 || $[5] !== x) {
+    t4 = <ValidateMemoization inputs={t3} output={x} />;
+    $[4] = t3;
+    $[5] = x;
+    $[6] = t4;
+  } else {
+    t4 = $[6];
+  }
+  return t4;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/typed-identity-function-frozen-input.js

@@ -1,5 +1,6 @@
 // @enableNewMutationAliasingModel
 
+import {useMemo} from 'react';
 import {
   identity,
   makeObject_Primitives,
@@ -10,7 +11,7 @@ import {
 
 function Component({a, b}) {
   // create a mutable value with input `a`
-  const x = makeObject_Primitives(a);
+  const x = useMemo(() => makeObject_Primitives(a), [a]);
 
   // freeze the value
   useIdentity(x);

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+import {useMemo} from 'react';
 import {ValidateMemoization} from 'shared-runtime';
 
 function Component({a, b, c}) {
@@ -13,9 +14,21 @@ function Component({a, b, c}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, c]} output={map} />
-      <ValidateMemoization inputs={[a, c]} output={mapAlias} />
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={map}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={mapAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }
@@ -44,6 +57,7 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
 import { ValidateMemoization } from "shared-runtime";
 
 function Component(t0) {
@@ -76,7 +90,9 @@ function Component(t0) {
   }
   let t2;
   if ($[7] !== map || $[8] !== t1) {
-    t2 = <ValidateMemoization inputs={t1} output={map} />;
+    t2 = (
+      <ValidateMemoization inputs={t1} output={map} onlyCheckCompiled={true} />
+    );
     $[7] = map;
     $[8] = t1;
     $[9] = t2;
@@ -94,7 +110,13 @@ function Component(t0) {
   }
   let t4;
   if ($[13] !== mapAlias || $[14] !== t3) {
-    t4 = <ValidateMemoization inputs={t3} output={mapAlias} />;
+    t4 = (
+      <ValidateMemoization
+        inputs={t3}
+        output={mapAlias}
+        onlyCheckCompiled={true}
+      />
+    );
     $[13] = mapAlias;
     $[14] = t3;
     $[15] = t4;
@@ -119,7 +141,9 @@ function Component(t0) {
   }
   let t7;
   if ($[20] !== t5 || $[21] !== t6) {
-    t7 = <ValidateMemoization inputs={t5} output={t6} />;
+    t7 = (
+      <ValidateMemoization inputs={t5} output={t6} onlyCheckCompiled={true} />
+    );
     $[20] = t5;
     $[21] = t6;
     $[22] = t7;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakmap-constructor.js

@@ -1,3 +1,4 @@
+import {useMemo} from 'react';
 import {ValidateMemoization} from 'shared-runtime';
 
 function Component({a, b, c}) {
@@ -9,9 +10,21 @@ function Component({a, b, c}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, c]} output={map} />
-      <ValidateMemoization inputs={[a, c]} output={mapAlias} />
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={map}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={mapAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.expect.md

@@ -2,6 +2,7 @@
 ## Input
 
 ```javascript
+import {useMemo} from 'react';
 import {ValidateMemoization} from 'shared-runtime';
 
 function Component({a, b, c}) {
@@ -13,9 +14,21 @@ function Component({a, b, c}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, c]} output={set} />
-      <ValidateMemoization inputs={[a, c]} output={setAlias} />
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={set}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={setAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }
@@ -44,6 +57,7 @@ export const FIXTURE_ENTRYPOINT = {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
 import { ValidateMemoization } from "shared-runtime";
 
 function Component(t0) {
@@ -76,7 +90,9 @@ function Component(t0) {
   }
   let t2;
   if ($[7] !== set || $[8] !== t1) {
-    t2 = <ValidateMemoization inputs={t1} output={set} />;
+    t2 = (
+      <ValidateMemoization inputs={t1} output={set} onlyCheckCompiled={true} />
+    );
     $[7] = set;
     $[8] = t1;
     $[9] = t2;
@@ -94,7 +110,13 @@ function Component(t0) {
   }
   let t4;
   if ($[13] !== setAlias || $[14] !== t3) {
-    t4 = <ValidateMemoization inputs={t3} output={setAlias} />;
+    t4 = (
+      <ValidateMemoization
+        inputs={t3}
+        output={setAlias}
+        onlyCheckCompiled={true}
+      />
+    );
     $[13] = setAlias;
     $[14] = t3;
     $[15] = t4;
@@ -119,7 +141,9 @@ function Component(t0) {
   }
   let t7;
   if ($[20] !== t5 || $[21] !== t6) {
-    t7 = <ValidateMemoization inputs={t5} output={t6} />;
+    t7 = (
+      <ValidateMemoization inputs={t5} output={t6} onlyCheckCompiled={true} />
+    );
     $[20] = t5;
     $[21] = t6;
     $[22] = t7;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/weakset-constructor.js

@@ -1,3 +1,4 @@
+import {useMemo} from 'react';
 import {ValidateMemoization} from 'shared-runtime';
 
 function Component({a, b, c}) {
@@ -9,9 +10,21 @@ function Component({a, b, c}) {
 
   return (
     <>
-      <ValidateMemoization inputs={[a, c]} output={set} />
-      <ValidateMemoization inputs={[a, c]} output={setAlias} />
-      <ValidateMemoization inputs={[b]} output={[hasB]} />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={set}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[a, c]}
+        output={setAlias}
+        onlyCheckCompiled={true}
+      />
+      <ValidateMemoization
+        inputs={[b]}
+        output={[hasB]}
+        onlyCheckCompiled={true}
+      />
     </>
   );
 }

compiler/packages/snap/src/sprout/shared-runtime.ts

@@ -269,12 +269,10 @@ export function ValidateMemoization({
   inputs,
   output: rawOutput,
   onlyCheckCompiled = false,
-  alwaysCheck = false,
 }: {
   inputs: Array<any>;
   output: any;
   onlyCheckCompiled?: boolean;
-  alwaysCheck?: boolean;
 }): React.ReactElement {
   'use no forget';
   // Wrap rawOutput as it might be a function, which useState would invoke.
@@ -282,7 +280,7 @@ export function ValidateMemoization({
   const [previousInputs, setPreviousInputs] = React.useState(inputs);
   const [previousOutput, setPreviousOutput] = React.useState(output);
   if (
-    alwaysCheck ||
+    !onlyCheckCompiled ||
     (onlyCheckCompiled &&
       (globalThis as any).__SNAP_EVALUATOR_MODE === 'forget')
   ) {