[compiler] Alternate pipeline for new mutability model

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]);
+      }
     }
-  }
-  inferMutationAliasingRanges(hir);
-  log({kind: 'hir', name: 'InferMutationAliasingRanges', value: hir});
-
-  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();
+      }
     }
   }
-  log({kind: 'hir', name: 'InferReferenceEffects', value: hir});
 
-  validateLocalsNotReassignedAfterRender(hir);
+  if (!env.config.enableNewMutationAliasingModel) {
+    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,15 +26,26 @@
   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) {
     for (const instr of block.instructions) {
       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 @@
   }
 }
 
+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 @@
     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>;
 };