[compiler] Add enableUseKeyedState flag and update setState validation

This commit adds a new compiler environment flag `enableUseKeyedState`, and
updates ValidateNoSetStateInRender to no longer throw on setState calls with primitive values. When `enableUseKeyedState` is enabled, calling setState with a primitive instead recommends using `useKeyedState` instead.

Author: poteto

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

@@ -318,6 +318,12 @@ export const EnvironmentConfigSchema = z.object({
    */
   validateNoSetStateInRender: z.boolean().default(true),
 
+  /**
+   * When enabled, changes the behavior of validateNoSetStateInRender to recommend
+   * using useKeyedState instead of calling setState directly in render.
+   */
+  enableUseKeyedState: z.boolean().default(false),
+
   /**
    * Validates that setState is not called synchronously within an effect (useEffect and friends).
    * Scheduling a setState (with an event listener, subscription, etc) is valid.

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

@@ -10,11 +10,116 @@ import {
   CompilerError,
   ErrorCategory,
 } from '../CompilerError';
-import {HIRFunction, IdentifierId, isSetStateType} from '../HIR';
+import {
+  HIRFunction,
+  Identifier,
+  IdentifierId,
+  Instruction,
+  InstructionValue,
+  isPrimitiveType,
+  isSetStateType,
+  Phi,
+  Place,
+  SpreadPattern,
+} from '../HIR';
 import {computeUnconditionalBlocks} from '../HIR/ComputeUnconditionalBlocks';
 import {eachInstructionValueOperand} from '../HIR/visitors';
 import {Result} from '../Utils/Result';
 
+function isPrimitiveSetArg(
+  arg: Place | SpreadPattern,
+  fn: HIRFunction,
+): boolean {
+  if (arg.kind !== 'Identifier') {
+    return false;
+  }
+
+  const visited = new Set<IdentifierId>();
+  const defs = buildDefinitionMap(fn);
+  return isPrimitiveIdentifier(arg.identifier, defs, visited);
+}
+
+type DefinitionMap = Map<IdentifierId, Instruction | Phi>;
+
+function buildDefinitionMap(fn: HIRFunction): DefinitionMap {
+  const defs: DefinitionMap = new Map();
+
+  for (const [, block] of fn.body.blocks) {
+    for (const phi of block.phis) {
+      defs.set(phi.place.identifier.id, phi);
+    }
+    for (const instr of block.instructions) {
+      defs.set(instr.lvalue.identifier.id, instr);
+    }
+  }
+
+  return defs;
+}
+
+function isPrimitiveIdentifier(
+  identifier: Identifier,
+  defs: DefinitionMap,
+  visited: Set<IdentifierId>,
+): boolean {
+  if (isPrimitiveType(identifier)) {
+    return true;
+  }
+
+  if (visited.has(identifier.id)) {
+    return false;
+  }
+  visited.add(identifier.id);
+
+  const def = defs.get(identifier.id);
+  if (def == null) {
+    return false;
+  }
+
+  if (isPhi(def)) {
+    return Array.from(def.operands.values()).every(operand =>
+      isPrimitiveIdentifier(operand.identifier, defs, visited),
+    );
+  }
+
+  return isPrimitiveInstruction(def.value, defs, visited);
+}
+
+function isPhi(def: Instruction | Phi): def is Phi {
+  return 'kind' in def && def.kind === 'Phi';
+}
+
+function isPrimitiveInstruction(
+  value: InstructionValue,
+  defs: DefinitionMap,
+  visited: Set<IdentifierId>,
+): boolean {
+  switch (value.kind) {
+    case 'Primitive':
+    case 'TemplateLiteral':
+    case 'JSXText':
+    case 'UnaryExpression':
+    case 'BinaryExpression':
+      return true;
+
+    case 'TypeCastExpression':
+      return isPrimitiveIdentifier(value.value.identifier, defs, visited);
+
+    case 'LoadLocal':
+    case 'LoadContext':
+      return isPrimitiveIdentifier(value.place.identifier, defs, visited);
+
+    case 'StoreLocal':
+    case 'StoreContext':
+      return isPrimitiveIdentifier(value.value.identifier, defs, visited);
+
+    case 'Await':
+      return isPrimitiveIdentifier(value.value.identifier, defs, visited);
+
+    default:
+      return false;
+  }
+}
+
 /**
  * Validates that the given function does not have an infinite update loop
  * caused by unconditionally calling setState during render. This validation
@@ -55,6 +160,7 @@ function validateNoSetStateInRenderImpl(
   unconditionalSetStateFunctions: Set<IdentifierId>,
 ): Result<void, CompilerError> {
   const unconditionalBlocks = computeUnconditionalBlocks(fn);
+  const enableUseKeyedState = fn.env.config.enableUseKeyedState;
   let activeManualMemoId: number | null = null;
   const errors = new CompilerError();
   for (const [, block] of fn.body.blocks) {
@@ -155,20 +261,48 @@ function validateNoSetStateInRenderImpl(
                 }),
               );
             } else if (unconditionalBlocks.has(block.id)) {
-              errors.pushDiagnostic(
-                CompilerDiagnostic.create({
-                  category: ErrorCategory.RenderSetState,
-                  reason:
-                    'Calling setState during render may trigger an infinite loop',
-                  description:
-                    'Calling setState during render will trigger another render, and can lead to infinite loops. (https://react.dev/reference/react/useState)',
-                  suggestions: null,
-                }).withDetails({
-                  kind: 'error',
-                  loc: callee.loc,
-                  message: 'Found setState() in render',
-                }),
-              );
+              let isArgPrimitive = false;
+
+              if (instr.value.args.length > 0) {
+                const arg = instr.value.args[0];
+                if (arg.kind === 'Identifier') {
+                  isArgPrimitive = isPrimitiveSetArg(arg, fn);
+                }
+              }
+
+              if (isArgPrimitive) {
+                if (enableUseKeyedState) {
+                  errors.pushDiagnostic(
+                    CompilerDiagnostic.create({
+                      category: ErrorCategory.RenderSetState,
+                      reason:
+                        'Calling setState during render may trigger an infinite loop',
+                      description:
+                        'Use useKeyedState instead of calling setState directly in render. Example: const [value, setValue] = useKeyedState(initialValue, key)',
+                      suggestions: null,
+                    }).withDetails({
+                      kind: 'error',
+                      loc: callee.loc,
+                      message: 'Found setState() in render',
+                    }),
+                  );
+                }
+              } else {
+                errors.pushDiagnostic(
+                  CompilerDiagnostic.create({
+                    category: ErrorCategory.RenderSetState,
+                    reason:
+                      'Calling setState during render may trigger an infinite loop',
+                    description:
+                      'Calling setState during render will trigger another render, and can lead to infinite loops. (https://react.dev/reference/react/useState)',
+                    suggestions: null,
+                  }).withDetails({
+                    kind: 'error',
+                    loc: callee.loc,
+                    message: 'Found setState() in render',
+                  }),
+                );
+              }
             }
           }
           break;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setState-in-render-unbound-state.expect.md

@@ -0,0 +1,42 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInRender @enableUseKeyedState
+import {useState} from 'react';
+
+function Component() {
+  const [total, setTotal] = useState(0);
+  setTotal(42);
+  return total;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+
+## Error
+
+```
+Found 1 error:
+
+Error: Calling setState during render may trigger an infinite loop
+
+Use useKeyedState instead of calling setState directly in render. Example: const [value, setValue] = useKeyedState(initialValue, key).
+
+error.invalid-setstate-enabled-use-keyed-state.ts:6:2
+  4 | function Component() {
+  5 |   const [total, setTotal] = useState(0);
+> 6 |   setTotal(42);
+    |   ^^^^^^^^ Found setState() in render
+  7 |   return total;
+  8 | }
+  9 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setstate-enabled-use-keyed-state.expect.md

@@ -0,0 +1,14 @@
+// @validateNoSetStateInRender @enableUseKeyedState
+import {useState} from 'react';
+
+function Component() {
+  const [total, setTotal] = useState(0);
+  setTotal(42);
+  return total;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setstate-enabled-use-keyed-state.js

@@ -0,0 +1,42 @@
+
+## Input
+
+```javascript
+// @validateNoSetStateInRender
+import {useState} from 'react';
+
+function Component() {
+  const [total, setTotal] = useState(0);
+  setTotal({count: 42});
+  return total;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};
+
+```
+
+
+## Error
+
+```
+Found 1 error:
+
+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-setstate-object-no-keyed-state.ts:6:2
+  4 | function Component() {
+  5 |   const [total, setTotal] = useState(0);
+> 6 |   setTotal({count: 42});
+    |   ^^^^^^^^ Found setState() in render
+  7 |   return total;
+  8 | }
+  9 |
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-setstate-object-no-keyed-state.expect.md

@@ -0,0 +1,14 @@
+// @validateNoSetStateInRender
+import {useState} from 'react';
+
+function Component() {
+  const [total, setTotal] = useState(0);
+  setTotal({count: 42});
+  return total;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [],
+  isComponent: true,
+};