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

Context variables — in the {Declare,Store,Load}Context sense — are conceptually like mutable boxes that are stored into and read out of at various points. Previously some fixtures were failing in the new inference bc i wasn't fully modeling them this way, so this PR moves toward more explicitly modeling these variables exactly like a mutable object that we're storing into and reading out of.

One catch is that unlike with regular variables, a `StoreContext Let x = ...` may not be the initial declaration of a value — for hoisted bindings, they may be a `DeclareContext HoistedLet x` first. So we first do a scan over the HIR to determine which declaration ids have hoisted let/const/function bindings. Then we model the "first" declaration of each context declaration id as the creation of fresh mutable value (box), then model subsequent reassignments as mutations of that box plus capturing of a value into that box, and model loads as CreateFrom from the box. Thus StoreContext assignments have equivalent effects to PropertyStore and LoadContext has equivalent effects to PropertyLoad.

ghstack-source-id: 8b8a6e3
Pull Request resolved: #33465

Author: josephsavona

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,
       });