[compiler] Switch to track setStates by aliasing and id instead of identifier names

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -277,7 +277,7 @@ function runWithEnvironment(
     }
 
     if (env.config.validateNoDerivedComputationsInEffects_exp) {
-      validateNoDerivedComputationsInEffects_exp(hir);
+      env.logErrors(validateNoDerivedComputationsInEffects_exp(hir));
     }
 
     if (env.config.validateNoSetStateInEffects) {

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

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import {Result} from '../Utils/Result';
 import {CompilerDiagnostic, CompilerError, Effect} from '..';
 import {ErrorCategory} from '../CompilerError';
 import {
@@ -20,7 +21,6 @@ import {
   isUseStateType,
   BasicBlock,
   isUseRefType,
-  GeneratedSource,
   SourceLocation,
 } from '../HIR';
 import {eachInstructionLValue, eachInstructionOperand} from '../HIR/visitors';
@@ -41,8 +41,8 @@ type ValidationContext = {
   readonly errors: CompilerError;
   readonly derivationCache: DerivationCache;
   readonly effects: Set<HIRFunction>;
-  readonly setStateCache: Map<string | undefined | null, Array<Place>>;
-  readonly effectSetStateCache: Map<string | undefined | null, Array<Place>>;
+  readonly setStateLoads: Map<IdentifierId, IdentifierId | null>;
+  readonly setStateUsages: Map<IdentifierId, Set<SourceLocation>>;
 };
 
 class DerivationCache {
@@ -173,25 +173,22 @@ function isNamedIdentifier(place: Place): place is Place & {
  */
 export function validateNoDerivedComputationsInEffects_exp(
   fn: HIRFunction,
-): void {
+): Result<void, CompilerError> {
   const functions: Map<IdentifierId, FunctionExpression> = new Map();
   const derivationCache = new DerivationCache();
   const errors = new CompilerError();
   const effects: Set<HIRFunction> = new Set();
 
-  const setStateCache: Map<string | undefined | null, Array<Place>> = new Map();
-  const effectSetStateCache: Map<
-    string | undefined | null,
-    Array<Place>
-  > = new Map();
+  const setStateLoads: Map<IdentifierId, IdentifierId> = new Map();
+  const setStateUsages: Map<IdentifierId, Set<SourceLocation>> = new Map();
 
   const context: ValidationContext = {
     functions,
     errors,
     derivationCache,
     effects,
-    setStateCache,
-    effectSetStateCache,
+    setStateLoads,
+    setStateUsages,
   };
 
   if (fn.fnType === 'Hook') {
@@ -236,9 +233,7 @@ export function validateNoDerivedComputationsInEffects_exp(
     validateEffect(effect, context);
   }
 
-  if (errors.hasAnyErrors()) {
-    throw errors;
-  }
+  return errors.asResult();
 }
 
 function recordPhiDerivations(
@@ -282,11 +277,61 @@ function joinValue(
   return 'fromPropsAndState';
 }
 
+function getRootSetState(
+  key: IdentifierId,
+  loads: Map<IdentifierId, IdentifierId | null>,
+  visited: Set<IdentifierId> = new Set(),
+): IdentifierId | null {
+  if (visited.has(key)) {
+    return null;
+  }
+  visited.add(key);
+
+  const parentId = loads.get(key);
+
+  if (parentId === undefined) {
+    return null;
+  }
+
+  if (parentId === null) {
+    return key;
+  }
+
+  return getRootSetState(parentId, loads, visited);
+}
+
+function maybeRecordSetState(
+  instr: Instruction,
+  loads: Map<IdentifierId, IdentifierId | null>,
+  usages: Map<IdentifierId, Set<SourceLocation>>,
+): void {
+  for (const operand of eachInstructionLValue(instr)) {
+    if (
+      instr.value.kind === 'LoadLocal' &&
+      loads.has(instr.value.place.identifier.id)
+    ) {
+      loads.set(operand.identifier.id, instr.value.place.identifier.id);
+    } else {
+      if (isSetStateType(operand.identifier)) {
+        // this is a root setState
+        loads.set(operand.identifier.id, null);
+      }
+    }
+
+    const rootSetState = getRootSetState(operand.identifier.id, loads);
+    if (rootSetState !== null && usages.get(rootSetState) === undefined) {
+      usages.set(rootSetState, new Set([operand.loc]));
+    }
+  }
+}
+
 function recordInstructionDerivations(
   instr: Instruction,
   context: ValidationContext,
   isFirstPass: boolean,
 ): void {
+  maybeRecordSetState(instr, context.setStateLoads, context.setStateUsages);
+
   let typeOfValue: TypeOfValue = 'ignored';
   let isSource: boolean = false;
   const sources: Set<IdentifierId> = new Set();
@@ -319,15 +364,13 @@ function recordInstructionDerivations(
   }
 
   for (const operand of eachInstructionOperand(instr)) {
-    if (
-      isSetStateType(operand.identifier) &&
-      operand.loc !== GeneratedSource &&
-      isFirstPass
-    ) {
-      if (context.setStateCache.has(operand.loc.identifierName)) {
-        context.setStateCache.get(operand.loc.identifierName)!.push(operand);
-      } else {
-        context.setStateCache.set(operand.loc.identifierName, [operand]);
+    if (context.setStateLoads.has(operand.identifier.id)) {
+      const rootSetStateId = getRootSetState(
+        operand.identifier.id,
+        context.setStateLoads,
+      );
+      if (rootSetStateId !== null) {
+        context.setStateUsages.get(rootSetStateId)?.add(operand.loc);
       }
     }
 
@@ -533,11 +576,16 @@ function validateEffect(
 
   const effectDerivedSetStateCalls: Array<{
     value: CallExpression;
-    loc: SourceLocation;
+    id: IdentifierId;
     sourceIds: Set<IdentifierId>;
     typeOfValue: TypeOfValue;
   }> = [];
 
+  const effectSetStateUsages: Map<
+    IdentifierId,
+    Set<SourceLocation>
+  > = new Map();
+
   const globals: Set<IdentifierId> = new Set();
   for (const block of effectFunction.body.blocks.values()) {
     for (const pred of block.preds) {
@@ -553,19 +601,16 @@ function validateEffect(
         return;
       }
 
+      maybeRecordSetState(instr, context.setStateLoads, effectSetStateUsages);
+
       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 (context.setStateLoads.has(operand.identifier.id)) {
+          const rootSetStateId = getRootSetState(
+            operand.identifier.id,
+            context.setStateLoads,
+          );
+          if (rootSetStateId !== null) {
+            effectSetStateUsages.get(rootSetStateId)?.add(operand.loc);
           }
         }
       }
@@ -583,7 +628,7 @@ function validateEffect(
         if (argMetadata !== undefined) {
           effectDerivedSetStateCalls.push({
             value: instr.value,
-            loc: instr.value.callee.loc,
+            id: instr.value.callee.identifier.id,
             sourceIds: argMetadata.sourcesIds,
             typeOfValue: argMetadata.typeOfValue,
           });
@@ -617,15 +662,17 @@ function validateEffect(
   }
 
   for (const derivedSetStateCall of effectDerivedSetStateCalls) {
+    const rootSetStateCall = getRootSetState(
+      derivedSetStateCall.id,
+      context.setStateLoads,
+    );
+
     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
+      rootSetStateCall !== null &&
+      effectSetStateUsages.has(rootSetStateCall) &&
+      context.setStateUsages.has(rootSetStateCall) &&
+      effectSetStateUsages.get(rootSetStateCall)!.size ===
+        context.setStateUsages.get(rootSetStateCall)!.size - 1
     ) {
       const propsSet = new Set<string>();
       const stateSet = new Set<string>();

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/effect-derived-computations/derived-state-conditionally-in-effect.expect.md

@@ -0,0 +1,86 @@
+
+## Input
+
+```javascript
+// @validateNoDerivedComputationsInEffects_exp @loggerTestOnly
+import {useEffect, useState} from 'react';
+
+function Component({value, enabled}) {
+  const [localValue, setLocalValue] = useState('');
+
+  useEffect(() => {
+    if (enabled) {
+      setLocalValue(value);
+    } else {
+      setLocalValue('disabled');
+    }
+  }, [value, enabled]);
+
+  return <div>{localValue}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{value: 'test', enabled: true}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @validateNoDerivedComputationsInEffects_exp @loggerTestOnly
+import { useEffect, useState } from "react";
+
+function Component(t0) {
+  const $ = _c(6);
+  const { value, enabled } = t0;
+  const [localValue, setLocalValue] = useState("");
+  let t1;
+  let t2;
+  if ($[0] !== enabled || $[1] !== value) {
+    t1 = () => {
+      if (enabled) {
+        setLocalValue(value);
+      } else {
+        setLocalValue("disabled");
+      }
+    };
+
+    t2 = [value, enabled];
+    $[0] = enabled;
+    $[1] = value;
+    $[2] = t1;
+    $[3] = t2;
+  } else {
+    t1 = $[2];
+    t2 = $[3];
+  }
+  useEffect(t1, t2);
+  let t3;
+  if ($[4] !== localValue) {
+    t3 = <div>{localValue}</div>;
+    $[4] = localValue;
+    $[5] = t3;
+  } else {
+    t3 = $[5];
+  }
+  return t3;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: "test", enabled: true }],
+};
+
+```
+
+## Logs
+
+```
+{"kind":"CompileError","detail":{"options":{"description":"Using an effect triggers an additional render which can hurt performance and user experience, potentially briefly showing stale values to the user\n\nThis setState call is setting a derived value that depends on the following reactive sources:\n\nProps: [value]\n\nData Flow Tree:\n└── value (Prop)\n\nSee: https://react.dev/learn/you-might-not-need-an-effect#updating-state-based-on-props-or-state","category":"EffectDerivationsOfState","reason":"You might not need an effect. Derive values in render, not effects.","details":[{"kind":"error","loc":{"start":{"line":9,"column":6,"index":244},"end":{"line":9,"column":19,"index":257},"filename":"derived-state-conditionally-in-effect.ts","identifierName":"setLocalValue"},"message":"This should be computed during render, not in an effect"}]}},"fnLoc":null}
+{"kind":"CompileSuccess","fnLoc":{"start":{"line":4,"column":0,"index":107},"end":{"line":16,"column":1,"index":378},"filename":"derived-state-conditionally-in-effect.ts"},"fnName":"Component","memoSlots":6,"memoBlocks":2,"memoValues":3,"prunedMemoBlocks":0,"prunedMemoValues":0}
+```
+
+### Eval output
+(kind: ok) <div>test</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/effect-derived-computations/error.derived-state-conditionally-in-effect.js → compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/effect-derived-computations/derived-state-conditionally-in-effect.js

@@ -1,4 +1,4 @@
-// @validateNoDerivedComputationsInEffects_exp
+// @validateNoDerivedComputationsInEffects_exp @loggerTestOnly
 import {useEffect, useState} from 'react';
 
 function Component({value, enabled}) {