[compiler] Improve ShallowMutable detection for frozen sources

The compiler currently fails to preserve manual memoization when components use rest spread destructuring for props:

```js
// OK
function Component(props) {
  const value = useMemo(() => compute(props.foo), [props.foo]);
}

// Manual memo could not be preserved
function Component({...props}) {
  const value = useMemo(() => compute(props.foo), [props.foo]);
}
```

The issue is that property accesses on rest-spread objects are treated as fully mutable.

This PR introduces a new `ShallowMutable` ValueKind for rest spreads that come specifically from props. When we access properties from a `ShallowMutable` value, they're treated as Frozen.

Closes #34313

Author: poteto

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

@@ -1445,6 +1445,7 @@ export enum ValueKind {
   Primitive = 'primitive',
   Global = 'global',
   Mutable = 'mutable',
+  ShallowMutable = 'shallowmutable',
   Context = 'context',
 }
 
@@ -1454,6 +1455,7 @@ export const ValueKindSchema = z.enum([
   ValueKind.Primitive,
   ValueKind.Global,
   ValueKind.Mutable,
+  ValueKind.ShallowMutable,
   ValueKind.Context,
 ]);
 

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

@@ -591,8 +591,14 @@ function applyEffect(
         };
         context.effectInstructionValueCache.set(effect, value);
       }
+
+      const outputKind =
+        fromValue.kind === ValueKind.ShallowMutable
+          ? ValueKind.Frozen
+          : fromValue.kind;
+
       state.initialize(value, {
-        kind: fromValue.kind,
+        kind: outputKind,
         reason: new Set(fromValue.reason),
       });
       state.define(effect.into, value);
@@ -607,10 +613,11 @@ function applyEffect(
           });
           break;
         }
