compiler: Add support for ref effects

ghstack-source-id: c534d4943f1433f79afdad7ed2d926589cfdea31
Pull Request resolved: #29733

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

@@ -291,6 +291,10 @@ export type FunctionEffect =
  kind: "GlobalMutation";
  error: CompilerErrorDetailOptions;
  }
+ | {
+ kind: "RefMutation";
+ error: CompilerErrorDetailOptions;
+ }
  | {
  kind: "ReactMutation";
  error: CompilerErrorDetailOptions;

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

@@ -29,6 +29,8 @@ import {
  isArrayType,
  isMutableEffect,
  isObjectType,
+ isRefValueType,
+ isUseRefType,
 } from "../HIR/HIR";
 import { FunctionSignature } from "../HIR/ObjectShape";
 import {
@@ -239,7 +241,8 @@ export default function inferReferenceEffects(
  functionEffects.forEach((eff) => {
  switch (eff.kind) {
  case "ReactMutation":
- case "GlobalMutation": {
+ case "GlobalMutation":
+ case "RefMutation": {
  CompilerError.throw(eff.error);
  }
  case "ContextMutation": {
@@ -407,7 +410,8 @@ class InferenceState {
  for (const effect of value.loweredFunc.func.effects) {
  if (
  effect.kind === "GlobalMutation" ||
- effect.kind === "ReactMutation"
+ effect.kind === "ReactMutation" ||
+ effect.kind === "RefMutation"
  ) {
  // Known effects are always propagated upwards
  functionEffects.push(effect);
@@ -539,10 +543,12 @@ class InferenceState {
  let reason = getWriteErrorReason(valueKind);
  functionEffect = {
  kind:
- valueKind.reason.size === 1 &&
- valueKind.reason.has(ValueReason.Global)
- ? "GlobalMutation"
- : "ReactMutation",
+ isRefValueType(place.identifier) || isUseRefType(place.identifier)
+ ? "RefMutation"
+ : valueKind.reason.size === 1 &&
+ valueKind.reason.has(ValueReason.Global)
+ ? "GlobalMutation"
+ : "ReactMutation",
  error: {
  reason,
  description:
@@ -578,10 +584,12 @@ class InferenceState {
  let reason = getWriteErrorReason(valueKind);
  functionEffect = {
  kind:
- valueKind.reason.size === 1 &&
- valueKind.reason.has(ValueReason.Global)
- ? "GlobalMutation"
- : "ReactMutation",
+ isRefValueType(place.identifier) || isUseRefType(place.identifier)
+ ? "RefMutation"
+ : valueKind.reason.size === 1 &&
+ valueKind.reason.has(ValueReason.Global)
+ ? "GlobalMutation"
+ : "ReactMutation",
  error: {
  reason,
  description:
@@ -1139,7 +1147,9 @@ function inferBlock(
  );
  functionEffects.push(
  ...propEffects.filter(
- (propEffect) => propEffect.kind !== "GlobalMutation"
+ (propEffect) =>
+ propEffect.kind !== "GlobalMutation" &&
+ propEffect.kind !== "RefMutation"
  )
  );
  }
@@ -1333,7 +1343,10 @@ function inferBlock(
  functionEffects.push(
  ...argumentEffects.filter(
  (argEffect) =>
- !isUseEffect || i !== 0 || argEffect.kind !== "GlobalMutation"
+ !isUseEffect ||
+ i !== 0 ||
+ (argEffect.kind !== "GlobalMutation" &&
+ argEffect.kind !== "RefMutation")
  )
  );
  hasCaptureArgument ||= place.effect === Effect.Capture;
@@ -1462,7 +1475,10 @@ function inferBlock(
  functionEffects.push(
  ...argumentEffects.filter(
  (argEffect) =>
- !isUseEffect || i !== 0 || argEffect.kind !== "GlobalMutation"
+ !isUseEffect ||
+ i !== 0 ||
+ (argEffect.kind !== "GlobalMutation" &&
+ argEffect.kind !== "RefMutation")
  )
  );
  hasCaptureArgument ||= place.effect === Effect.Capture;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-mutating-ref-in-callback-passed-to-jsx-indirect.expect.md

@@ -0,0 +1,108 @@
+
+## Input
+
+```javascript
+import { useRef } from "react";
+
+function Component() {
+ const ref = useRef(null);
+
+ const setRef = () => {
+ if (ref.current !== null) {
+ ref.current = "";
+ }
+ };
+
+ const onClick = () => {
+ setRef();
+ };
+
+ return (
+ <>
+ <input ref={ref} />
+ <button onClick={onClick} />
+ </>
+ );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useRef } from "react";
+
+function Component() {
+ const $ = _c(10);
+ const ref = useRef(null);
+ let t0;
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+ t0 = () => {
+ if (ref.current !== null) {
+ ref.current = "";
+ }
+ };
+ $[0] = t0;
+ } else {
+ t0 = $[0];
+ }
+ const setRef = t0;
+ let t1;
+ if ($[1] !== setRef) {
+ t1 = () => {
+ setRef();
+ };
+ $[1] = setRef;
+ $[2] = t1;
+ } else {
+ t1 = $[2];
+ }
+ const onClick = t1;
+ let t2;
+ if ($[3] !== ref) {
+ t2 = <input ref={ref} />;
+ $[3] = ref;
+ $[4] = t2;
+ } else {
+ t2 = $[4];
+ }
+ let t3;
+ if ($[5] !== onClick) {
+ t3 = <button onClick={onClick} />;
+ $[5] = onClick;
+ $[6] = t3;
+ } else {
+ t3 = $[6];
+ }
+ let t4;
+ if ($[7] !== t2 || $[8] !== t3) {
+ t4 = (
+ <>
+ {t2}
+ {t3}
+ </>
+ );
+ $[7] = t2;
+ $[8] = t3;
+ $[9] = t4;
+ } else {
+ t4 = $[9];
+ }
+ return t4;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{}],
+};
+
+```
+  
+### Eval output
+(kind: ok) <input><button></button>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-mutating-ref-in-callback-passed-to-jsx-indirect.tsx

@@ -0,0 +1,27 @@
+import { useRef } from "react";
+
+function Component() {
+ const ref = useRef(null);
+
+ const setRef = () => {
+ if (ref.current !== null) {
+ ref.current = "";
+ }
+ };
+
+ const onClick = () => {
+ setRef();
+ };
+
+ return (
+ <>
+ <input ref={ref} />
+ <button onClick={onClick} />
+ </>
+ );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-mutating-ref-in-callback-passed-to-jsx.expect.md

@@ -0,0 +1,93 @@
+
+## Input
+
+```javascript
+import { useRef } from "react";
+
+function Component() {
+ const ref = useRef(null);
+
+ const onClick = () => {
+ if (ref.current !== null) {
+ ref.current = "";
+ }
+ };
+
+ return (
+ <>
+ <input ref={ref} />
+ <button onClick={onClick} />
+ </>
+ );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useRef } from "react";
+
+function Component() {
+ const $ = _c(8);
+ const ref = useRef(null);
+ let t0;
+ if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+ t0 = () => {
+ if (ref.current !== null) {
+ ref.current = "";
+ }
+ };
+ $[0] = t0;
+ } else {
+ t0 = $[0];
+ }
+ const onClick = t0;
+ let t1;
+ if ($[1] !== ref) {
+ t1 = <input ref={ref} />;
+ $[1] = ref;
+ $[2] = t1;
+ } else {
+ t1 = $[2];
+ }
+ let t2;
+ if ($[3] !== onClick) {
+ t2 = <button onClick={onClick} />;
+ $[3] = onClick;
+ $[4] = t2;
+ } else {
+ t2 = $[4];
+ }
+ let t3;
+ if ($[5] !== t1 || $[6] !== t2) {
+ t3 = (
+ <>
+ {t1}
+ {t2}
+ </>
+ );
+ $[5] = t1;
+ $[6] = t2;
+ $[7] = t3;
+ } else {
+ t3 = $[7];
+ }
+ return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{}],
+};
+
+```
+  
+### Eval output
+(kind: ok) <input><button></button>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-mutating-ref-in-callback-passed-to-jsx.tsx

@@ -0,0 +1,23 @@
+import { useRef } from "react";
+
+function Component() {
+ const ref = useRef(null);
+
+ const onClick = () => {
+ if (ref.current !== null) {
+ ref.current = "";
+ }
+ };
+
+ return (
+ <>
+ <input ref={ref} />
+ <button onClick={onClick} />
+ </>
+ );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+ fn: Component,
+ params: [{}],
+};