[compiler] Option to treat "set-" prefixed callees as setState functions (#34505)

Calling setState functions during render can lead to extraneous renders
or even infinite loops. We also have runtime detection for loops, but
static detection is obviously even better.

This PR adds an option to infer identifers as setState functions if both
the following conditions are met:
- The identifier is named starting with "set"
- The identifier is used as the callee of a call expression

By inferring values as SetState type, this allows our existing
ValidateNoSetStateInRender rule to flag calls during render, disallowing
examples like the following:

```js
function Component({setParentState}) {
  setParentState(...);
  ^^^^^^^^^^^^^^ Error: Cannot call setState in render
}
```

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

@@ -621,6 +621,13 @@ export const EnvironmentConfigSchema = z.object({
    */
   enableTreatRefLikeIdentifiersAsRefs: z.boolean().default(true),
 
+  /**
+   * Treat identifiers as SetState type if both
+   * - they are named with a "set-" prefix
+   * - they are called somewhere
+   */
+  enableTreatSetIdentifiersAsStateSetters: z.boolean().default(false),
+
   /*
    * If specified a value, the compiler lowers any calls to `useContext` to use
    * this value as the callee.

compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts

@@ -31,6 +31,7 @@ import {
   BuiltInObjectId,
   BuiltInPropsId,
   BuiltInRefValueId,
+  BuiltInSetStateId,
   BuiltInUseRefId,
 } from '../HIR/ObjectShape';
 import {eachInstructionLValue, eachInstructionOperand} from '../HIR/visitors';
@@ -276,9 +277,16 @@ function* generateInstructionTypes(
        * We should change Hook to a subtype of Function or change unifier logic.
        * (see https://github.com/facebook/react-forget/pull/1427)
        */
+      let shapeId: string | null = null;
+      if (env.config.enableTreatSetIdentifiersAsStateSetters) {
+        const name = getName(names, value.callee.identifier.id);
+        if (name.startsWith('set')) {
+          shapeId = BuiltInSetStateId;
+        }
+      }
       yield equation(value.callee.identifier.type, {
         kind: 'Function',
-        shapeId: null,
+        shapeId,
         return: returnType,
         isConstructor: false,
       });

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-unconditional-set-state-hook-return-in-render.expect.md

@@ -0,0 +1,55 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInRender @enableTreatSetIdentifiersAsStateSetters
+function Component() {
+  const [state, setState] = useCustomState(0);
+  const aliased = setState;
+
+  setState(1);
+  aliased(2);
+
+  return state;
+}
+
+function useCustomState(init) {
+  return useState(init);
+}
+
+```
+
+
+## Error
+
+```
+Found 2 errors:
+
+Error: Calling setState during render may trigger an infinite loop
+
+Calling setState during render will trigger another render, and can lead to infinite loops. (https://react.dev/reference/react/useState).
+
+error.invalid-unconditional-set-state-hook-return-in-render.ts:6:2
+  4 |   const aliased = setState;
+  5 |
+> 6 |   setState(1);
+    |   ^^^^^^^^ Found setState() in render
+  7 |   aliased(2);
+  8 |
+  9 |   return state;
+
+Error: Calling setState during render may trigger an infinite loop
+
+Calling setState during render will trigger another render, and can lead to infinite loops. (https://react.dev/reference/react/useState).
+
+error.invalid-unconditional-set-state-hook-return-in-render.ts:7:2
+   5 |
+   6 |   setState(1);
+>  7 |   aliased(2);
+     |   ^^^^^^^ Found setState() in render
+   8 |
+   9 |   return state;
+  10 | }
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-unconditional-set-state-hook-return-in-render.js

@@ -0,0 +1,14 @@
+// @validateNoSetStateInRender @enableTreatSetIdentifiersAsStateSetters
+function Component() {
+  const [state, setState] = useCustomState(0);
+  const aliased = setState;
+
+  setState(1);
+  aliased(2);
+
+  return state;
+}
+
+function useCustomState(init) {
+  return useState(init);
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-unconditional-set-state-prop-in-render.expect.md

@@ -0,0 +1,50 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInRender @enableTreatSetIdentifiersAsStateSetters
+function Component({setX}) {
+  const aliased = setX;
+
+  setX(1);
+  aliased(2);
+
+  return x;
+}
+
+```
+
+
+## Error
+
+```
+Found 2 errors:
+
+Error: Calling setState during render may trigger an infinite loop
+
+Calling setState during render will trigger another render, and can lead to infinite loops. (https://react.dev/reference/react/useState).
+
+error.invalid-unconditional-set-state-prop-in-render.ts:5:2
+  3 |   const aliased = setX;
+  4 |
+> 5 |   setX(1);
+    |   ^^^^ Found setState() in render
+  6 |   aliased(2);
+  7 |
+  8 |   return x;
+
+Error: Calling setState during render may trigger an infinite loop
+
+Calling setState during render will trigger another render, and can lead to infinite loops. (https://react.dev/reference/react/useState).
+
+error.invalid-unconditional-set-state-prop-in-render.ts:6:2
+  4 |
+  5 |   setX(1);
+> 6 |   aliased(2);
+    |   ^^^^^^^ Found setState() in render
+  7 |
+  8 |   return x;
+  9 | }
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-unconditional-set-state-prop-in-render.js

@@ -0,0 +1,9 @@
+// @validateNoSetStateInRender @enableTreatSetIdentifiersAsStateSetters
+function Component({setX}) {
+  const aliased = setX;
+
+  setX(1);
+  aliased(2);
+
+  return x;
+}