+        case ValueKind.ShallowMutable:
         case ValueKind.Frozen: {
           effects.push({
             kind: 'Create',
-            value: fromValue.kind,
+            value: outputKind,
             into: effect.into,
             reason: [...fromValue.reason][0] ?? ValueReason.Other,
           });
@@ -720,11 +727,14 @@ function applyEffect(
        * copy-on-write semantics, then we can prune the effect
        */
       const intoKind = state.kind(effect.into).kind;
+      const fromKind = state.kind(effect.from).kind;
+
       let isMutableDesination: boolean;
       switch (intoKind) {
         case ValueKind.Context:
         case ValueKind.Mutable:
-        case ValueKind.MaybeFrozen: {
+        case ValueKind.MaybeFrozen:
+        case ValueKind.ShallowMutable: {
           isMutableDesination = true;
           break;
         }
@@ -733,14 +743,14 @@ function applyEffect(
           break;
         }
       }
-      const fromKind = state.kind(effect.from).kind;
       let isMutableReferenceType: boolean;
       switch (fromKind) {
         case ValueKind.Global:
         case ValueKind.Primitive: {
           isMutableReferenceType = false;
           break;
         }
+        case ValueKind.ShallowMutable:
         case ValueKind.Frozen: {
           isMutableReferenceType = false;
           applyEffect(
@@ -781,6 +791,7 @@ function applyEffect(
       const fromValue = state.kind(effect.from);
       const fromKind = fromValue.kind;
       switch (fromKind) {
+        case ValueKind.ShallowMutable:
         case ValueKind.Frozen: {
           applyEffect(
             context,
@@ -1267,6 +1278,7 @@ class InferenceState {
     switch (value.kind) {
       case ValueKind.Context:
       case ValueKind.Mutable:
+      case ValueKind.ShallowMutable:
       case ValueKind.MaybeFrozen: {
         const values = this.values(place);
         for (const instrValue of values) {
@@ -1315,13 +1327,30 @@ class InferenceState {
     if (isRefOrRefValue(place.identifier)) {
       return 'mutate-ref';
     }
-    const kind = this.kind(place).kind;
+    const abstractValue = this.kind(place);
+    const kind = abstractValue.kind;
+
+    // Downgrade ShallowMutable to Mutable when mutated
+    if (kind === ValueKind.ShallowMutable) {
+      const values = this.values(place);
+      for (const value of values) {
+        const valueInfo = this.#values.get(value);
+        if (valueInfo && valueInfo.kind === ValueKind.ShallowMutable) {
+          this.#values.set(value, {
+            kind: ValueKind.Mutable,
+            reason: valueInfo.reason,
+          });
+        }
+      }
+    }
+
     switch (variant) {
       case 'MutateConditionally':
       case 'MutateTransitiveConditionally': {
         switch (kind) {
           case ValueKind.Mutable:
-          case ValueKind.Context: {
+          case ValueKind.Context:
+          case ValueKind.ShallowMutable: {
             return 'mutate';
           }
           default: {
@@ -1333,7 +1362,8 @@ class InferenceState {
       case 'MutateTransitive': {
         switch (kind) {
           case ValueKind.Mutable:
-          case ValueKind.Context: {
+          case ValueKind.Context:
+          case ValueKind.ShallowMutable: {
             return 'mutate';
           }
           case ValueKind.Primitive: {
@@ -2349,6 +2379,7 @@ function areArgumentsImmutableAndNonMutating(
     const kind = state.kind(place).kind;
     switch (kind) {
       case ValueKind.Primitive:
+      case ValueKind.ShallowMutable:
       case ValueKind.Frozen: {
         /*
          * Only immutable values, or frozen lambdas are allowed.

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/array-spread-from-hook.expect.md

@@ -0,0 +1,90 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {ValidateMemoization} from 'shared-runtime';
+
+function useData() {
+  return ['a', 'b', 'c'];
+}
+
+function Component() {
+  const [first, ...rest] = useData();
+  
+  const result = useMemo(() => {
+    return rest.join('-');
+  }, [rest]);
+  
+  return <ValidateMemoization inputs={[rest]} output={result} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+import { ValidateMemoization } from "shared-runtime";
+
+function useData() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = ["a", "b", "c"];
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+function Component() {
+  const $ = _c(7);
+  const t0 = useData();
+  let rest;
+  if ($[0] !== t0) {
+    [, ...rest] = t0;
+    $[0] = t0;
+    $[1] = rest;
+  } else {
+    rest = $[1];
+  }
+
+  const result = rest.join("-");
+  let t1;
+  if ($[2] !== rest) {
+    t1 = [rest];
+    $[2] = rest;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  let t2;
+  if ($[4] !== result || $[5] !== t1) {
+    t2 = <ValidateMemoization inputs={t1} output={result} />;
+    $[4] = result;
+    $[5] = t1;
+    $[6] = t2;
+  } else {
+    t2 = $[6];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[["b","c"]],"output":"b-c"}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/array-spread-from-hook.tsx

@@ -0,0 +1,22 @@
+import {useMemo} from 'react';
+import {ValidateMemoization} from 'shared-runtime';
+
+function useData() {
+  return ['a', 'b', 'c'];
+}
+
+function Component() {
+  const [first, ...rest] = useData();
+
+  const result = useMemo(() => {
+    return rest.join('-');
+  }, [rest]);
+
+  return <ValidateMemoization inputs={[rest]} output={result} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/preserve-memo-validation/destructured-with-rest-props.expect.md

@@ -0,0 +1,120 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+
+function useTheme() {
+  return {primary: '#blue', secondary: '#green'};
+}
+
+function computeStyles(
+  specialProp: string | undefined,
+  restProps: any,
+  theme: any,
+) {
+  return {
+    color: specialProp ? theme.primary : theme.secondary,
+    ...restProps.style,
+  };
+}
+
+export function SpecialButton({
+  specialProp,
+  ...restProps
+}: {
+  specialProp?: string;
+  style?: Record<string, string>;
+  onClick?: () => void;
+}) {
+  const theme = useTheme();
+
+  const styles = useMemo(
+    () => computeStyles(specialProp, restProps, theme),
+    [specialProp, restProps, theme],
+  );
+
+  return (
+    <button style={styles} onClick={restProps.onClick}>
+      Click me
+    </button>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: SpecialButton,
+  params: [{specialProp: 'test', style: {fontSize: '16px'}, onClick: () => {}}],
+  isComponent: true,
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+
+function useTheme() {
+  const $ = _c(1);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = { primary: "#blue", secondary: "#green" };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+
+function computeStyles(specialProp, restProps, theme) {
+  const $ = _c(3);
+
+  const t0 = specialProp ? theme.primary : theme.secondary;
+  let t1;
+  if ($[0] !== restProps.style || $[1] !== t0) {
+    t1 = { color: t0, ...restProps.style };
+    $[0] = restProps.style;
+    $[1] = t0;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export function SpecialButton(t0) {
+  const $ = _c(3);
+  const { specialProp, ...restProps } = t0;
+
+  const theme = useTheme();
+
+  const styles = computeStyles(specialProp, restProps, theme);
+  let t1;
+  if ($[0] !== restProps.onClick || $[1] !== styles) {
+    t1 = (
+      <button style={styles} onClick={restProps.onClick}>
+        Click me
+      </button>
+    );
+    $[0] = restProps.onClick;
+    $[1] = styles;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: SpecialButton,
+  params: [
+    { specialProp: "test", style: { fontSize: "16px" }, onClick: () => {} },
+  ],
+  isComponent: true,
+};
+
+```
+      
+### Eval output
+(kind: ok) <button style="font-size: 16px;">Click me</button>