[compiler] Handle legacy mutableIfOperandsMutable signatures

ghstack-source-id: 3478c26d27
Pull Request resolved: https://github.com/facebook/react/pull/33379

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

@@ -36,7 +36,6 @@ import type {
 } from './HIR';
 import {GotoVariant, InstructionKind} from './HIR';
 import {
-  AliasedPlace,
   AliasingEffect,
   AliasingSignature,
 } from '../Inference/InferMutationAliasingEffects';
@@ -965,10 +964,27 @@ export function printAliasingEffect(effect: AliasingEffect): string {
       return `Create ${printPlaceForAliasEffect(effect.into)} = kindOf(${printPlaceForAliasEffect(effect.from)})`;
     }
     case 'Apply': {
-      const params = effect.params.map(printAliasedPlace).join(', ');
-      const rest = effect.rest != null ? printAliasedPlace(effect.rest) : '';
-      const returns = printAliasedPlace(effect.returns);
-      return `Apply ${returns} = ${printAliasedPlace(effect.function)} as ${effect.receiver} ( ${params} ${params.length !== 0 && rest !== '' ? ', ...' : ''}${rest})`;
+      const receiverCallee =
+        effect.receiver.identifier.id === effect.function.identifier.id
+          ? printPlaceForAliasEffect(effect.receiver)
+          : `${printPlaceForAliasEffect(effect.receiver)}.${printPlaceForAliasEffect(effect.function)}`;
+      const args = effect.args
+        .map(arg => {
+          if (arg.kind === 'Identifier') {
+            return printPlaceForAliasEffect(arg);
+          }
+          return `...${printPlaceForAliasEffect(arg.place)}`;
+        })
+        .join(', ');
+      let signature = '';
+      if (effect.signature != null) {
+        if (effect.signature.aliasing != null) {
+          signature = printAliasingSignature(effect.signature.aliasing);
+        } else {
+          signature = JSON.stringify(effect.signature, null, 2);
+        }
+      }
+      return `Apply ${printPlaceForAliasEffect(effect.into)} = ${receiverCallee}(${args})${signature != '' ? '\n     ' : ''}${signature}`;
     }
     case 'Freeze': {
       return `Freeze ${printPlaceForAliasEffect(effect.value)} ${effect.reason}`;
@@ -989,10 +1005,6 @@ function printPlaceForAliasEffect(place: Place): string {
   return printIdentifier(place.identifier);
 }
 
-function printAliasedPlace(place: AliasedPlace): string {
-  return place.kind + ' ' + printPlaceForAliasEffect(place.place);
-}
-
 export function printAliasingSignature(signature: AliasingSignature): string {
   const tokens: Array<string> = ['function '];
   tokens.push('(');

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

@@ -84,8 +84,10 @@ function lowerWithMutationAliasing(fn: HIRFunction): void {
         break;
       }
       case 'Apply': {
-        capturedOrMutated.add(effect.function.place.identifier.id);
-        break;
+        CompilerError.invariant(false, {
+          reason: `[AnalyzeFunctions] Expected Apply effects to be replaced with more precise effects`,
+          loc: effect.function.loc,
+        });
       }
       case 'Mutate':
       case 'MutateConditionally':

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

@@ -30,6 +30,7 @@ import {
 import {Ok, Result} from '../Utils/Result';
 import {
   getFunctionCallSignature,
+  isKnownMutableEffect,
   mergeValueKinds,
 } from './InferReferenceEffects';
 import {
@@ -284,6 +285,10 @@ function applySignature(
    */
   const aliased = new Set<IdentifierId>();
 
+  if (DEBUG) {
+    console.log(printInstruction(instruction));
+  }
+
   const effects: Array<AliasingEffect> = [];
   for (const effect of signature.effects) {
     applyEffect(
@@ -296,11 +301,6 @@ function applySignature(
     );
   }
   if (DEBUG) {
-    console.log(printInstruction(instruction));
-    console.log('input effects');
-    console.log(signature.effects.map(printAliasingEffect).join('\n'));
-    console.log('refined effects');
-    console.log(effects.map(printAliasingEffect).join('\n'));
     console.log(
       prettyFormat(state.debugAbstractValue(state.kind(instruction.lvalue))),
     );
@@ -324,6 +324,9 @@ function applyEffect(
   aliased: Set<IdentifierId>,
   effects: Array<AliasingEffect>,
 ): void {
+  if (DEBUG) {
+    console.log(printAliasingEffect(effect));
+  }
   switch (effect.kind) {
     case 'Freeze': {
       const didFreeze = state.freeze(effect.value, effect.reason);
@@ -506,23 +509,116 @@ function applyEffect(
       break;
     }
     case 'Apply': {
-      const values = state.values(effect.function.place);
-      if (values.length !== 1 || values[0].kind !== 'FunctionExpression') {
-        const mutationKind = state.mutate(
-          'MutateTransitiveConditionally',
-          effect.function.place,
+      const signatureEffects =
+        effect.signature?.aliasing != null
+          ? computeEffectsForSignature(
+              effect.signature.aliasing,
+              effect.into,
+              effect.receiver,
+              effect.args,
+            )
+          : null;
+      if (signatureEffects != null) {
+        for (const signatureEffect of signatureEffects) {
+          applyEffect(
+            state,
+            signatureEffect,
+            instruction,
+            effectInstructionValueCache,
+            aliased,
+            effects,
+          );
+        }
+      } else if (effect.signature != null) {
+        const legacyEffects = computeEffectsForLegacySignature(
+          state,
+          effect.signature,
+          effect.into,
+          effect.receiver,
+          effect.args,
         );
-        if (mutationKind === 'mutate') {
-          effects.push({
-            kind: 'MutateTransitiveConditionally',
-            value: effect.function.place,
-          });
+        for (const legacyEffect of legacyEffects) {
+          applyEffect(
+            state,
+            legacyEffect,
+            instruction,
+            effectInstructionValueCache,
+            aliased,
+            effects,
+          );
         }
       } else {
-        CompilerError.throwTodo({
-          reason: `Support ${effect.kind} effects`,
-          loc: instruction.loc,
-        });
+        applyEffect(
+          state,
+          {
+            kind: 'Create',
+            into: effect.into,
+            value: ValueKind.Mutable,
+          },
+          instruction,
+          effectInstructionValueCache,
+          aliased,
+          effects,
+        );
+        /*
+         * If no signature then by default:
+         * - All operands are conditionally mutated, except some instruction
+         *   variants are assumed to not mutate the callee (such as `new`)
+         * - All operands are captured into (but not directly aliased as)
+         *   every other argument.
+         */
+        for (const arg of [effect.function, ...effect.args]) {
+          const operand = arg.kind === 'Identifier' ? arg : arg.place;
+          if (operand !== effect.function || effect.mutatesFunction) {
+            applyEffect(
+              state,
+              {
+                kind: 'MutateTransitiveConditionally',
+                value: operand,
+              },
+              instruction,
+              effectInstructionValueCache,
+              aliased,
+              effects,
+            );
+          }
+          /*
+           * TODO: this should be Alias, since the function could be identity.
+           * Ie local mutation of the result could change the input.
+           * But if we emit multiple Alias calls, currently the last one will win
+           * when we update the inferencestate in applySignature. So we may need to group
+           * them here, or coalesce them in applySignature
+           *
+           * maybe make `from: Place | Array<Place>`
+           */
+          applyEffect(
+            state,
+            {kind: 'Capture', from: operand, into: effect.into},
+            instruction,
+            effectInstructionValueCache,
+            aliased,
+            effects,
+          );
+          for (const otherArg of [effect.function, ...effect.args]) {
+            const other =
+              otherArg.kind === 'Identifier' ? otherArg : otherArg.place;
+            if (other === arg) {
+              continue;
+            }
+            applyEffect(
+              state,
+              {
+                kind: 'Capture',
+                from: operand,
+                into: other,
+              },
+              instruction,
+              effectInstructionValueCache,
+              aliased,
+              effects,
+            );
+          }
+        }
       }
       break;
     }
@@ -1050,7 +1146,7 @@ function computeSignatureForInstruction(
     case 'MethodCall': {
       let callee;
       let receiver;
-      let mutatesCallee = false;
+      let mutatesCallee;
       if (value.kind === 'NewExpression') {
         callee = value.callee;
         receiver = value.callee;
@@ -1070,72 +1166,15 @@ function computeSignatureForInstruction(
         );
       }
       const signature = getFunctionCallSignature(env, callee.identifier.type);
-      const signatureEffects =
-        signature != null && signature.aliasing != null
-          ? computeEffectsForSignature(
-              signature.aliasing,
-              lvalue,
-              receiver,
-              value.args,
-            )
-          : null;
-      /*
-       * TODO: we need to know the kinds of the receiver/operands in order to replicate
-       * legacy function behaviors. In particular things like mutableOnlyIfOperandsAreMutable.
-       * we could add alias signatures for all of these, but it is also very reasonable to
-       * emit an Apply effect here — including the signature! — and then let applySignature
-       * do the work of interpreting the signature. It would mean applySignature has to recurse,
-       * but that actually feels inevitable. We already had an Apply effect anyway!
-       */
-      if (signatureEffects != null) {
-        effects.push(...signatureEffects);
-      } else if (signature != null) {
-        effects.push(
-          ...computeEffectsForLegacySignature(
-            signature,
-            lvalue,
-            receiver,
-            value.args,
-          ),
-        );
-      } else {
-        effects.push({kind: 'Create', into: lvalue, value: ValueKind.Mutable});
-        /**
-         * If no signature then by default:
-         * - All operands are conditionally mutated, except some instruction
-         *   variants are assumed to not mutate the callee (such as `new`)
-         * - All operands are captured into (but not directly aliased as)
-         *   every other argument.
-         */
-        for (const operand of eachInstructionValueOperand(value)) {
-          if (operand !== callee || mutatesCallee) {
-            effects.push({
-              kind: 'MutateTransitiveConditionally',
-              value: operand,
-            });
-          }
-          /*
-           * TODO: this should be Alias, since the function could be identity.
-           * Ie local mutation of the result could change the input.
-           * But if we emit multiple Alias calls, currently the last one will win
-           * when we update the inferencestate in applySignature. So we may need to group
-           * them here, or coalesce them in applySignature
-           *
-           * maybe make `from: Place | Array<Place>`
-           */
-          effects.push({kind: 'Capture', from: operand, into: lvalue});
-          for (const other of eachInstructionValueOperand(value)) {
-            if (other === operand) {
-              continue;
-            }
-            effects.push({
-              kind: 'Capture',
-              from: operand,
-              into: other,
-            });
-          }
-        }
-      }
+      effects.push({
+        kind: 'Apply',
+        receiver,
+        function: callee,
+        mutatesFunction: mutatesCallee,
+        args: value.args,
+        into: lvalue,
+        signature,
+      });
       break;
     }
     case 'PropertyDelete':
@@ -1437,6 +1476,7 @@ function computeSignatureForInstruction(
  * compilation output.
  */
 function computeEffectsForLegacySignature(
+  state: InferenceState,
   signature: FunctionSignature,
   lvalue: Place,
   receiver: Place,
@@ -1448,19 +1488,6 @@ function computeEffectsForLegacySignature(
     into: lvalue,
     value: signature.returnValueKind,
   });
-  /*
-   * InferReferenceEffects and FunctionSignature have an implicit assumption that the receiver
-   * is captured into the return value. Consider for example the signature for Array.proto.pop:
-   * the calleeEffect is Store, since it's a known mutation but non-transitive. But the return
-   * of the pop() captures from the receiver! This isn't specified explicitly. So we add this
-   * here, and rely on applySignature() to downgrade this to ImmutableCapture (or prune) if
-   * the type doesn't actually need to be captured based on the input and return type.
-   */
-  effects.push({
-    kind: 'Capture',
-    from: receiver,
-    into: lvalue,
-  });
   const stores: Array<Place> = [];
   const captures: Array<Place> = [];
   const returnValueReason =
@@ -1521,6 +1548,41 @@ function computeEffectsForLegacySignature(
     }
   }
 
+  if (
+    signature.mutableOnlyIfOperandsAreMutable &&
+    areArgumentsImmutableAndNonMutating(state, args)
+  ) {
+    effects.push({
+      kind: 'Capture',
+      from: receiver,
+      into: lvalue,
+    });
+    for (const arg of args) {
+      const place = arg.kind === 'Identifier' ? arg : arg.place;
+      effects.push({
+        kind: 'ImmutableCapture',
+        from: place,
+        into: lvalue,
+      });
+    }
+    return effects;
+  }
+
+  if (signature.calleeEffect !== Effect.Capture) {
+    /*
+     * InferReferenceEffects and FunctionSignature have an implicit assumption that the receiver
+     * is captured into the return value. Consider for example the signature for Array.proto.pop:
+     * the calleeEffect is Store, since it's a known mutation but non-transitive. But the return
+     * of the pop() captures from the receiver! This isn't specified explicitly. So we add this
+     * here, and rely on applySignature() to downgrade this to ImmutableCapture (or prune) if
+     * the type doesn't actually need to be captured based on the input and return type.
+     */
+    effects.push({
+      kind: 'Capture',
+      from: receiver,
+      into: lvalue,
+    });
+  }
   visit(receiver, signature.calleeEffect);
   for (let i = 0; i < args.length; i++) {
     const arg = args[i];
@@ -1549,6 +1611,66 @@ function computeEffectsForLegacySignature(
   return effects;
 }
 
+/**
+ * Returns true if all of the arguments are both non-mutable (immutable or frozen)
+ * _and_ are not functions which might mutate their arguments. Note that function
+ * expressions count as frozen so long as they do not mutate free variables: this
+ * function checks that such functions also don't mutate their inputs.
+ */
+function areArgumentsImmutableAndNonMutating(
+  state: InferenceState,
+  args: Array<Place | SpreadPattern>,
+): boolean {
+  for (const arg of args) {
+    if (arg.kind === 'Identifier' && arg.identifier.type.kind === 'Function') {
+      const fnShape = state.env.getFunctionSignature(arg.identifier.type);
+      if (fnShape != null) {
+        return (
+          !fnShape.positionalParams.some(isKnownMutableEffect) &&
+          (fnShape.restParam == null ||
+            !isKnownMutableEffect(fnShape.restParam))
+        );
+      }
+    }
+    const place = arg.kind === 'Identifier' ? arg : arg.place;
+
+    const kind = state.kind(place).kind;
+    switch (kind) {
+      case ValueKind.Primitive:
+      case ValueKind.Frozen: {
+        /*
+         * Only immutable values, or frozen lambdas are allowed.
+         * A lambda may appear frozen even if it may mutate its inputs,
+         * so we have a second check even for frozen value types
+         */
+        break;
+      }
+      default: {
+        /**
+         * Globals, module locals, and other locally defined functions may
+         * mutate their arguments.
+         */
+        return false;
+      }
+    }
+    const values = state.values(place);
+    for (const value of values) {
+      if (
+        value.kind === 'FunctionExpression' &&
+        value.loweredFunc.func.params.some(param => {
+          const place = param.kind === 'Identifier' ? param : param.place;
+          const range = place.identifier.mutableRange;
+          return range.end > range.start + 1;
+        })
+      ) {
+        // This is a function which may mutate its inputs
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
 function computeEffectsForSignature(
   signature: AliasingSignature,
   lvalue: Place,
@@ -1779,11 +1901,12 @@ export type AliasingEffect =
    */
   | {
       kind: 'Apply';
-      function: AliasedPlace;
-      receiver: AliasedPlace;
-      params: Array<AliasedPlace>;
-      rest: AliasedPlace | null;
-      returns: AliasedPlace;
+      receiver: Place;
+      function: Place;
+      mutatesFunction: boolean;
+      args: Array<Place | SpreadPattern>;
+      into: Place;
+      signature: FunctionSignature | null;
     };
 
 export type AliasingSignature = {

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

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import {CompilerError} from '..';
 import {
   Effect,
   HIRFunction,
@@ -181,11 +182,10 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
             break;
           }
           case 'Apply': {
-            operandEffects.set(
-              effect.function.place.identifier.id,
-              Effect.ConditionallyMutate,
-            );
-            break;
+            CompilerError.invariant(false, {
+              reason: `[AnalyzeFunctions] Expected Apply effects to be replaced with more precise effects`,
+              loc: effect.function.loc,
+            });
           }
           case 'MutateTransitive':
           case 'MutateConditionally':

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/array-filter.expect.md

@@ -0,0 +1,93 @@
+
+## Input
+
+```javascript
+// @enableNewMutationAliasingModel
+function Component({value}) {
+  const arr = [{value: 'foo'}, {value: 'bar'}, {value}];
+  useIdentity(null);
+  const derived = arr.filter(Boolean);
+  return (
+    <Stringify>
+      {derived.at(0)}
+      {derived.at(-1)}
+    </Stringify>
+  );
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enableNewMutationAliasingModel
+function Component(t0) {
+  const $ = _c(13);
+  const { value } = t0;
+  let t1;
+  let t2;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = { value: "foo" };
+    t2 = { value: "bar" };
+    $[0] = t1;
+    $[1] = t2;
+  } else {
+    t1 = $[0];
+    t2 = $[1];
+  }
+  let t3;
+  if ($[2] !== value) {
+    t3 = [t1, t2, { value }];
+    $[2] = value;
+    $[3] = t3;
+  } else {
+    t3 = $[3];
+  }
+  const arr = t3;
+  useIdentity(null);
+  let t4;
+  if ($[4] !== arr) {
+    t4 = arr.filter(Boolean);
+    $[4] = arr;
+    $[5] = t4;
+  } else {
+    t4 = $[5];
+  }
+  const derived = t4;
+  let t5;
+  if ($[6] !== derived) {
+    t5 = derived.at(0);
+    $[6] = derived;
+    $[7] = t5;
+  } else {
+    t5 = $[7];
+  }
+  let t6;
+  if ($[8] !== derived) {
+    t6 = derived.at(-1);
+    $[8] = derived;
+    $[9] = t6;
+  } else {
+    t6 = $[9];
+  }
+  let t7;
+  if ($[10] !== t5 || $[11] !== t6) {
+    t7 = (
+      <Stringify>
+        {t5}
+        {t6}
+      </Stringify>
+    );
+    $[10] = t5;
+    $[11] = t6;
+    $[12] = t7;
+  } else {
+    t7 = $[12];
+  }
+  return t7;
+}
+
+```
+      
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/new-mutability/array-filter.js

@@ -0,0 +1,12 @@
+// @enableNewMutationAliasingModel
+function Component({value}) {
+  const arr = [{value: 'foo'}, {value: 'bar'}, {value}];
+  useIdentity(null);
+  const derived = arr.filter(Boolean);
+  return (
+    <Stringify>
+      {derived.at(0)}
+      {derived.at(-1)}
+    </Stringify>
+  );
+}