[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

@@ -20,6 +20,8 @@ import {
   isUseStateType,
   BasicBlock,
   isUseRefType,
+  GeneratedSource,
+  SourceLocation,
 } from '../HIR';
 import {eachInstructionLValue, eachInstructionOperand} from '../HIR/visitors';
 import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables';
@@ -36,6 +38,9 @@ type DerivationMetadata = {
 type ValidationContext = {
   readonly functions: Map<IdentifierId, FunctionExpression>;
   readonly derivationCache: Map<IdentifierId, DerivationMetadata>;
+  readonly setStateCache: Map<string | undefined | null, Array<Place>>;
+  readonly effectSetStateCache: Map<string | undefined | null, Array<Place>>;
+  readonly effects: Array<HIRFunction>;
   readonly errors: CompilerError;
 };
 
@@ -69,9 +74,19 @@ export function validateNoDerivedComputationsInEffects(fn: HIRFunction): void {
 
   const errors = new CompilerError();
 
+  const setStateCache: Map<string | undefined | null, Array<Place>> = new Map();
+  const effectSetStateCache: Map<
+    string | undefined | null,
+    Array<Place>
+  > = new Map();
+  const effects: Array<HIRFunction> = [];
+
   const context: ValidationContext = {
     functions,
     derivationCache,
+    setStateCache,
+    effectSetStateCache,
+    effects,
     errors,
   };
 
@@ -109,6 +124,10 @@ export function validateNoDerivedComputationsInEffects(fn: HIRFunction): void {
     }
   }
 
+  for (const effect of effects) {
+    validateEffect(effect, context);
+  }
+
   if (errors.hasAnyErrors()) {
     throw errors;
   }
@@ -228,7 +247,7 @@ function recordInstructionDerivations(
     ) {
       const effectFunction = context.functions.get(value.args[0].identifier.id);
       if (effectFunction != null) {
-        validateEffect(effectFunction.loweredFunc.func, context);
+        context.effects.push(effectFunction.loweredFunc.func);
       }
     } else if (isUseStateType(lvalue.identifier) && value.args.length > 0) {
       const stateValueSource = value.args[0];
@@ -240,6 +259,14 @@ function recordInstructionDerivations(
   }
 
   for (const operand of eachInstructionOperand(instr)) {
+    if (isSetStateType(operand.identifier) && operand.loc !== GeneratedSource) {
+      if (context.setStateCache.has(operand.loc.identifierName)) {
+        context.setStateCache.get(operand.loc.identifierName)!.push(operand);
+      } else {
+        context.setStateCache.set(operand.loc.identifierName, [operand]);
+      }
+    }
+
     const operandMetadata = context.derivationCache.get(operand.identifier.id);
 
     if (operandMetadata === undefined) {
@@ -316,6 +343,7 @@ function validateEffect(
 
   const effectDerivedSetStateCalls: Array<{
     value: CallExpression;
+    loc: SourceLocation;
     sourceIds: Set<IdentifierId>;
   }> = [];
 
@@ -334,6 +362,23 @@ function validateEffect(
         return;
       }
 
+      for (const operand of eachInstructionOperand(instr)) {
+        if (
+          isSetStateType(operand.identifier) &&
+          operand.loc !== GeneratedSource
+        ) {
+          if (context.effectSetStateCache.has(operand.loc.identifierName)) {
+            context.effectSetStateCache
+              .get(operand.loc.identifierName)!
+              .push(operand);
+          } else {
+            context.effectSetStateCache.set(operand.loc.identifierName, [
+              operand,
+            ]);
+          }
+        }
+      }
+
       if (
         instr.value.kind === 'CallExpression' &&
         isSetStateType(instr.value.callee.identifier) &&
@@ -347,6 +392,7 @@ function validateEffect(
         if (argMetadata !== undefined) {
           effectDerivedSetStateCalls.push({
             value: instr.value,
+            loc: instr.value.callee.loc,
             sourceIds: argMetadata.sourcesIds,
           });
         }
@@ -379,13 +425,24 @@ function validateEffect(
   }
 
   for (const derivedSetStateCall of effectDerivedSetStateCalls) {
-    context.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 &&
+      context.effectSetStateCache.has(derivedSetStateCall.loc.identifierName) &&
+      context.setStateCache.has(derivedSetStateCall.loc.identifierName) &&
+      context.effectSetStateCache.get(derivedSetStateCall.loc.identifierName)!
+        .length ===
+        context.setStateCache.get(derivedSetStateCall.loc.identifierName)!
+          .length -
+          1
+    ) {
+      context.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>