[compiler] Alternate pipeline for new mutability model

This PR gets a first fixture working end-to-end with the new mutability and aliasing model. Key changes:

* Add a feature flag to enable the model. When enabled we no longer call InferReferenceEffects or InferMutableRanges, and instead use the new equivalents.
* Adds a pass that infers Place-specific effects based on mutable ranges and instruction effects. This is necessary to satisfy existing code that requires operand effects to be populated.
* Adds a pass that infers the outwardly-visible capturing/aliasing behavior of a function expression. The idea is that this can bubble up and be used in conjunction with the `Apply` effect to get precise inference of things like `array.map(() => { ... })`.

ghstack-source-id: 03427439ec
Pull Request resolved: https://github.com/facebook/react/pull/33364

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

@@ -228,25 +228,27 @@ function runWithEnvironment(
   analyseFunctions(hir);
   log({kind: 'hir', name: 'AnalyseFunctions', value: hir});
 
-  const mutabilityAliasingErrors = inferMutationAliasingEffects(hir);
-  log({kind: 'hir', name: 'InferMutationAliasingEffects', value: hir});
-  if (env.isInferredMemoEnabled) {
-    if (mutabilityAliasingErrors.isErr()) {
-      throw mutabilityAliasingErrors.unwrapErr();
+  if (!env.config.enableNewMutationAliasingModel) {
+    const fnEffectErrors = inferReferenceEffects(hir);
+    if (env.isInferredMemoEnabled) {
+      if (fnEffectErrors.length > 0) {
+        CompilerError.throw(fnEffectErrors[0]);
+      }
+    }
+    log({kind: 'hir', name: 'InferReferenceEffects', value: hir});
+  } else {
+    const mutabilityAliasingErrors = inferMutationAliasingEffects(hir);
+    log({kind: 'hir', name: 'InferMutationAliasingEffects', value: hir});
+    if (env.isInferredMemoEnabled) {
+      if (mutabilityAliasingErrors.isErr()) {
+        throw mutabilityAliasingErrors.unwrapErr();
+      }
     }
   }
-  inferMutationAliasingRanges(hir);
-  log({kind: 'hir', name: 'InferMutationAliasingRanges', value: hir});
 
-  const fnEffectErrors = inferReferenceEffects(hir);
-  if (env.isInferredMemoEnabled) {
-    if (fnEffectErrors.length > 0) {
-      CompilerError.throw(fnEffectErrors[0]);
-    }
+  if (!env.config.enableNewMutationAliasingModel) {
+    validateLocalsNotReassignedAfterRender(hir);
   }
-  log({kind: 'hir', name: 'InferReferenceEffects', value: hir});
-
-  validateLocalsNotReassignedAfterRender(hir);
 
   // Note: Has to come after infer reference effects because "dead" code may still affect inference
   deadCodeElimination(hir);
@@ -260,8 +262,13 @@ function runWithEnvironment(
   pruneMaybeThrows(hir);
   log({kind: 'hir', name: 'PruneMaybeThrows', value: hir});
 
-  inferMutableRanges(hir);
-  log({kind: 'hir', name: 'InferMutableRanges', value: hir});
+  if (!env.config.enableNewMutationAliasingModel) {
+    inferMutableRanges(hir);
+    log({kind: 'hir', name: 'InferMutableRanges', value: hir});
+  } else {
+    inferMutationAliasingRanges(hir);
+    log({kind: 'hir', name: 'InferMutationAliasingRanges', value: hir});
+  }
 
   if (env.isInferredMemoEnabled) {
     if (env.config.assertValidMutableRanges) {

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

@@ -226,6 +226,7 @@ export function lower(
     loc: func.node.loc ?? GeneratedSource,
     env,
     effects: null,
+    aliasingEffects: null,
     directives,
   });
 }

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

@@ -243,6 +243,11 @@ export const EnvironmentConfigSchema = z.object({
    */
   enableUseTypeAnnotations: z.boolean().default(false),
 
+  /**
+   * Enable a new model for mutability and aliasing inference
+   */
+  enableNewMutationAliasingModel: z.boolean().default(false),
+
   /**
    * Enables inference of optional dependency chains. Without this flag
    * a property chain such as `props?.items?.foo` will infer as a dep on

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

@@ -287,6 +287,7 @@ export type HIRFunction = {
   generator: boolean;
   async: boolean;
   directives: Array<string>;
+  aliasingEffects?: Array<AliasingEffect> | null;
 };
 
 export type FunctionEffect =

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

@@ -75,9 +75,9 @@ export function printFunction(fn: HIRFunction): string {
   if (definition.length !== 0) {
     output.push(definition);
   }
-  output.push(printType(fn.returnType));
-  output.push(printHIR(fn.body));
+  output.push(`: ${printType(fn.returnType)} @ ${printPlace(fn.returns)}`);
   output.push(...fn.directives);
+  output.push(printHIR(fn.body));
   return output.join('\n');
 }
 

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

@@ -26,6 +26,12 @@ import {
   eachInstructionValueOperand,
 } from '../HIR/visitors';
 import {Iterable_some} from '../Utils/utils';
+import {
+  AliasingEffect,
+  inferMutationAliasingEffects,
+} from './InferMutationAliasingEffects';
+import {inferMutationAliasingFunctionEffects} from './InferMutationAliasingFunctionEffects';
+import {inferMutationAliasingRanges} from './InferMutationAliasingRanges';
 
 export default function analyseFunctions(func: HIRFunction): void {
   for (const [_, block] of func.body.blocks) {
@@ -33,8 +39,13 @@ export default function analyseFunctions(func: HIRFunction): void {
       switch (instr.value.kind) {
         case 'ObjectMethod':
         case 'FunctionExpression': {
-          const aliases = lower(instr.value.loweredFunc.func);
-          infer(instr.value.loweredFunc, aliases);
+          if (!func.env.config.enableNewMutationAliasingModel) {
+            const aliases = lower(instr.value.loweredFunc.func);
+            infer(instr.value.loweredFunc, aliases);
+          } else {
+            lowerWithMutationAliasing(instr.value.loweredFunc.func);
+            infer(instr.value.loweredFunc, new DisjointSet());
+          }
 
           /**
            * Reset mutable range for outer inferReferenceEffects
@@ -51,6 +62,22 @@ export default function analyseFunctions(func: HIRFunction): void {
   }
 }
 
+function lowerWithMutationAliasing(fn: HIRFunction): void {
+  analyseFunctions(fn);
+  inferMutationAliasingEffects(fn, {isFunctionExpression: true});
+  deadCodeElimination(fn);
+  inferMutationAliasingRanges(fn);
+  rewriteInstructionKindsBasedOnReassignment(fn);
+  inferReactiveScopeVariables(fn);
+  fn.env.logger?.debugLogIRs?.({
+    kind: 'hir',
+    name: 'AnalyseFunction (inner)',
+    value: fn,
+  });
+  const effects = inferMutationAliasingFunctionEffects(fn);
+  fn.aliasingEffects = effects;
+}
+
 function lower(func: HIRFunction): DisjointSet<Identifier> {
   analyseFunctions(func);
   inferReferenceEffects(func, {isFunctionExpression: true});
@@ -64,7 +91,7 @@ function lower(func: HIRFunction): DisjointSet<Identifier> {
     value: func,
   });
   inferAliasesForCapturing(func, aliases);
-  return aliases;
+  return aliases ?? new DisjointSet();
 }
 
 /**

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

@@ -38,9 +38,7 @@ import {
   Set_isSuperset,
 } from '../Utils/utils';
 import {
-  printAliasingEffect,
   printIdentifier,
-  printInstruction,
   printInstructionValue,
   printPlace,
   printSourceLocation,
@@ -51,7 +49,7 @@ export function inferMutationAliasingEffects(
   {isFunctionExpression}: {isFunctionExpression: boolean} = {
     isFunctionExpression: false,
   },
-): Result<Array<AliasingEffect>, CompilerError> {
+): Result<void, CompilerError> {
   const initialState = InferenceState.empty(fn.env, isFunctionExpression);
 
   // Map of blocks to the last (merged) incoming state that was processed
@@ -169,7 +167,7 @@ export function inferMutationAliasingEffects(
       }
     }
   }
-  return Ok([]);
+  return Ok(undefined);
 }
 
 function inferParam(
@@ -203,9 +201,11 @@ function inferBlock(
       instructionSignature = computeSignatureForInstruction(state.env, instr);
       instructionSignatureCache.set(instr, instructionSignature);
     }
-    // console.log(
-    //   printInstruction({...instr, effects: [...instructionSignature.effects]}),
-    // );
+    /*
+     * console.log(
+     *   printInstruction({...instr, effects: [...instructionSignature.effects]}),
+     * );
+     */
     const effects = applySignature(
       state,
       instructionSignature,
@@ -304,6 +304,14 @@ function applySignature(
             isCopyByReferenceValue = false;
             break;
           }
+          case ValueKind.Frozen: {
+            /*
+             * TODO: add a separate "ImmutableAlias" effect to downgrade to, that doesn't impact mutable ranges
+             * We want to remember that the data flow occurred for PruneNonEscapingScopes
+             */
+            isCopyByReferenceValue = false;
+            break;
+          }
           default: {
             isCopyByReferenceValue = true;
             break;
@@ -321,6 +329,12 @@ function applySignature(
          */
         const fromKind = state.kind(effect.from).kind;
         switch (fromKind) {
+          /*
+           * TODO: add a separate "ImmutableAlias" effect to downgrade to, that doesn't impact mutable ranges
+           * We want to remember that the data flow occurred for PruneNonEscapingScopes
+           * (use this to replace the ValueKind.Frozen case)
+           */
+          case ValueKind.Frozen:
           case ValueKind.Global:
           case ValueKind.Primitive: {
             let value = effectInstructionValueCache.get(effect);
@@ -1184,8 +1198,8 @@ function computeEffectsForSignature(
   }
   // Build substitutions
   const substitutions: Map<IdentifierId, Array<Place>> = new Map();
-  substitutions.set(signature.receiver.identifier.id, [receiver]);
-  substitutions.set(signature.returns.identifier.id, [lvalue]);
+  substitutions.set(signature.receiver, [receiver]);
+  substitutions.set(signature.returns, [lvalue]);
   const params = signature.params;
   for (let i = 0; i < args.length; i++) {
     const arg = args[i];
@@ -1194,14 +1208,10 @@ function computeEffectsForSignature(
         return null;
       }
       const place = arg.kind === 'Identifier' ? arg : arg.place;
-      getOrInsertWith(
-        substitutions,
-        signature.rest.identifier.id,
-        () => [],
-      ).push(place);
+      getOrInsertWith(substitutions, signature.rest, () => []).push(place);
     } else {
       const param = params[i];
-      substitutions.set(param.identifier.id, [arg]);
+      substitutions.set(param, [arg]);
     }
   }
 
@@ -1414,10 +1424,10 @@ export type AliasingEffect =
     };
 
 export type AliasingSignature = {
-  receiver: Place;
-  params: Array<Place>;
-  rest: Place | null;
-  returns: Place;
+  receiver: IdentifierId;
+  params: Array<IdentifierId>;
+  rest: IdentifierId | null;
+  returns: IdentifierId;
   effects: Array<AliasingEffect>;
 };
 

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

@@ -0,0 +1,136 @@
+/**
+ * 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} from '../HIR';
+import {getOrInsertDefault} from '../Utils/utils';
+import {AliasingEffect} from './InferMutationAliasingEffects';
+
+export function inferMutationAliasingFunctionEffects(
+  fn: HIRFunction,
+): Array<AliasingEffect> | null {
+  const effects: Array<AliasingEffect> = [];
+  /*
+   * Quick hack to infer "mutation" of context vars and params. The problem is that we really want
+   * to figure out more precisely what changed. Is there a known mutation, or just a conditional
+   * mutation?
+   * once we assign scope ids, we can just look through all the effects in the entire body
+   * and find the maximum level of mutation on each scope (for scopes that are on context refs/params)
+   */
+  const tracked = new Map<IdentifierId, Place>();
+  tracked.set(fn.returns.identifier.id, fn.returns);
+  for (const operand of fn.context) {
+    tracked.set(operand.identifier.id, operand);
+    if (operand.identifier.mutableRange.end > 0) {
+      effects.push({kind: 'MutateTransitiveConditionally', value: operand});
+    }
+  }
+  for (const param of fn.params) {
+    const place = param.kind === 'Identifier' ? param : param.place;
+    tracked.set(place.identifier.id, place);
+    if (place.identifier.mutableRange.end > 0) {
+      effects.push({kind: 'MutateTransitiveConditionally', value: place});
+    }
+  }
+  type AliasedIdentifier = {
+    kind: 'Capture' | 'Alias' | 'CreateFrom';
+    place: Place;
+  };
+  const dataFlow = new Map<IdentifierId, Array<AliasedIdentifier>>();
+
+  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 === '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);
+            }
+          }
+        }
+      }
+    }
+    if (block.terminal.kind === 'return') {
+      const from = lookup(block.terminal.value, 'Alias');
+      if (from != null) {
+        getOrInsertDefault(dataFlow, fn.returns.identifier.id, []).push(
+          ...from,
+        );
+      }
+    }
+  }
+
+  for (const [into, from] of dataFlow) {
+    const input = tracked.get(into);
+    if (input == null) {
+      continue;
+    }
+    for (const aliased of from) {
+      const effect = {kind: aliased.kind, from: aliased.place, into: input};
+      effects.push(effect);
+    }
+  }
+
+  return effects;
+}
+
+type AliasingKind = 'Alias' | 'Capture' | 'CreateFrom';
+function joinEffects(
+  effect1: AliasingKind,
+  effect2: AliasingKind,
+): AliasingKind {
+  if (effect1 === 'Capture' || effect2 === 'Capture') {
+    return 'Capture';
+  } else {
+    return 'Alias';
+  }
+}

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

@@ -5,8 +5,20 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {HIRFunction, Identifier, makeInstructionId} from '../HIR/HIR';
+import {
+  Effect,
+  HIRFunction,
+  Identifier,
+  IdentifierId,
+  makeInstructionId,
+} from '../HIR/HIR';
+import {
+  eachInstructionLValue,
+  eachInstructionValueOperand,
+  eachTerminalOperand,
+} from '../HIR/visitors';
 import DisjointSet from '../Utils/DisjointSet';
+import {assertExhaustive} from '../Utils/utils';
 import {inferMutableRangesForAlias} from './InferMutableRangesForAlias';
 
 export function inferMutationAliasingRanges(fn: HIRFunction): void {
@@ -27,6 +39,13 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
     }
 
     for (const instr of block.instructions) {
+      for (const lvalue of eachInstructionLValue(instr)) {
+        lvalue.identifier.mutableRange.start = instr.id;
+        lvalue.identifier.mutableRange.end = makeInstructionId(
+          Math.max(instr.id + 1, lvalue.identifier.mutableRange.end),
+        );
+      }
+
       if (instr.effects == null) continue;
       for (const effect of instr.effects) {
         if (
@@ -72,4 +91,99 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
     }
   }
   inferMutableRangesForAlias(fn, aliases);
+
+  /**
+   * Part 3
+   * Add legacy operand-specific effects based on instruction effects and mutable ranges
+   */
+  for (const block of fn.body.blocks.values()) {
+    for (const phi of block.phis) {
+      // TODO: we don't actually set these effects today!
+      phi.place.effect = Effect.Store;
+      const isPhiMutatedAfterCreation: boolean =
+        phi.place.identifier.mutableRange.end >
+        (block.instructions.at(0)?.id ?? block.terminal.id);
+      for (const operand of phi.operands.values()) {
+        operand.effect = isPhiMutatedAfterCreation
+          ? Effect.Capture
+          : Effect.Read;
+      }
+    }
+    for (const instr of block.instructions) {
+      if (instr.effects == null) {
+        for (const lvalue of eachInstructionLValue(instr)) {
+          lvalue.effect = Effect.ConditionallyMutate;
+        }
+        for (const operand of eachInstructionValueOperand(instr.value)) {
+          operand.effect = Effect.Read;
+        }
+        continue;
+      }
+      const operandEffects = new Map<IdentifierId, Effect>();
+      for (const effect of instr.effects) {
+        switch (effect.kind) {
+          case 'Alias':
+          case 'Capture':
+          case 'CreateFrom': {
+            const isMutatedOrReassigned =
+              effect.into.identifier.mutableRange.end > instr.id;
+            if (isMutatedOrReassigned) {
+              operandEffects.set(effect.from.identifier.id, Effect.Capture);
+              operandEffects.set(effect.into.identifier.id, Effect.Store);
+            } else {
+              operandEffects.set(effect.from.identifier.id, Effect.Read);
+              operandEffects.set(effect.into.identifier.id, Effect.Store);
+            }
+            break;
+          }
+          case 'Create': {
+            break;
+          }
+          case 'Mutate': {
+            operandEffects.set(effect.value.identifier.id, Effect.Store);
+            break;
+          }
+          case 'Apply': {
+            operandEffects.set(
+              effect.function.place.identifier.id,
+              Effect.ConditionallyMutate,
+            );
+            break;
+          }
+          case 'MutateTransitive':
+          case 'MutateConditionally':
+          case 'MutateTransitiveConditionally': {
+            operandEffects.set(
+              effect.value.identifier.id,
+              Effect.ConditionallyMutate,
+            );
+            break;
+          }
+          case 'Freeze': {
+            operandEffects.set(effect.value.identifier.id, Effect.Freeze);
+            break;
+          }
+          default: {
+            assertExhaustive(
+              effect,
+              `Unexpected effect kind ${(effect as any).kind}`,
+            );
+          }
+        }
+      }
+      for (const lvalue of eachInstructionLValue(instr)) {
+        const effect =
+          operandEffects.get(lvalue.identifier.id) ??
+          Effect.ConditionallyMutate;
+        lvalue.effect = effect;
+      }
+      for (const operand of eachInstructionValueOperand(instr.value)) {
+        const effect = operandEffects.get(operand.identifier.id) ?? Effect.Read;
+        operand.effect = effect;
+      }
+    }
+    for (const operand of eachTerminalOperand(block.terminal)) {
+      operand.effect = Effect.Read;
+    }
+  }
 }

compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts

@@ -148,6 +148,19 @@ export function Iterable_some<T>(
   return false;
 }
 
+export function Iterable_filter<T>(
+  iter: Iterable<T>,
+  pred: (item: T) => boolean,
+): Array<T> {
+  const result: Array<T> = [];
+  for (const item of iter) {
+    if (pred(item)) {
+      result.push(item);
+    }
+  }
+  return result;
+}
+
 export function nonNull<T extends NonNullable<U>, U>(
   value: T | null | undefined,
 ): value is T {

compiler/packages/snap/src/sprout/index.ts

@@ -42,6 +42,7 @@ export function runSprout(
     (globalThis as any).__SNAP_EVALUATOR_MODE = undefined;
   }
   if (forgetResult.kind === 'UnexpectedError') {
+    console.log(forgetCode);
     return makeError('Unexpected error in Forget runner', forgetResult.value);
   }
   if (originalCode.indexOf('@disableNonForgetInSprout') === -1) {