[compiler][newinference] Fixes for transitive function capturing, mutation via property loads

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

@@ -66,12 +66,12 @@ function lowerWithMutationAliasing(fn: HIRFunction): void {
   inferMutationAliasingRanges(fn);
   rewriteInstructionKindsBasedOnReassignment(fn);
   inferReactiveScopeVariables(fn);
+  const effects = inferMutationAliasingFunctionEffects(fn);
   fn.env.logger?.debugLogIRs?.({
     kind: 'hir',
     name: 'AnalyseFunction (inner)',
     value: fn,
   });
-  const effects = inferMutationAliasingFunctionEffects(fn);
   if (effects != null) {
     fn.aliasingEffects ??= [];
     fn.aliasingEffects?.push(...effects);

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

@@ -155,7 +155,7 @@ export function inferMutationAliasingEffects(
   }
   queue(fn.body.entry, initialState);
 
-  const context = new Context(isFunctionExpression);
+  const context = new Context(isFunctionExpression, fn);
 
   let count = 0;
   while (queuedStates.size !== 0) {
@@ -193,9 +193,11 @@ class Context {
     new Map();
   catchHandlers: Map<BlockId, Place> = new Map();
   isFuctionExpression: boolean;
+  fn: HIRFunction;
 
-  constructor(isFunctionExpression: boolean) {
+  constructor(isFunctionExpression: boolean, fn: HIRFunction) {
     this.isFuctionExpression = isFunctionExpression;
+    this.fn = fn;
   }
 }
 
@@ -309,8 +311,14 @@ function applySignature(
   ) {
     const aliasingEffects =
       instruction.value.loweredFunc.func.aliasingEffects ?? [];
+    const context = new Set(
+      instruction.value.loweredFunc.func.context.map(p => p.identifier.id),
+    );
     for (const effect of aliasingEffects) {
       if (effect.kind === 'Mutate' || effect.kind === 'MutateTransitive') {
+        if (!context.has(effect.value.identifier.id)) {
+          continue;
+        }
         const value = state.kind(effect.value);
         switch (value.kind) {
           case ValueKind.Global:
@@ -458,7 +466,7 @@ function applyEffect(
         default: {
           effects.push({
             // OK: recording information flow
-            kind: 'Alias',
+            kind: 'CreateFrom', // prev Alias
             from: effect.from,
             into: effect.into,
           });
@@ -467,6 +475,7 @@ function applyEffect(
       break;
     }
     case 'CreateFunction': {
+      effects.push(effect);
       const isMutable = effect.captures.some(capture => {
         switch (state.kind(capture).kind) {
           case ValueKind.Context:
@@ -644,6 +653,13 @@ function applyEffect(
           if (DEBUG) {
             console.log('apply function expression effects');
           }
+          applyEffect(
+            context,
+            state,
+            {kind: 'MutateTransitiveConditionally', value: effect.function},
+            aliased,
+            effects,
+          );
           for (const signatureEffect of signatureEffects) {
             applyEffect(context, state, signatureEffect, aliased, effects);
           }
@@ -1587,10 +1603,26 @@ function computeSignatureForInstruction(
       break;
     }
     case 'StoreGlobal': {
-      CompilerError.throwTodo({
-        reason: `Handle StoreGlobal in new inference`,
-        loc: instr.loc,
+      effects.push({
+        kind: 'MutateGlobal',
+        place: value.value,
+        error: {
+          severity: ErrorSeverity.InvalidReact,
+          reason: getWriteErrorReason({
+            kind: ValueKind.Global,
+            reason: new Set([ValueReason.Global]),
+            context: new Set(),
+          }),
+          description:
+            value.value.identifier.name !== null &&
+            value.value.identifier.name.kind === 'named'
+              ? `Found mutation of \`${value.value.identifier.name}\``
+              : null,
+          loc: value.value.loc,
+          suggestions: null,
+        },
       });
+      break;
     }
     case 'TypeCastExpression': {
       effects.push({kind: 'Assign', from: value.value, into: lvalue});

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

@@ -20,6 +20,10 @@
    */
   const tracked = new Map<IdentifierId, Place>();
   tracked.set(fn.returns.identifier.id, fn.returns);
+  for (const operand of [...fn.context, ...fn.params]) {
+    const place = operand.kind === 'Identifier' ? operand : operand.place;
+    tracked.set(place.identifier.id, place);
+  }
 
   /**
    * Track capturing/aliasing of context vars and params into each other and into the return.
@@ -95,6 +99,13 @@
               ).push(...from);
             }
           }
+        } else if (
+          effect.kind === 'Create' ||
+          effect.kind === 'CreateFunction'
+        ) {
+          getOrInsertDefault(dataFlow, effect.into.identifier.id, [
+            {kind: 'Alias', place: effect.into},
+          ]);
         } else if (
           effect.kind === 'MutateFrozen' ||
           effect.kind === 'MutateGlobal'
@@ -121,6 +132,12 @@
       continue;
     }
     for (const aliased of from) {
+      if (
+        aliased.place.identifier.id === input.identifier.id ||
+        !tracked.has(aliased.place.identifier.id)
+      ) {
+        continue;
+      }
       const effect = {kind: aliased.kind, from: aliased.place, into: input};
       effects.push(effect);
       if (
@@ -133,7 +150,7 @@
   }
   // TODO: more precise return effect inference
   if (!hasReturn) {
-    effects.push({
+    effects.unshift({
       kind: 'Create',
       into: fn.returns,
       value:
@@ -143,6 +160,9 @@
     });
   }
 
+  // console.log(pretty(tracked));
+  // console.log(pretty(dataFlow));
+  // console.log('FunctionEffects', effects.map(printAliasingEffect).join('\n'));
   return effects;
 }
 

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

@@ -107,11 +107,9 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
       for (const effect of instr.effects) {
         if (effect.kind === 'Create' || effect.kind === 'CreateFunction') {
           state.create(effect.into);
-        } else if (
-          effect.kind === 'Assign' ||
-          effect.kind === 'Alias' ||
-          effect.kind === 'CreateFrom'
-        ) {
+        } else if (effect.kind === 'CreateFrom') {
+          state.createFrom(index++, effect.from, effect.into);
+        } else if (effect.kind === 'Assign' || effect.kind === 'Alias') {
           state.assign(index++, effect.from, effect.into);
         } else if (effect.kind === 'Capture') {
           state.capture(index++, effect.from, effect.into);
@@ -181,7 +179,7 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
   for (const mutation of mutations) {
     state.mutate(
       mutation.index,
-      mutation.place,
+      mutation.place.identifier,
       makeInstructionId(mutation.id + 1),
       mutation.transitive,
       mutation.kind,
@@ -191,8 +189,9 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
     console.log(state.debug());
   }
   fn.aliasingEffects ??= [];
-  for (const param of fn.context) {
-    const node = state.nodes.get(param.identifier);
+  for (const param of [...fn.context, ...fn.params]) {
+    const place = param.kind === 'Identifier' ? param : param.place;
+    const node = state.nodes.get(place.identifier);
     if (node == null) {
       continue;
     }
@@ -201,30 +200,30 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
       mutated = true;
       fn.aliasingEffects.push({
         kind: 'MutateConditionally',
-        value: param,
+        value: place,
       });
     } else if (node.local === MutationKind.Definite) {
       mutated = true;
       fn.aliasingEffects.push({
         kind: 'Mutate',
-        value: param,
+        value: place,
       });
     }
     if (node.transitive === MutationKind.Conditional) {
       mutated = true;
       fn.aliasingEffects.push({
         kind: 'MutateTransitiveConditionally',
-        value: param,
+        value: place,
       });
     } else if (node.transitive === MutationKind.Definite) {
       mutated = true;
       fn.aliasingEffects.push({
         kind: 'MutateTransitive',
-        value: param,
+        value: place,
       });
     }
     if (mutated) {
-      param.effect = Effect.Capture;
+      place.effect = Effect.Capture;
     }
   }
 
@@ -247,6 +246,21 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
           ? Effect.Capture
           : Effect.Read;
       }
+      if (
+        isPhiMutatedAfterCreation &&
+        phi.place.identifier.mutableRange.start === 0
+      ) {
+        /*
+         * TODO: ideally we'd construct a precise start range, but what really
+         * matters is that the phi's range appears mutable (end > start + 1)
+         * so we just set the start to the previous instruction before this block
+         */
+        const firstInstructionIdOfBlock =
+          block.instructions.at(0)?.id ?? block.terminal.id;
+        phi.place.identifier.mutableRange.start = makeInstructionId(
+          firstInstructionIdOfBlock - 1,
+        );
+      }
     }
     for (const instr of block.instructions) {
       for (const lvalue of eachInstructionLValue(instr)) {
@@ -357,6 +371,7 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
 
 type Node = {
   id: Identifier;
+  createdFrom: Map<Identifier, number>;
   captures: Map<Identifier, number>;
   aliases: Map<Identifier, number>;
   edges: Array<{index: number; node: Identifier; kind: 'capture' | 'alias'}>;
@@ -369,6 +384,7 @@ class AliasingState {
   create(place: Place): void {
     this.nodes.set(place.identifier, {
       id: place.identifier,
+      createdFrom: new Map(),
       captures: new Map(),
       aliases: new Map(),
       edges: [],
@@ -377,6 +393,24 @@ class AliasingState {
     });
   }
 
+  createFrom(index: number, from: Place, into: Place): void {
+    this.create(into);
+    const fromNode = this.nodes.get(from.identifier);
+    const toNode = this.nodes.get(into.identifier);
+    if (fromNode == null || toNode == null) {
+      if (DEBUG) {
+        console.log(
+          `skip: createFrom ${printPlace(from)}${!!fromNode} -> ${printPlace(into)}${!!toNode}`,
+        );
+      }
+      return;
+    }
+    fromNode.edges.push({index, node: into.identifier, kind: 'alias'});
+    if (!toNode.createdFrom.has(from.identifier)) {
+      toNode.createdFrom.set(from.identifier, index);
+    }
+  }
+
   capture(index: number, from: Place, into: Place): void {
     const fromNode = this.nodes.get(from.identifier);
     const toNode = this.nodes.get(into.identifier);
@@ -406,22 +440,22 @@ class AliasingState {
       return;
     }
     fromNode.edges.push({index, node: into.identifier, kind: 'alias'});
-    if (!toNode.captures.has(from.identifier)) {
+    if (!toNode.aliases.has(from.identifier)) {
       toNode.aliases.set(from.identifier, index);
     }
   }
 
   mutate(
     index: number,
-    start: Place,
+    start: Identifier,
     end: InstructionId,
     transitive: boolean,
     kind: MutationKind,
   ): void {
     const seen = new Set<Identifier>();
-    const queue = [start.identifier];
+    const queue = [{place: start, transitive}];
     while (queue.length !== 0) {
-      const current = queue.pop()!;
+      const {place: current, transitive} = queue.pop()!;
       if (seen.has(current)) {
         continue;
       }
@@ -430,14 +464,14 @@ class AliasingState {
       if (node == null) {
         if (DEBUG) {
           console.log(
-            `no node! ${printPlace(start)} for identifier ${printIdentifier(current)}`,
+            `no node! ${printIdentifier(start)} for identifier ${printIdentifier(current)}`,
           );
         }
         continue;
       }
       if (DEBUG) {
         console.log(
-          `mutate $${node.id.id} via ${printPlace(start)} at [${end}]`,
+          `mutate $${node.id.id} via ${printIdentifier(start)} at [${end}]`,
         );
       }
       node.id.mutableRange.end = makeInstructionId(
@@ -457,7 +491,13 @@ class AliasingState {
         if (edge.index >= index) {
           break;
         }
-        queue.push(edge.node);
+        queue.push({place: edge.node, transitive});
+      }
+      for (const [alias, when] of node.createdFrom) {
+        if (when >= index) {
+          continue;
+        }
+        queue.push({place: alias, transitive: true});
       }
       /**
        * all mutations affect backward alias edges by the rules:
@@ -468,7 +508,7 @@ class AliasingState {
         if (when >= index) {
           continue;
         }
-        queue.push(alias);
+        queue.push({place: alias, transitive});
       }
       /**
        * but only transitive mutations affect captures
@@ -478,7 +518,7 @@ class AliasingState {
           if (when >= index) {
             continue;
           }
-          queue.push(capture);
+          queue.push({place: capture, transitive});
         }
       }
     }
@@ -489,7 +529,7 @@ class AliasingState {
   }
 }
 
-function pretty(v: any): string {
+export function pretty(v: any): string {
   return prettyFormat(v, {
     plugins: [
       {