[compiler] Optimize props spread for common cases

As part of the new inference model we updated to (correctly) treat destructuring spread as creating a new mutable object. This had the unfortunate side-effect of reducing precision on destructuring of props, though:

```js
function Component({x, ...rest}) {
  const z = rest.z;
  identity(z);
  return <Stringify x={x} z={z} />;
}
```

Memoized as the following, where we don't realize that `z` is actually frozen:

```js
function Component(t0) {
  const $ = _c(6);
  let x;
  let z;
  if ($[0] !== t0) {
    const { x: t1, ...rest } = t0;
    x = t1;
    z = rest.z;
    identity(z);
...
```

#34341 was our first thought of how to do this (thanks @poteto for exploring this idea!). But during review it became clear that it was a bit more complicated than I had thought. So this PR explores a more conservative alternative. The idea is:

* Track known sources of frozen values: component props, hook params, and hook return values.
* Find all object spreads where the rvalue is a known frozen value.
* Look at how such objects are used, and if they are only used to access properties (PropertyLoad/Destructure), pass to hooks, or pass to jsx then we can be very confident the object is not mutated. We consider any such objects to be frozen, even though technically spread creates a new object.

See new fixtures for more examples.

Author: josephsavona

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

@@ -19,6 +19,7 @@ import {
   Environment,
   FunctionExpression,
   GeneratedSource,
+  getHookKind,
   HIRFunction,
   Hole,
   IdentifierId,
@@ -198,6 +199,7 @@ export function inferMutationAliasingEffects(
     isFunctionExpression,
     fn,
     hoistedContextDeclarations,
+    findNonMutatedDestructureSpreads(fn),
   );
 
   let iterationCount = 0;
@@ -287,15 +289,18 @@ class Context {
   isFuctionExpression: boolean;
   fn: HIRFunction;
   hoistedContextDeclarations: Map<DeclarationId, Place | null>;
+  nonMutatingSpreads: Set<IdentifierId>;
 
   constructor(
     isFunctionExpression: boolean,
     fn: HIRFunction,
     hoistedContextDeclarations: Map<DeclarationId, Place | null>,
+    nonMutatingSpreads: Set<IdentifierId>,
   ) {
     this.isFuctionExpression = isFunctionExpression;
     this.fn = fn;
     this.hoistedContextDeclarations = hoistedContextDeclarations;
+    this.nonMutatingSpreads = nonMutatingSpreads;
   }
 
   cacheApplySignature(
@@ -322,6 +327,161 @@ class Context {
   }
 }
 
+/**
+ * Finds objects created via ObjectPattern spread destructuring
+ * (`const {x, ...spread} = ...`) where a) the rvalue is known frozen and
+ * b) the spread value cannot possibly be directly mutated. The idea is that
+ * for this set of values, we can treat the spread object as frozen.
+ *
+ * The primary use case for this is props spreading:
+ *
+ * ```
+ * function Component({prop, ...otherProps}) {
+ *   const transformedProp = transform(prop, otherProps.foo);
+ *   // pass `otherProps` down:
+ *   return <Foo {...otherProps} prop={transformedProp} />;
+ * }
+ * ```
+ *
+ * Here we know that since `otherProps` cannot be mutated, we don't have to treat
+ * it as mutable: `otherProps.foo` only reads a value that must be frozen, so it
+ * can be treated as frozen too.
+ */
+function findNonMutatedDestructureSpreads(fn: HIRFunction): Set<IdentifierId> {
+  const knownFrozen = new Set<IdentifierId>();
+  if (fn.fnType === 'Component') {
+    const [props] = fn.params;
+    if (props != null && props.kind === 'Identifier') {
+      knownFrozen.add(props.identifier.id);
+    }
+  } else {
+    for (const param of fn.params) {
+      if (param.kind === 'Identifier') {
+        knownFrozen.add(param.identifier.id);
+      }
+    }
+  }
+
+  // Map of temporaries to identifiers for spread objects
+  const candidateNonMutatingSpreads = new Map<IdentifierId, IdentifierId>();
+  for (const block of fn.body.blocks.values()) {
+    if (candidateNonMutatingSpreads.size !== 0) {
+      for (const phi of block.phis) {
+        for (const operand of phi.operands.values()) {
+          const spread = candidateNonMutatingSpreads.get(operand.identifier.id);
+          if (spread != null) {
+            candidateNonMutatingSpreads.delete(spread);
+          }
+        }
+      }
+    }
+    for (const instr of block.instructions) {
+      const {lvalue, value} = instr;
+      switch (value.kind) {
+        case 'Destructure': {
+          if (
+            !knownFrozen.has(value.value.identifier.id) ||
+            !(
+              value.lvalue.kind === InstructionKind.Let ||
+              value.lvalue.kind === InstructionKind.Const
+            ) ||
+            value.lvalue.pattern.kind !== 'ObjectPattern'
+          ) {
+            continue;
+          }
+          for (const item of value.lvalue.pattern.properties) {
+            if (item.kind !== 'Spread') {
+              continue;
+            }
+            candidateNonMutatingSpreads.set(
+              item.place.identifier.id,
+              item.place.identifier.id,
+            );
+          }
+          break;
+        }
+        case 'LoadLocal': {
+          const spread = candidateNonMutatingSpreads.get(
+            value.place.identifier.id,
+          );
+          if (spread != null) {
+            candidateNonMutatingSpreads.set(lvalue.identifier.id, spread);
+          }
+          break;
+        }
+        case 'StoreLocal': {
+          const spread = candidateNonMutatingSpreads.get(
+            value.value.identifier.id,
+          );
+          if (spread != null) {
+            candidateNonMutatingSpreads.set(lvalue.identifier.id, spread);
+            candidateNonMutatingSpreads.set(
+              value.lvalue.place.identifier.id,
+              spread,
+            );
+          }
+          break;
+        }
+        case 'JsxFragment':
+        case 'JsxExpression': {
+          // Passing objects created with spread to jsx can't mutate them
+          break;
+        }
+        case 'PropertyLoad': {
+          // Properties must be frozen since the original value was frozen
+          break;
+        }
+        case 'CallExpression':
+        case 'MethodCall': {
+          const callee =
+            value.kind === 'CallExpression' ? value.callee : value.property;
+          if (getHookKind(fn.env, callee.identifier) != null) {
+            // Hook calls have frozen arguments, and non-ref returns are frozen
+            if (!isRefOrRefValue(lvalue.identifier)) {
+              knownFrozen.add(lvalue.identifier.id);
+            }
+          } else {
+            // Non-hook calls check their operands, since they are potentially mutable
+            if (candidateNonMutatingSpreads.size !== 0) {
+              // Otherwise any reference to the spread object itself may mutate
+              for (const operand of eachInstructionValueOperand(value)) {
+                const spread = candidateNonMutatingSpreads.get(
+                  operand.identifier.id,
+                );
+                if (spread != null) {
+                  candidateNonMutatingSpreads.delete(spread);
+                }
+              }
+            }
+          }
+          break;
+        }
+        default: {
+          if (candidateNonMutatingSpreads.size !== 0) {
+            // Otherwise any reference to the spread object itself may mutate
+            for (const operand of eachInstructionValueOperand(value)) {
+              const spread = candidateNonMutatingSpreads.get(
+                operand.identifier.id,
+              );
+              if (spread != null) {
+                candidateNonMutatingSpreads.delete(spread);
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
+  const nonMutatingSpreads = new Set<IdentifierId>();
+  for (const [key, value] of candidateNonMutatingSpreads) {
+    if (key === value) {
+      nonMutatingSpreads.add(key);
+    }
+  }
+  return nonMutatingSpreads;
+}
+
 function inferParam(
   param: Place | SpreadPattern,
   initialState: InferenceState,
@@ -2054,7 +2214,9 @@ function computeSignatureForInstruction(
             kind: 'Create',
             into: place,
             reason: ValueReason.Other,
-            value: ValueKind.Mutable,
+            value: context.nonMutatingSpreads.has(place.identifier.id)
+              ? ValueKind.Frozen
+              : ValueKind.Mutable,
           });
           effects.push({
             kind: 'Capture',

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nonmutated-spread-hook-return.expect.md

@@ -0,0 +1,63 @@
+
+## Input
+
+```javascript
+import {identity, Stringify, useIdentity} from 'shared-runtime';
+
+function Component(props) {
+  const {x, ...rest} = useIdentity(props);
+  const z = rest.z;
+  identity(z);
+  return <Stringify x={x} z={z} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{x: 'Hello', z: 'World'}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { identity, Stringify, useIdentity } from "shared-runtime";
+
+function Component(props) {
+  const $ = _c(6);
+  const t0 = useIdentity(props);
+  let rest;
+  let x;
+  if ($[0] !== t0) {
+    ({ x, ...rest } = t0);
+    $[0] = t0;
+    $[1] = rest;
+    $[2] = x;
+  } else {
+    rest = $[1];
+    x = $[2];
+  }
+  const z = rest.z;
+  identity(z);
+  let t1;
+  if ($[3] !== x || $[4] !== z) {
+    t1 = <Stringify x={x} z={z} />;
+    $[3] = x;
+    $[4] = z;
+    $[5] = t1;
+  } else {
+    t1 = $[5];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: "Hello", z: "World" }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"x":"Hello","z":"World"}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nonmutated-spread-hook-return.js

@@ -0,0 +1,13 @@
+import {identity, Stringify, useIdentity} from 'shared-runtime';
+
+function Component(props) {
+  const {x, ...rest} = useIdentity(props);
+  const z = rest.z;
+  identity(z);
+  return <Stringify x={x} z={z} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{x: 'Hello', z: 'World'}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nonmutated-spread-props-jsx.expect.md

@@ -0,0 +1,57 @@
+
+## Input
+
+```javascript
+import {identity, Stringify} from 'shared-runtime';
+
+function Component({x, ...rest}) {
+  return <Stringify {...rest} x={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{x: 'Hello', z: 'World'}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { identity, Stringify } from "shared-runtime";
+
+function Component(t0) {
+  const $ = _c(6);
+  let rest;
+  let x;
+  if ($[0] !== t0) {
+    ({ x, ...rest } = t0);
+    $[0] = t0;
+    $[1] = rest;
+    $[2] = x;
+  } else {
+    rest = $[1];
+    x = $[2];
+  }
+  let t1;
+  if ($[3] !== rest || $[4] !== x) {
+    t1 = <Stringify {...rest} x={x} />;
+    $[3] = rest;
+    $[4] = x;
+    $[5] = t1;
+  } else {
+    t1 = $[5];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ x: "Hello", z: "World" }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"z":"World","x":"Hello"}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/nonmutated-spread-props-jsx.js

@@ -0,0 +1,10 @@
+import {identity, Stringify} from 'shared-runtime';
+
+function Component({x, ...rest}) {
+  return <Stringify {...rest} x={x} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{x: 'Hello', z: 'World'}],
+};