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

@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/Validation/ValidateNoSetStateInEffects.ts

@@ -17,6 +17,8 @@ import {
   isUseEffectHookType,
   isUseInsertionEffectHookType,
   isUseLayoutEffectHookType,
+  isUseRefType,
+  isRefValueType,
   Place,
 } from '../HIR';
 import {eachInstructionValueOperand} from '../HIR/visitors';
@@ -130,6 +132,7 @@ function getSetStateCall(
   fn: HIRFunction,
   setStateFunctions: Map<IdentifierId, Place>,
 ): Place | null {
+  const refDerivedValues: Set<IdentifierId> = new Set();
   for (const [, block] of fn.body.blocks) {
     for (const instr of block.instructions) {
       switch (instr.value.kind) {
@@ -140,6 +143,9 @@ function getSetStateCall(
               instr.value.place,
             );
           }
+          if (refDerivedValues.has(instr.value.place.identifier.id)) {
+            refDerivedValues.add(instr.lvalue.identifier.id);
+          }
           break;
         }
         case 'StoreLocal': {
@@ -153,6 +159,37 @@ function getSetStateCall(
               instr.value.value,
             );
           }
+          if (refDerivedValues.has(instr.value.value.identifier.id)) {
+            refDerivedValues.add(instr.value.lvalue.place.identifier.id);
+            refDerivedValues.add(instr.lvalue.identifier.id);
+          }
+          break;
+        }
+        case 'PropertyLoad': {
+          if (
+            instr.value.property === 'current' &&
+            (isUseRefType(instr.value.object.identifier) ||
+              isRefValueType(instr.value.object.identifier))
+          ) {
+            refDerivedValues.add(instr.lvalue.identifier.id);
+          } else if (refDerivedValues.has(instr.value.object.identifier.id)) {
+            refDerivedValues.add(instr.lvalue.identifier.id);
+          }
+          break;
+        }
+        case 'Destructure': {
+          if (refDerivedValues.has(instr.value.value.identifier.id)) {
+            for (const place of eachInstructionValueOperand(instr.value)) {
+              refDerivedValues.add(place.identifier.id);
+            }
+            refDerivedValues.add(instr.lvalue.identifier.id);
+          }
+          break;
+        }
+        case 'MethodCall': {
+          if (refDerivedValues.has(instr.value.receiver.identifier.id)) {
+            refDerivedValues.add(instr.lvalue.identifier.id);
+          }
           break;
         }
         case 'CallExpression': {
@@ -161,6 +198,18 @@ function getSetStateCall(
             isSetStateType(callee.identifier) ||
             setStateFunctions.has(callee.identifier.id)
           ) {
+            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.
+               */
+              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-useLayoutEffect-from-ref.expect.md

@@ -0,0 +1,62 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInEffects
+import { useState, useRef, useLayoutEffect } from 'react';
+
+function Tooltip() {
+  const ref = useRef(null);
+  const [tooltipHeight, setTooltipHeight] = useState(0);
+
+  useLayoutEffect(() => {
+    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, useLayoutEffect } 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];
+  }
+  useLayoutEffect(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-setState-in-useLayoutEffect-from-ref.js

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