[compiler][newinference] Explicitly model StoreContext variables as mutable boxes

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

@@ -302,7 +302,10 @@ function runWithEnvironment(
       validateNoImpureFunctionsInRender(hir).unwrap();
     }
 
-    if (env.config.validateNoFreezingKnownMutableFunctions) {
+    if (
+      env.config.validateNoFreezingKnownMutableFunctions ||
+      env.config.enableNewMutationAliasingModel
+    ) {
       validateNoFreezingKnownMutableFunctions(hir).unwrap();
     }
   }

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

@@ -5,13 +5,14 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import invariant from 'invariant';
-import {HIRFunction, Identifier, MutableRange} from './HIR';
+import {HIRFunction, MutableRange, Place} from './HIR';
 import {
   eachInstructionLValue,
   eachInstructionOperand,
   eachTerminalOperand,
 } from './visitors';
+import {CompilerError} from '..';
+import {printPlace} from './PrintHIR';
 
 /*
  * Checks that all mutable ranges in the function are well-formed, with
@@ -20,38 +21,43 @@ import {
 export function assertValidMutableRanges(fn: HIRFunction): void {
   for (const [, block] of fn.body.blocks) {
     for (const phi of block.phis) {
-      visitIdentifier(phi.place.identifier);
-      for (const [, operand] of phi.operands) {
-        visitIdentifier(operand.identifier);
+      visit(phi.place, `phi for block bb${block.id}`);
+      for (const [pred, operand] of phi.operands) {
+        visit(operand, `phi predecessor bb${pred} for block bb${block.id}`);
       }
     }
     for (const instr of block.instructions) {
       for (const operand of eachInstructionLValue(instr)) {
-        visitIdentifier(operand.identifier);
+        visit(operand, `instruction [${instr.id}]`);
       }
       for (const operand of eachInstructionOperand(instr)) {
-        visitIdentifier(operand.identifier);
+        visit(operand, `instruction [${instr.id}]`);
       }
     }
     for (const operand of eachTerminalOperand(block.terminal)) {
-      visitIdentifier(operand.identifier);
+      visit(operand, `terminal [${block.terminal.id}]`);
     }
   }
 }
 
-function visitIdentifier(identifier: Identifier): void {
-  validateMutableRange(identifier.mutableRange);
-  if (identifier.scope !== null) {
-    validateMutableRange(identifier.scope.range);
+function visit(place: Place, description: string): void {
+  validateMutableRange(place, place.identifier.mutableRange, description);
+  if (place.identifier.scope !== null) {
+    validateMutableRange(place, place.identifier.scope.range, description);
   }
 }
 
-function validateMutableRange(mutableRange: MutableRange): void {
-  invariant(
-    (mutableRange.start === 0 && mutableRange.end === 0) ||
-      mutableRange.end > mutableRange.start,
-    'Identifier scope mutableRange was invalid: [%s:%s]',
-    mutableRange.start,
-    mutableRange.end,
+function validateMutableRange(
+  place: Place,
+  range: MutableRange,
+  description: string,
+): void {
+  CompilerError.invariant(
+    (range.start === 0 && range.end === 0) || range.end > range.start,
+    {
+      reason: `Invalid mutable range: [${range.start}:${range.end}]`,
+      description: `${printPlace(place)} in ${description}`,
+      loc: place.loc,
+    },
   );
 }

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

@@ -15,6 +15,7 @@ import {
 import {
   BasicBlock,
   BlockId,
+  DeclarationId,
   Environment,
   FunctionExpression,
   HIRFunction,
@@ -157,7 +158,13 @@ export function inferMutationAliasingEffects(
   }
   queue(fn.body.entry, initialState);
 
-  const context = new Context(isFunctionExpression, fn);
+  const hoistedContextDeclarations = findHoistedContextDeclarations(fn);
+
+  const context = new Context(
+    isFunctionExpression,
+    fn,
+    hoistedContextDeclarations,
+  );
 
   let count = 0;
   while (queuedStates.size !== 0) {
@@ -189,6 +196,25 @@ export function inferMutationAliasingEffects(
   return Ok(undefined);
 }
 
+function findHoistedContextDeclarations(fn: HIRFunction): Set<DeclarationId> {
+  const hoisted = new Set<DeclarationId>();
+  for (const block of fn.body.blocks.values()) {
+    for (const instr of block.instructions) {
+      if (instr.value.kind === 'DeclareContext') {
+        const kind = instr.value.lvalue.kind;
+        if (
+          kind == InstructionKind.HoistedConst ||
+          kind == InstructionKind.HoistedFunction ||
+          kind == InstructionKind.HoistedLet
+        ) {
+          hoisted.add(instr.value.lvalue.place.identifier.declarationId);
+        }
+      }
+    }
+  }
+  return hoisted;
+}
+
 class Context {
   internedEffects: Map<string, AliasingEffect> = new Map();
   instructionSignatureCache: Map<Instruction, InstructionSignature> = new Map();
@@ -197,10 +223,16 @@ class Context {
   catchHandlers: Map<BlockId, Place> = new Map();
   isFuctionExpression: boolean;
   fn: HIRFunction;
+  hoistedContextDeclarations: Set<DeclarationId>;
 
-  constructor(isFunctionExpression: boolean, fn: HIRFunction) {
+  constructor(
+    isFunctionExpression: boolean,
+    fn: HIRFunction,
+    hoistedContextDeclarations: Set<DeclarationId>,
+  ) {
     this.isFuctionExpression = isFunctionExpression;
     this.fn = fn;
+    this.hoistedContextDeclarations = hoistedContextDeclarations;
   }
 
   internEffect(effect: AliasingEffect): AliasingEffect {
@@ -241,7 +273,11 @@ function inferBlock(
   for (const instr of block.instructions) {
     let instructionSignature = context.instructionSignatureCache.get(instr);
     if (instructionSignature == null) {
-      instructionSignature = computeSignatureForInstruction(state.env, instr);
+      instructionSignature = computeSignatureForInstruction(
+        context,
+        state.env,
+        instr,
+      );
       context.instructionSignatureCache.set(instr, instructionSignature);
     }
     const effects = applySignature(context, state, instructionSignature, instr);
@@ -1265,6 +1301,7 @@ type InstructionSignature = {
  * an instruction.
  */
 function computeSignatureForInstruction(
+  context: Context,
   env: Environment,
   instr: Instruction,
 ): InstructionSignature {
@@ -1603,26 +1640,37 @@ function computeSignatureForInstruction(
       break;
     }
     case 'LoadContext': {
-      // We load from the context
-      effects.push({kind: 'Assign', from: value.place, into: lvalue});
+      /*
+       * Context variables are like mutable boxes. Loading from one
+       * is equivalent to a PropertyLoad from the box, so we model it
+       * with the same effect we use there (CreateFrom)
+       */
+      effects.push({kind: 'CreateFrom', from: value.place, into: lvalue});
       break;
     }
     case 'StoreContext': {
-      // We're storing into the conceptual box
-      if (value.lvalue.kind === InstructionKind.Reassign) {
+      /*
+       * Context variables are like mutable boxes, so semantically
+       * we're either creating (let/const) or mutating (reassign) a box,
+       * and then capturing the value into it.
+       */
+      if (
+        value.lvalue.kind === InstructionKind.Reassign ||
+        context.hoistedContextDeclarations.has(
+          value.lvalue.place.identifier.declarationId,
+        )
+      ) {
         effects.push({kind: 'Mutate', value: value.lvalue.place});
       } else {
-        // Context variables are conceptually like mutable boxes
         effects.push({
           kind: 'Create',
           into: value.lvalue.place,
           value: ValueKind.Mutable,
           reason: ValueReason.Other,
         });
       }
-      // Which aliases the value
       effects.push({
-        kind: 'Alias',
+        kind: 'Capture',
         from: value.value,
         into: value.lvalue.place,
       });