[compiler] Allow setStates in use{Layout,Insertion}Effect where the set value is derived from a ref (facebook#34462)

@stipsan found this issue where the compiler would bailout on the
`useLayoutEffect` examples in the React docs. While setState in an
effect is typically an anti-pattern due to the fact that it hurts
performance through cascading renders, the one scenario where it _is_
allowed is if the value being set flows from a ref.

Author: poteto

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -276,7 +276,7 @@ function runWithEnvironment(
     }
 
     if (env.config.validateNoSetStateInEffects) {
-      env.logErrors(validateNoSetStateInEffects(hir));
+      env.logErrors(validateNoSetStateInEffects(hir, env));
     }
 
     if (env.config.validateNoJSXInTryStatements) {

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

@@ -660,6 +660,13 @@ export const EnvironmentConfigSchema = z.object({
    * while its parent function remains uncompiled.
    */
   validateNoDynamicallyCreatedComponentsOrHooks: z.boolean().default(false),
+
+  /**
+   * When enabled, allows setState calls in effects when the value being set is
+   * derived from a ref. This is useful for patterns where initial layout measurements
+   * from refs need to be stored in state during mount.
+   */
+  enableAllowSetStateFromRefsInEffects: z.boolean().default(true),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateNoSetStateInEffects.ts

@@ -11,16 +11,23 @@ import {
   ErrorCategory,
 } from '../CompilerError';
 import {
+  Environment,
   HIRFunction,
   IdentifierId,
   isSetStateType,
   isUseEffectHookType,
   isUseInsertionEffectHookType,
   isUseLayoutEffectHookType,
+  isUseRefType,
+  isRefValueType,
   Place,
 } from '../HIR';
-import {eachInstructionValueOperand} from '../HIR/visitors';
+import {
+  eachInstructionLValue,
+  eachInstructionValueOperand,
+} from '../HIR/visitors';
 import {Result} from '../Utils/Result';
+import {Iterable_some} from '../Utils/utils';
 
 /**
  * Validates against calling setState in the body of an effect (useEffect and friends),
@@ -32,6 +39,7 @@ import {Result} from '../Utils/Result';
  */
 export function validateNoSetStateInEffects(
   fn: HIRFunction,
+  env: Environment,
 ): Result<void, CompilerError> {
   const setStateFunctions: Map<IdentifierId, Place> = new Map();
   const errors = new CompilerError();
@@ -72,6 +80,7 @@ export function validateNoSetStateInEffects(
             const callee = getSetStateCall(
               instr.value.loweredFunc.func,
               setStateFunctions,
+              env,
             );
             if (callee !== null) {
               setStateFunctions.set(instr.lvalue.identifier.id, callee);
@@ -129,9 +138,42 @@ export function validateNoSetStateInEffects(
 function getSetStateCall(
   fn: HIRFunction,
   setStateFunctions: Map<IdentifierId, Place>,
+  env: Environment,
 ): Place | null {
+  const refDerivedValues: Set<IdentifierId> = new Set();
+
+  const isDerivedFromRef = (place: Place): boolean => {
+    return (
+      refDerivedValues.has(place.identifier.id) ||
+      isUseRefType(place.identifier) ||
+      isRefValueType(place.identifier)
+    );
+  };
+
   for (const [, block] of fn.body.blocks) {
     for (const instr of block.instructions) {
+      if (env.config.enableAllowSetStateFromRefsInEffects) {
+        const hasRefOperand = Iterable_some(
+          eachInstructionValueOperand(instr.value),
+          isDerivedFromRef,
+        );
+
+        if (hasRefOperand) {
+          for (const lvalue of eachInstructionLValue(instr)) {
+            refDerivedValues.add(lvalue.identifier.id);
+          }
+        }
+
+        if (
+          instr.value.kind === 'PropertyLoad' &&
+          instr.value.property === 'current' &&
+          (isUseRefType(instr.value.object.identifier) ||
+            isRefValueType(instr.value.object.identifier))
+        ) {
+          refDerivedValues.add(instr.lvalue.identifier.id);
+        }
+      }
+
       switch (instr.value.kind) {
         case 'LoadLocal': {
           if (setStateFunctions.has(instr.value.place.identifier.id)) {
@@ -161,6 +203,21 @@ function getSetStateCall(
             isSetStateType(callee.identifier) ||
             setStateFunctions.has(callee.identifier.id)
           ) {
+            if (env.config.enableAllowSetStateFromRefsInEffects) {
+              const arg = instr.value.args.at(0);
+              if (
+                arg !== undefined &&
+                arg.kind === 'Identifier' &&
+                refDerivedValues.has(arg.identifier.id)
+              ) {
+                /**
+                 * The one special case where we allow setStates in effects is in the very specific
+                 * scenario where the value being set is derived from a ref. For example this may
+                 * be needed when initial layout measurements from refs need to be stored in state.
+                 */
+                return null;
+              }
+            }
             /*
              * TODO: once we support multiple locations per error, we should link to the
              * original Place in the case that setStateFunction.has(callee)

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/valid-set-state-in-useEffect-from-ref.expect.md

@@ -0,0 +1,63 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInEffects
+import {useState, useRef, useEffect} from 'react';
+
+function Tooltip() {
+  const ref = useRef(null);
+  const [tooltipHeight, setTooltipHeight] = useState(0);
+
+  useEffect(() => {
+    const {height} = ref.current.getBoundingClientRect();
+    setTooltipHeight(height);
+  }, []);
+
+  return tooltipHeight;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Tooltip,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateNoSetStateInEffects
+import { useState, useRef, useEffect } from "react";
+
+function Tooltip() {
+  const $ = _c(2);
+  const ref = useRef(null);
+  const [tooltipHeight, setTooltipHeight] = useState(0);
+  let t0;
+  let t1;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = () => {
+      const { height } = ref.current.getBoundingClientRect();
+      setTooltipHeight(height);
+    };
+    t1 = [];
+    $[0] = t0;
+    $[1] = t1;
+  } else {
+    t0 = $[0];
+    t1 = $[1];
+  }
+  useEffect(t0, t1);
+  return tooltipHeight;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Tooltip,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: exception) Cannot read properties of null (reading 'getBoundingClientRect')

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/valid-set-state-in-useEffect-from-ref.js

@@ -0,0 +1,19 @@
+// @validateNoSetStateInEffects
+import {useState, useRef, useEffect} from 'react';
+
+function Tooltip() {
+  const ref = useRef(null);
+  const [tooltipHeight, setTooltipHeight] = useState(0);
+
+  useEffect(() => {
+    const {height} = ref.current.getBoundingClientRect();
+    setTooltipHeight(height);
+  }, []);
+
+  return tooltipHeight;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Tooltip,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/valid-setState-in-effect-from-ref-arithmetic.expect.md

@@ -0,0 +1,68 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects
+import {useState, useRef, useLayoutEffect} from 'react';
+
+function Component() {
+  const ref = useRef({size: 5});
+  const [computedSize, setComputedSize] = useState(0);
+
+  useLayoutEffect(() => {
+    setComputedSize(ref.current.size * 10);
+  }, []);
+
+  return computedSize;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects
+import { useState, useRef, useLayoutEffect } from "react";
+
+function Component() {
+  const $ = _c(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = { size: 5 };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const ref = useRef(t0);
+  const [computedSize, setComputedSize] = useState(0);
+  let t1;
+  let t2;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = () => {
+      setComputedSize(ref.current.size * 10);
+    };
+    t2 = [];
+    $[1] = t1;
+    $[2] = t2;
+  } else {
+    t1 = $[1];
+    t2 = $[2];
+  }
+  useLayoutEffect(t1, t2);
+  return computedSize;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) 50

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/valid-setState-in-effect-from-ref-arithmetic.js

@@ -0,0 +1,18 @@
+// @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects
+import {useState, useRef, useLayoutEffect} from 'react';
+
+function Component() {
+  const ref = useRef({size: 5});
+  const [computedSize, setComputedSize] = useState(0);
+
+  useLayoutEffect(() => {
+    setComputedSize(ref.current.size * 10);
+  }, []);
+
+  return computedSize;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/valid-setState-in-effect-from-ref-array-index.expect.md

@@ -0,0 +1,69 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects
+import {useState, useRef, useEffect} from 'react';
+
+function Component() {
+  const ref = useRef([1, 2, 3, 4, 5]);
+  const [value, setValue] = useState(0);
+
+  useEffect(() => {
+    const index = 2;
+    setValue(ref.current[index]);
+  }, []);
+
+  return value;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects
+import { useState, useRef, useEffect } from "react";
+
+function Component() {
+  const $ = _c(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = [1, 2, 3, 4, 5];
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const ref = useRef(t0);
+  const [value, setValue] = useState(0);
+  let t1;
+  let t2;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = () => {
+      setValue(ref.current[2]);
+    };
+    t2 = [];
+    $[1] = t1;
+    $[2] = t2;
+  } else {
+    t1 = $[1];
+    t2 = $[2];
+  }
+  useEffect(t1, t2);
+  return value;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) 3

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/valid-setState-in-effect-from-ref-array-index.js

@@ -0,0 +1,19 @@
+// @validateNoSetStateInEffects @enableAllowSetStateFromRefsInEffects
+import {useState, useRef, useEffect} from 'react';
+
+function Component() {
+  const ref = useRef([1, 2, 3, 4, 5]);
+  const [value, setValue] = useState(0);
+
+  useEffect(() => {
+    const index = 2;
+    setValue(ref.current[index]);
+  }, []);
+
+  return value;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+};