[Compiler] Don't throw calculate in render if the blamed setter is used outside of the effect

Summary:
If the setter is used both inside and outside the effect then usually the solution is more complex and requires hoisting state up to a parent component since we can't just remove the local state.

To do this, we now have 2 caches that track setState usages (not just calls) since if the effect is passed as an argument or called outside the effect the solution gets more complex which we are trying to avoid for now

Author: jorge-cab

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

@@ -19,6 +19,8 @@ import {
   Instruction,
   isUseStateType,
   isUseRefType,
+  GeneratedSource,
+  SourceLocation,
 } from '../HIR';
 import {eachInstructionLValue, eachInstructionOperand} from '../HIR/visitors';
 import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables';
@@ -59,6 +61,10 @@ export function validateNoDerivedComputationsInEffects(fn: HIRFunction): void {
   const functions: Map<IdentifierId, FunctionExpression> = new Map();
 
   const derivationCache: Map<IdentifierId, DerivationMetadata> = new Map();
+  const setStateCache: Map<string | undefined | null, Array<Place>> = new Map();
+
+  const effects: Array<HIRFunction> = [];
+
   if (fn.fnType === 'Hook') {
     for (const param of fn.params) {
       if (param.kind === 'Identifier') {
@@ -127,11 +133,7 @@ export function validateNoDerivedComputationsInEffects(fn: HIRFunction): void {
           ) {
             const effectFunction = functions.get(value.args[0].identifier.id);
             if (effectFunction != null) {
-              validateEffect(
-                effectFunction.loweredFunc.func,
-                errors,
-                derivationCache,
-              );
+              effects.push(effectFunction.loweredFunc.func);
             }
           } else if (
             isUseStateType(lvalue.identifier) &&
@@ -146,6 +148,25 @@ export function validateNoDerivedComputationsInEffects(fn: HIRFunction): void {
         }
 
         for (const operand of eachInstructionOperand(instr)) {
+          // Record setState usages everywhere
+          switch (instr.value.kind) {
+            case 'JsxExpression':
+            case 'CallExpression':
+            case 'MethodCall':
+              if (
+                isSetStateType(operand.identifier) &&
+                operand.loc !== GeneratedSource
+              ) {
+                if (setStateCache.has(operand.loc.identifierName)) {
+                  setStateCache.get(operand.loc.identifierName)!.push(operand);
+                } else {
+                  setStateCache.set(operand.loc.identifierName, [operand]);
+                }
+              }
+              break;
+            default:
+          }
+
           const operandMetadata = derivationCache.get(operand.identifier.id);
 
           if (operandMetadata === undefined) {
@@ -214,6 +235,10 @@ export function validateNoDerivedComputationsInEffects(fn: HIRFunction): void {
     }
   }
 
+  for (const effect of effects) {
+    validateEffect(effect, errors, derivationCache, setStateCache);
+  }
+
   if (errors.hasAnyErrors()) {
     throw errors;
   }
@@ -275,11 +300,17 @@ function validateEffect(
   effectFunction: HIRFunction,
   errors: CompilerError,
   derivationCache: Map<IdentifierId, DerivationMetadata>,
+  setStateCache: Map<string | undefined | null, Array<Place>>,
 ): void {
+  const effectSetStateCache: Map<
+    string | undefined | null,
+    Array<Place>
+  > = new Map();
   const seenBlocks: Set<BlockId> = new Set();
 
   const effectDerivedSetStateCalls: Array<{
     value: CallExpression;
+    loc: SourceLocation;
     sourceIds: Set<IdentifierId>;
   }> = [];
 
@@ -298,6 +329,28 @@ function validateEffect(
         return;
       }
 
+      for (const operand of eachInstructionOperand(instr)) {
+        switch (instr.value.kind) {
+          case 'JsxExpression':
+          case 'CallExpression':
+          case 'MethodCall':
+            if (
+              isSetStateType(operand.identifier) &&
+              operand.loc !== GeneratedSource
+            ) {
+              if (effectSetStateCache.has(operand.loc.identifierName)) {
+                effectSetStateCache
+                  .get(operand.loc.identifierName)!
+                  .push(operand);
+              } else {
+                effectSetStateCache.set(operand.loc.identifierName, [operand]);
+              }
+            }
+            break;
+          default:
+        }
+      }
+
       if (
         instr.value.kind === 'CallExpression' &&
         isSetStateType(instr.value.callee.identifier) &&
@@ -311,6 +364,7 @@ function validateEffect(
         if (argMetadata !== undefined) {
           effectDerivedSetStateCalls.push({
             value: instr.value,
+            loc: instr.value.callee.loc,
             sourceIds: argMetadata.sourcesIds,
           });
         }
@@ -343,13 +397,22 @@ function validateEffect(
   }
 
   for (const derivedSetStateCall of effectDerivedSetStateCalls) {
-    errors.push({
-      category: ErrorCategory.EffectDerivationsOfState,
-      reason:
-        'Values derived from props and state should be calculated during render, not in an effect. (https://react.dev/learn/you-might-not-need-an-effect#updating-state-based-on-props-or-state)',
-      description: null,
-      loc: derivedSetStateCall.value.callee.loc,
-      suggestions: null,
-    });
+    if (
+      derivedSetStateCall.loc !== GeneratedSource &&
+      effectSetStateCache.has(derivedSetStateCall.loc.identifierName) &&
+      setStateCache.has(derivedSetStateCall.loc.identifierName) &&
+      effectSetStateCache.get(derivedSetStateCall.loc.identifierName)!
+        .length ===
+        setStateCache.get(derivedSetStateCall.loc.identifierName)!.length
+    ) {
+      errors.push({
+        category: ErrorCategory.EffectDerivationsOfState,
+        reason:
+          'Values derived from props and state should be calculated during render, not in an effect. (https://react.dev/learn/you-might-not-need-an-effect#updating-state-based-on-props-or-state)',
+        description: null,
+        loc: derivedSetStateCall.value.callee.loc,
+        suggestions: null,
+      });
+    }
   }
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/effect-derived-computations/derived-state-from-prop-setter-call-outside-effect-no-error.expect.md

@@ -0,0 +1,84 @@
+
+## Input
+
+```javascript
+// @validateNoDerivedComputationsInEffects
+import {useEffect, useState} from 'react';
+
+function Component({initialName}) {
+  const [name, setName] = useState('');
+
+  useEffect(() => {
+    setName(initialName);
+  }, [initialName]);
+
+  return (
+    <div>
+      <input value={name} onChange={e => setName(e.target.value)} />
+    </div>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{initialName: 'John'}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateNoDerivedComputationsInEffects
+import { useEffect, useState } from "react";
+
+function Component(t0) {
+  const $ = _c(6);
+  const { initialName } = t0;
+  const [name, setName] = useState("");
+  let t1;
+  let t2;
+  if ($[0] !== initialName) {
+    t1 = () => {
+      setName(initialName);
+    };
+    t2 = [initialName];
+    $[0] = initialName;
+    $[1] = t1;
+    $[2] = t2;
+  } else {
+    t1 = $[1];
+    t2 = $[2];
+  }
+  useEffect(t1, t2);
+  let t3;
+  if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
+    t3 = (e) => setName(e.target.value);
+    $[3] = t3;
+  } else {
+    t3 = $[3];
+  }
+  let t4;
+  if ($[4] !== name) {
+    t4 = (
+      <div>
+        <input value={name} onChange={t3} />
+      </div>
+    );
+    $[4] = name;
+    $[5] = t4;
+  } else {
+    t4 = $[5];
+  }
+  return t4;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ initialName: "John" }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div><input value="John"></div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/effect-derived-computations/error.derived-state-from-prop-setter-call-outside-effect-no-error.js renamed to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/effect-derived-computations/derived-state-from-prop-setter-call-outside-effect-no-error.js

@@ -0,0 +1,85 @@
+
+## Input
+
+```javascript
+// @validateNoDerivedComputationsInEffects
+import {useEffect, useState} from 'react';
+
+function MockComponent({onSet}) {
+  return <div onClick={() => onSet('clicked')}>Mock Component</div>;
+}
+
+function Component({propValue}) {
+  const [value, setValue] = useState(null);
+  useEffect(() => {
+    setValue(propValue);
+  }, [propValue]);
+
+  return <MockComponent onSet={setValue} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{propValue: 'test'}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateNoDerivedComputationsInEffects
+import { useEffect, useState } from "react";
+
+function MockComponent(t0) {
+  const $ = _c(2);
+  const { onSet } = t0;
+  let t1;
+  if ($[0] !== onSet) {
+    t1 = <div onClick={() => onSet("clicked")}>Mock Component</div>;
+    $[0] = onSet;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  return t1;
+}
+
+function Component(t0) {
+  const $ = _c(4);
+  const { propValue } = t0;
+  const [, setValue] = useState(null);
+  let t1;
+  let t2;
+  if ($[0] !== propValue) {
+    t1 = () => {
+      setValue(propValue);
+    };
+    t2 = [propValue];
+    $[0] = propValue;
+    $[1] = t1;
+    $[2] = t2;
+  } else {
+    t1 = $[1];
+    t2 = $[2];
+  }
+  useEffect(t1, t2);
+  let t3;
+  if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
+    t3 = <MockComponent onSet={setValue} />;
+    $[3] = t3;
+  } else {
+    t3 = $[3];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ propValue: "test" }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>Mock Component</div>