[compiler] Allow setState after await in effect callbacks

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

@@ -24,6 +24,7 @@ import {
   Place,
   Effect,
   BlockId,
+  InstructionId,
 } from '../HIR';
 import {
   eachInstructionLValue,
@@ -193,6 +194,7 @@ function getSetStateCall(
   const enableAllowSetStateFromRefsInEffects =
     env.config.enableAllowSetStateFromRefsInEffects;
   const refDerivedValues: Set<IdentifierId> = new Set();
+  const postAwaitInstructions = collectPostAwaitInstructions(fn);
 
   const isDerivedFromRef = (place: Place): boolean => {
     return (
@@ -316,6 +318,9 @@ function getSetStateCall(
             isSetStateType(callee.identifier) ||
             setStateFunctions.has(callee.identifier.id)
           ) {
+            if (postAwaitInstructions.has(instr.id)) {
+              break;
+            }
             if (enableAllowSetStateFromRefsInEffects) {
               const arg = instr.value.args.at(0);
               if (
@@ -345,3 +350,59 @@ function getSetStateCall(
   }
   return null;
 }
+
+function collectPostAwaitInstructions(fn: HIRFunction): Set<InstructionId> {
+  const postAwaitInstructions = new Set<InstructionId>();
+  if (!fn.async) {
+    return postAwaitInstructions;
+  }
+
+  const startAfterAwait: Map<BlockId, boolean> = new Map();
+  for (const [id] of fn.body.blocks) {
+    startAfterAwait.set(id, id !== fn.body.entry);
+  }
+  startAfterAwait.set(fn.body.entry, false);
+
+  let changed = true;
+  while (changed) {
+    changed = false;
+    const endAfterAwait: Map<BlockId, boolean> = new Map();
+
+    for (const [id, block] of fn.body.blocks) {
+      let afterAwait = startAfterAwait.get(id) ?? false;
+      for (const instr of block.instructions) {
+        if (instr.value.kind === 'Await') {
+          afterAwait = true;
+        }
+      }
+      endAfterAwait.set(id, afterAwait);
+    }
+
+    for (const [id, block] of fn.body.blocks) {
+      if (id === fn.body.entry) {
+        continue;
+      }
+      const startsAfterAwait =
+        block.preds.size > 0 &&
+        [...block.preds].every(pred => endAfterAwait.get(pred) === true);
+      if ((startAfterAwait.get(id) ?? false) !== startsAfterAwait) {
+        startAfterAwait.set(id, startsAfterAwait);
+        changed = true;
+      }
+    }
+  }
+
+  for (const [id, block] of fn.body.blocks) {
+    let afterAwait = startAfterAwait.get(id) ?? false;
+    for (const instr of block.instructions) {
+      if (afterAwait) {
+        postAwaitInstructions.add(instr.id);
+      }
+      if (instr.value.kind === 'Await') {
+        afterAwait = true;
+      }
+    }
+  }
+
+  return postAwaitInstructions;
+}

packages/eslint-plugin-react-hooks/__tests__/ReactCompilerRuleTypescript-test.ts

@@ -195,7 +195,66 @@ const tests: CompilerTestCases = {
   ],
 };
 
+const setStateInEffectTests: CompilerTestCases = {
+  valid: [
+    {
+      name: 'Allows setState after an await in an effect-called async callback',
+      filename: 'test.tsx',
+      code: normalizeIndent`
+        import {useCallback, useEffect, useState} from 'react';
+
+        function Component() {
+          const [ready, setReady] = useState(false);
+          const load = useCallback(async () => {
+            await fetch('/data');
+            setReady(true);
+          }, []);
+
+          useEffect(() => {
+            load();
+          }, [load]);
+
+          return <div>{ready ? 'Ready' : 'Loading'}</div>;
+        }
+      `,
+    },
+  ],
+  invalid: [
+    {
+      name: 'Reports setState before an await in an effect-called async callback',
+      filename: 'test.tsx',
+      code: normalizeIndent`
+        import {useCallback, useEffect, useState} from 'react';
+
+        function Component() {
+          const [ready, setReady] = useState(false);
+          const load = useCallback(async () => {
+            setReady(true);
+            await fetch('/data');
+          }, []);
+
+          useEffect(() => {
+            load();
+          }, [load]);
+
+          return <div>{ready ? 'Ready' : 'Loading'}</div>;
+        }
+      `,
+      errors: [
+        {
+          message: /Avoid calling setState\(\) directly within an effect/,
+        },
+      ],
+    },
+  ],
+};
+
 const eslintTester = new ESLintTesterV8({
   parser: require.resolve('@typescript-eslint/parser-v5'),
 });
 eslintTester.run('react-compiler', allRules['immutability'].rule, tests);
+eslintTester.run(
+  'react-compiler set-state-in-effect',
+  allRules['set-state-in-effect'].rule,
+  setStateInEffectTests,
+);