[compiler] Improve inference of function expression mutation/aliasing effects

Distinguish the various forms of mutation, and do basic replaying of these effects when the function is created (imperfect, temporary).

ghstack-source-id: c6d6d897a1
Pull Request resolved: https://github.com/facebook/react/pull/33365

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

@@ -71,6 +71,8 @@ export function printFunction(fn: HIRFunction): string {
         })
         .join(', ') +
       ')';
+  } else {
+    definition += '()';
   }
   if (definition.length !== 0) {
     output.push(definition);
@@ -573,8 +575,11 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
             }
           })
           .join(', ') ?? '';
-      const type = printType(instrValue.loweredFunc.func.returnType).trim();
-      value = `${kind} ${name} @context[${context}] @effects[${effects}]${type !== '' ? ` return${type}` : ''}:\n${fn}`;
+      const aliasingEffects =
+        instrValue.loweredFunc.func.aliasingEffects
+          ?.map(printAliasingEffect)
+          ?.join(', ') ?? '';
+      value = `${kind} ${name} @context[${context}] @effects[${effects}] @aliasingEffects=[${aliasingEffects}]\n${fn}`;
       break;
     }
     case 'TaggedTemplateExpression': {

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

@@ -998,6 +998,9 @@ function computeSignatureForInstruction(
         into: lvalue,
         value: ValueKind.Mutable,
       });
+      if (value.loweredFunc.func.aliasingEffects != null) {
+        effects.push(...value.loweredFunc.func.aliasingEffects);
+      }
       break;
     }
     case 'GetIterator': {

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

@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {HIRFunction, IdentifierId, Place} from '../HIR';
+import {HIRFunction, IdentifierId, Place, ScopeId} from '../HIR';
 import {getOrInsertDefault} from '../Utils/utils';
 import {AliasingEffect} from './InferMutationAliasingEffects';
 
@@ -13,34 +13,61 @@ 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)
+
+  /**
+   * Map used to identify tracked variables: params, context vars, return value
+   * This is used to detect mutation/capturing/aliasing of params/context vars
    */
   const tracked = new Map<IdentifierId, Place>();
   tracked.set(fn.returns.identifier.id, fn.returns);
+
+  /**
+   * For each reactive scope we track whether there are known/conditional local and transitive
+   * mutations. This is used to recover precise mutation effects for each of the params and
+   * context variables.
+   */
+  const trackedScopes = new Map<
+    ScopeId,
+    {local: MutationKind; transitive: MutationKind}
+  >();
   for (const operand of fn.context) {
     tracked.set(operand.identifier.id, operand);
-    if (operand.identifier.mutableRange.end > 0) {
-      effects.push({kind: 'MutateTransitiveConditionally', value: operand});
+    if (operand.identifier.scope != null) {
+      getOrInsertDefault(trackedScopes, operand.identifier.scope.id, {
+        local: MutationKind.None,
+        transitive: MutationKind.None,
+      });
     }
   }
   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});
+    if (place.identifier.scope != null) {
+      getOrInsertDefault(trackedScopes, place.identifier.scope.id, {
+        local: MutationKind.None,
+        transitive: MutationKind.None,
+      });
     }
   }
+
+  /**
+   * Track capturing/aliasing of context vars and params into each other and into the return.
+   * We don't need to track locals and intermediate values, since we're only concerned with effects
+   * as they relate to arguments visible outside the function.
+   *
+   * For each aliased identifier we track capture/alias/createfrom and then merge this with how
+   * the value is used. Eg capturing an alias => capture. See joinEffects() helper.
+   */
   type AliasedIdentifier = {
     kind: 'Capture' | 'Alias' | 'CreateFrom';
     place: Place;
   };
   const dataFlow = new Map<IdentifierId, Array<AliasedIdentifier>>();
 
+  /*
+   * Check for aliasing of tracked values. Also joins the effects of how the value is
+   * used (@param kind) with the aliasing type of each value
+   */
   function lookup(
     place: Place,
     kind: AliasedIdentifier['kind'],
@@ -96,6 +123,37 @@ export function inferMutationAliasingFunctionEffects(
               ).push(...from);
             }
           }
+        } else if (
+          effect.kind === 'Mutate' &&
+          effect.value.identifier.scope != null &&
+          trackedScopes.has(effect.value.identifier.scope.id)
+        ) {
+          const scope = trackedScopes.get(effect.value.identifier.scope.id)!;
+          scope.local = MutationKind.Definite;
+        } else if (
+          effect.kind === 'MutateConditionally' &&
+          effect.value.identifier.scope != null &&
+          trackedScopes.has(effect.value.identifier.scope.id)
+        ) {
+          const scope = trackedScopes.get(effect.value.identifier.scope.id)!;
+          scope.local = Math.max(MutationKind.Conditional, scope.local);
+        } else if (
+          effect.kind === 'MutateTransitive' &&
+          effect.value.identifier.scope != null &&
+          trackedScopes.has(effect.value.identifier.scope.id)
+        ) {
+          const scope = trackedScopes.get(effect.value.identifier.scope.id)!;
+          scope.transitive = MutationKind.Definite;
+        } else if (
+          effect.kind === 'MutateTransitiveConditionally' &&
+          effect.value.identifier.scope != null &&
+          trackedScopes.has(effect.value.identifier.scope.id)
+        ) {
+          const scope = trackedScopes.get(effect.value.identifier.scope.id)!;
+          scope.transitive = Math.max(
+            MutationKind.Conditional,
+            scope.transitive,
+          );
         }
       }
     }
@@ -109,6 +167,27 @@ export function inferMutationAliasingFunctionEffects(
     }
   }
 
+  // Create mutation effects based on observed mutation types
+  for (const value of tracked.values()) {
+    if (
+      value.identifier.id === fn.returns.identifier.id ||
+      value.identifier.scope == null
+    ) {
+      continue;
+    }
+    const scope = trackedScopes.get(value.identifier.scope.id)!;
+    if (scope.local === MutationKind.Definite) {
+      effects.push({kind: 'Mutate', value});
+    } else if (scope.local === MutationKind.Conditional) {
+      effects.push({kind: 'MutateConditionally', value});
+    }
+    if (scope.transitive === MutationKind.Definite) {
+      effects.push({kind: 'MutateTransitive', value});
+    } else if (scope.transitive === MutationKind.Conditional) {
+      effects.push({kind: 'MutateTransitiveConditionally', value});
+    }
+  }
+  // Create aliasing effects based on observed data flow
   for (const [into, from] of dataFlow) {
     const input = tracked.get(into);
     if (input == null) {
@@ -123,6 +202,12 @@ export function inferMutationAliasingFunctionEffects(
   return effects;
 }
 
+enum MutationKind {
+  None = 0,
+  Conditional = 1,
+  Definite = 2,
+}
+
 type AliasingKind = 'Alias' | 'Capture' | 'CreateFrom';
 function joinEffects(
   effect1: AliasingKind,

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/InferReactiveScopeVariables.ts

@@ -376,12 +376,12 @@ export function findDisjointMutableValues(
       } else {
         for (const operand of eachInstructionOperand(instr)) {
           if (
-            isMutable(instr, operand) &&
+            isMutable(instr, operand)
             /*
              * exclude global variables from being added to scopes, we can't recreate them!
              * TODO: improve handling of module-scoped variables and globals
              */
-            operand.identifier.mutableRange.start > 0
+            // && operand.identifier.mutableRange.start > 0
           ) {
             if (
               instr.value.kind === 'FunctionExpression' ||