Skip to content

[compiler] Improve ShallowMutable detection for frozen sources #34341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 1 commit into from
+845 −7
Closed

[compiler] Improve ShallowMutable detection for frozen sources #34341

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1445,6 +1445,7 @@ export enum ValueKind {
Primitive = 'primitive',
Global = 'global',
Mutable = 'mutable',
ShallowMutable = 'shallowmutable',
Copy link
Member

@josephsavona josephsavona Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

you'll need to update mergeValueKinds() in InferMutationAliasingEffects to account for the new type

Context = 'context',
}

Expand All @@ -1454,6 +1455,7 @@ export const ValueKindSchema = z.enum([
ValueKind.Primitive,
ValueKind.Global,
ValueKind.Mutable,
ValueKind.ShallowMutable,
ValueKind.Context,
]);

Expand Down
44 changes: 37 additions & 7 deletions compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutationAliasingEffects.ts
View file
Open in desktop
Copy link
Member

@josephsavona josephsavona Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I must be overlooking things — where do we actually create new instances that are ShallowMutable?

Copy link
Member

@josephsavona josephsavona Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yeah, the bug from the PR description still occurs on the latest version of the PR, and the fixtures don't seem different than before the change. We don't actually create ShallowMutable instances anywhere afaict

Copy link
Member

@josephsavona josephsavona Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe i'm looking too early, this is just exciting!!!!

Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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,
});
Expand Down Expand Up @@ -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;
}
Expand All @@ -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(
Expand Down Expand Up @@ -781,6 +791,7 @@ function applyEffect(
const fromValue = state.kind(effect.from);
const fromKind = fromValue.kind;
switch (fromKind) {
case ValueKind.ShallowMutable:
Copy link
Member

@josephsavona josephsavona Sep 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm this doesn't seem right, we should keep this as an Assign. The new rule for shallow mutable should be[1]:

Given
  x is ShallowMutable
  CreateFrom y <- x
Then
  y is Frozen

Whereas this basically says

Given
  x is ShallowMutable
  Assign y = x
Then
  y is Frozen (bc immutable capture)

[1] would be nice to update MUTABILITY_ALIASING_MODEL for this, btw

case ValueKind.Frozen: {
applyEffect(
context,
Expand Down Expand Up @@ -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) {
Expand Down Expand Up @@ -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,
});
}
Comment on lines +1338 to +1343
Copy link
Member

@josephsavona josephsavona Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks right. See my comment elsewhere about updating mergeValueKinds() to account for ShallowMutability. The logic there should basically be to treat ShallowMutable like Mutable when it joins with other things (if Mutable | x => Mutable, then ShallowMutable | x => ShallowMutable). Except ShallowMutable | Mutable => Mutable.

Copy link
Member

@josephsavona josephsavona Sep 3, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The fact that you didn't hit the invariant in mergeValueKinds() also suggests we need some more tests. Something like

function Component(props) {
  let x;
  if (props.cond) {
    x = [props.item]; // shallow mutable
  } else {
    x = []; // regular mutable
  }
  const z = x[0]; // z is mutable
}

}
}

switch (variant) {
case 'MutateConditionally':
case 'MutateTransitiveConditionally': {
switch (kind) {
case ValueKind.Mutable:
case ValueKind.Context: {
case ValueKind.Context:
case ValueKind.ShallowMutable: {
return 'mutate';
}
default: {
Expand All @@ -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: {
Expand Down
91 changes: 91 additions & 0 deletions ...s__/fixtures/compiler/preserve-memo-validation/array-spread-from-hook.expect.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,91 @@

## 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>
22 changes: 22 additions & 0 deletions ...piler/src/__tests__/fixtures/compiler/preserve-memo-validation/array-spread-from-hook.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
};
120 changes: 120 additions & 0 deletions ...xtures/compiler/preserve-memo-validation/destructured-with-rest-props.expect.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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>
Loading
Loading