Mirror of upstream PR #33184

compiler/.gitignore

@@ -6,6 +6,9 @@
 debug/
 target/
 
+nocommit*.js
+nocommit*.expect.md
+
 # These are backup files generated by rustfmt
 **/*.rs.bk
 

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

@@ -300,6 +300,10 @@ export type FunctionEffect =
       places: ReadonlySet<Place>;
       effect: Effect;
       loc: SourceLocation;
+    }
+  | {
+      kind: 'CaptureEffect';
+      places: ReadonlySet<Place>;
     };
 
 /*
@@ -1430,6 +1434,46 @@ export const ValueKindSchema = z.enum([
   ValueKind.Context,
 ]);
 
+export type AliasingEffect =
+  /**
+   * Freezes the operand if not already frozen
+   */
+  | {kind: 'Freeze'; place: Place}
+  /**
+   * Known mutation of a value, which may effect that value or values that it contains.
+   * This is rare!
+   */
+  | {kind: 'MutateTransitive'; place: Place}
+  /**
+   * Known mutation of a specific value, targeting only that specific value but not
+   * values that it contains.
+   *
+   * Example: `array.push(item)` mutates the array but does not mutate items stored in the array.
+   */
+  | {kind: 'MutateLocal'; place: Place}
+  /**
+   * Possible mutation of a specific value
+   */
+  | {kind: 'ConditionallyMutate'; place: Place}
+  /**
+   * Direct aliasing of one identifier by another identifier
+   * Examples: `x = y` (from y -> to x) or phis (from operand -> to phi)
+   */
+  | {kind: 'Alias'; from: Place; to: Place}
+  /**
+   * Direct aliasing of an identifier (or a sub-path), or storing a value/sub-path
+   * into a part of another object.
+   *
+   * One of from.path and/or to.path must be non-null (else this is equivalent to Alias)
+   */
+  | {
+      kind: 'Capture';
+      from: {place: Place; path: '*' | null};
+      to: {place: Place; path: '*' | null};
+    }
+  // Known mutation of a global
+  | {kind: 'MutateGlobal'; place: Place};
+
 // The effect with which a value is modified.
 export enum Effect {
   // Default value: not allowed after lifetime inference

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

@@ -546,12 +546,23 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       const effects =
         instrValue.loweredFunc.func.effects
           ?.map(effect => {
-            if (effect.kind === 'ContextMutation') {
-              return `ContextMutation places=[${[...effect.places]
-                .map(place => printPlace(place))
-                .join(', ')}] effect=${effect.effect}`;
-            } else {
-              return `GlobalMutation`;
+            switch (effect.kind) {
+              case 'ContextMutation': {
+                return `ContextMutation places=[${[...effect.places]
+                  .map(place => printPlace(place))
+                  .join(', ')}] effect=${effect.effect}`;
+              }
+              case 'GlobalMutation': {
+                return 'GlobalMutation';
+              }
+              case 'ReactMutation': {
+                return 'ReactMutation';
+              }
+              case 'CaptureEffect': {
+                return `CaptureEffect places=[${[...effect.places]
+                  .map(place => printPlace(place))
+                  .join(', ')}]`;
+              }
             }
           })
           .join(', ') ?? '';
@@ -720,7 +731,7 @@ function isMutable(range: MutableRange): boolean {
 }
 
 const DEBUG_MUTABLE_RANGES = false;
-function printMutableRange(identifier: Identifier): string {
+export function printMutableRange(identifier: Identifier): string {
   if (DEBUG_MUTABLE_RANGES) {
     // if debugging, print both the identifier and scope range if they differ
     const range = identifier.mutableRange;

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

@@ -11,6 +11,7 @@ import {
   HIRFunction,
   Identifier,
   LoweredFunction,
+  Place,
   isRefOrRefValue,
   makeInstructionId,
 } from '../HIR';
@@ -19,15 +20,21 @@ import {inferReactiveScopeVariables} from '../ReactiveScopes';
 import {rewriteInstructionKindsBasedOnReassignment} from '../SSA';
 import {inferMutableRanges} from './InferMutableRanges';
 import inferReferenceEffects from './InferReferenceEffects';
+import DisjointSet from '../Utils/DisjointSet';
+import {
+  eachInstructionLValue,
+  eachInstructionValueOperand,
+} from '../HIR/visitors';
+import {Iterable_some} from '../Utils/utils';
 
 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': {
-          lower(instr.value.loweredFunc.func);
-          infer(instr.value.loweredFunc);
+          const aliases = lower(instr.value.loweredFunc.func);
+          infer(instr.value.loweredFunc, aliases);
 
           /**
            * Reset mutable range for outer inferReferenceEffects
@@ -44,21 +51,73 @@ export default function analyseFunctions(func: HIRFunction): void {
   }
 }
 
-function lower(func: HIRFunction): void {
+function lower(func: HIRFunction): DisjointSet<Identifier> {
   analyseFunctions(func);
   inferReferenceEffects(func, {isFunctionExpression: true});
   deadCodeElimination(func);
-  inferMutableRanges(func);
+  const aliases = inferMutableRanges(func);
   rewriteInstructionKindsBasedOnReassignment(func);
   inferReactiveScopeVariables(func);
   func.env.logger?.debugLogIRs?.({
     kind: 'hir',
     name: 'AnalyseFunction (inner)',
     value: func,
   });
+  inferAliasesForCapturing(func, aliases);
+  return aliases;
+}
+
+/**
+ * The alias sets returned by InferMutableRanges() accounts only for aliases that
+ * are known to mutate together. Notably this skips cases where a value is captured
+ * into some other value, but neither is subsequently mutated. An example is pushing
+ * a mutable value onto an array, where neither the array or value are subsequently
+ * mutated.
+ *
+ * This function extends the aliases sets to account for such capturing, so that we
+ * can detect cases where one of the values in a set is mutated later (in an outer function)
+ * we can correctly infer them as mutating together.
+ */
+function inferAliasesForCapturing(
+  fn: HIRFunction,
+  aliases: DisjointSet<Identifier>,
+): void {
+  for (const block of fn.body.blocks.values()) {
+    for (const instr of block.instructions) {
+      const {lvalue, value} = instr;
+      const hasCapture =
+        lvalue.effect === Effect.Store ||
+        Iterable_some(
+          eachInstructionValueOperand(value),
+          operand => operand.effect === Effect.Capture,
+        );
+      if (!hasCapture) {
+        continue;
+      }
+      const operands: Array<Identifier> = [];
+      for (const lvalue of eachInstructionLValue(instr)) {
+        operands.push(lvalue.identifier);
+      }
+      for (const operand of eachInstructionValueOperand(instr.value)) {
+        if (
+          operand.effect === Effect.Store ||
+          operand.effect === Effect.Mutate ||
+          operand.effect === Effect.Capture
+        ) {
+          operands.push(operand.identifier);
+        }
+      }
+      if (operands.length > 1) {
+        aliases.union(operands);
+      }
+    }
+  }
 }
 
-function infer(loweredFunc: LoweredFunction): void {
+function infer(
+  loweredFunc: LoweredFunction,
+  aliases: DisjointSet<Identifier>,
+): void {
   for (const operand of loweredFunc.func.context) {
     const identifier = operand.identifier;
     CompilerError.invariant(operand.effect === Effect.Unknown, {
@@ -85,6 +144,23 @@ function infer(loweredFunc: LoweredFunction): void {
       operand.effect = Effect.Read;
     }
   }
+  const contextIdentifiers = new Map(
+    loweredFunc.func.context.map(place => [place.identifier, place]),
+  );
+  for (const set of aliases.buildSets()) {
+    const contextOperands: Set<Place> = new Set(
+      [...set]
+        .map(identifier => contextIdentifiers.get(identifier))
+        .filter(place => place != null) as Array<Place>,
+    );
+    if (contextOperands.size !== 0) {
+      loweredFunc.func.effects ??= [];
+      loweredFunc.func.effects?.push({
+        kind: 'CaptureEffect',
+        places: contextOperands,
+      });
+    }
+  }
 }
 
 function isMutatedOrReassigned(id: Identifier): boolean {

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

@@ -60,6 +60,10 @@ function inferInstr(
       alias = instrValue.value;
       break;
     }
+    case 'IteratorNext': {
+      alias = instrValue.collection;
+      break;
+    }
     default:
       return;
   }

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

@@ -0,0 +1,33 @@
+/**
+ * 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, Identifier} from '../HIR/HIR';
+import DisjointSet from '../Utils/DisjointSet';
+
+export function inferAliasForFunctionCaptureEffects(
+  func: HIRFunction,
+  aliases: DisjointSet<Identifier>,
+): void {
+  for (const [_, block] of func.body.blocks) {
+    for (const instr of block.instructions) {
+      const {value} = instr;
+      if (
+        value.kind !== 'FunctionExpression' &&
+        value.kind !== 'ObjectMethod'
+      ) {
+        continue;
+      }
+      const loweredFunction = value.loweredFunc.func;
+      for (const effect of loweredFunction.effects ?? []) {
+        if (effect.kind !== 'CaptureEffect') {
+          continue;
+        }
+        aliases.union([...effect.places].map(place => place.identifier));
+      }
+    }
+  }
+}