[compiler] proof of concept to validate relay derived setStates

Summary:
Experimenting with this a bit. Not a serious change yet

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

@@ -39,6 +39,7 @@ import {TypeConfig} from './TypeSchema';
  * The React team is open to collaborating with library authors to help develop compatible versions of these APIs,
  * and we have already reached out to the teams who own any API listed here to ensure they are aware of the issue.
  */
+
 export function defaultModuleTypeProvider(
   moduleName: string,
 ): TypeConfig | null {

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

@@ -41,6 +41,7 @@ import {
 } from './HIR';
 import {
   BuiltInMixedReadonlyId,
+  BuiltInUseFragmentId,
   DefaultMutatingHook,
   DefaultNonmutatingHook,
   FunctionSignature,
@@ -648,6 +649,12 @@ export const EnvironmentConfigSchema = z.object({
    */
   enableTreatSetIdentifiersAsStateSetters: z.boolean().default(false),
 
+  /**
+   * Treat hook calls named "useFragment" as returning BuiltInUseFragment type.
+   * This enables tracking of Relay fragment-derived values through the program.
+   */
+  enableTreatUseFragmentAsRelay: z.boolean().default(false),
+
   /*
    * If specified a value, the compiler lowers any calls to `useContext` to use
    * this value as the callee.
@@ -819,14 +826,22 @@ export class Environment {
         ],
         suggestions: null,
       });
+      // Use BuiltInUseFragmentId for useFragment to enable tracking of fragment-derived values
+      const returnTypeShapeId =
+        hookName === 'useFragment'
+          ? BuiltInUseFragmentId
+          : hook.transitiveMixedData
+            ? BuiltInMixedReadonlyId
+            : null;
       this.#globals.set(
         hookName,
         addHook(this.#shapes, {
           positionalParams: [],
           restParam: hook.effectKind,
-          returnType: hook.transitiveMixedData
-            ? {kind: 'Object', shapeId: BuiltInMixedReadonlyId}
-            : {kind: 'Poly'},
+          returnType:
+            returnTypeShapeId != null
+              ? {kind: 'Object', shapeId: returnTypeShapeId}
+              : {kind: 'Poly'},
           returnValueKind: hook.valueKind,
           calleeEffect: Effect.Read,
           hookKind: 'Custom',

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

@@ -1886,6 +1886,10 @@ export function isUseStateType(id: Identifier): boolean {
   return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInUseState';
 }
 
+export function isUseFragmentType(id: Identifier): boolean {
+  return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInUseFragment';
+}
+
 export function isJsxType(type: Type): boolean {
   return type.kind === 'Object' && type.shapeId === 'BuiltInJsx';
 }

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

@@ -392,6 +392,7 @@ export const BuiltInSetActionStateId = 'BuiltInSetActionState';
 export const BuiltInUseRefId = 'BuiltInUseRefId';
 export const BuiltInRefValueId = 'BuiltInRefValue';
 export const BuiltInMixedReadonlyId = 'BuiltInMixedReadonly';
+export const BuiltInUseFragmentId = 'BuiltInUseFragment';
 export const BuiltInUseEffectHookId = 'BuiltInUseEffectHook';
 export const BuiltInUseLayoutEffectHookId = 'BuiltInUseLayoutEffectHook';
 export const BuiltInUseInsertionEffectHookId = 'BuiltInUseInsertionEffectHook';
@@ -1475,6 +1476,16 @@ addObject(BUILTIN_SHAPES, BuiltInMixedReadonlyId, [
   ['*', {kind: 'Object', shapeId: BuiltInMixedReadonlyId}],
 ]);
 
+/**
+ * BuiltInUseFragment represents values returned from Relay's useFragment hook.
+ * The catch-all property ensures that property accesses on useFragment data
+ * are also typed as BuiltInUseFragmentId, enabling tracking of fragment-derived
+ * values through the program.
+ */
+addObject(BUILTIN_SHAPES, BuiltInUseFragmentId, [
+  ['*', {kind: 'Object', shapeId: BuiltInUseFragmentId}],
+]);
+
 addObject(BUILTIN_SHAPES, BuiltInJsxId, []);
 addObject(BUILTIN_SHAPES, BuiltInFunctionId, []);
 

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

@@ -33,6 +33,7 @@ import {
   BuiltInPropsId,
   BuiltInRefValueId,
   BuiltInSetStateId,
+  BuiltInUseFragmentId,
   BuiltInUseRefId,
 } from '../HIR/ObjectShape';
 import {eachInstructionLValue, eachInstructionOperand} from '../HIR/visitors';
@@ -265,7 +266,18 @@ function* generateInstructionTypes(
 
     case 'LoadGlobal': {
       const globalType = env.getGlobalDeclaration(value.binding, value.loc);
-      if (globalType) {
+      // If the binding is useFragment and the flag is enabled, override the type
+      if (
+        env.config.enableTreatUseFragmentAsRelay &&
+        value.binding.name === 'useFragment'
+      ) {
+        yield equation(left, {
+          kind: 'Function',
+          shapeId: BuiltInUseFragmentId,
+          return: {kind: 'Object', shapeId: BuiltInUseFragmentId},
+          isConstructor: false,
+        });
+      } else if (globalType) {
         yield equation(left, globalType);
       }
       break;
@@ -279,9 +291,9 @@ function* generateInstructionTypes(
        * (see https://github.com/facebook/react-forget/pull/1427)
        */
       let shapeId: string | null = null;
+      const calleeName = getName(names, value.callee.identifier.id);
       if (env.config.enableTreatSetIdentifiersAsStateSetters) {
-        const name = getName(names, value.callee.identifier.id);
-        if (name.startsWith('set')) {
+        if (calleeName.startsWith('set')) {
           shapeId = BuiltInSetStateId;
         }
       }
@@ -730,26 +742,32 @@ class Unifier {
          *   T2
          *   Phi [
          *     T3
-         *     Phi [
-         *       T4
-         *     ]
+         *     Phi [T4]
          *   ]
          * ]
          *
-         * Which avoids the cycle
+         * Which then resolves to T1, T2, T3, T4 as candidates (because
+         * they're all resolved Phi types).
          */
-        const operands = [];
+        const operands: Array<Type> = [];
         for (const operand of type.operands) {
-          if (operand.kind === 'Type' && operand.id === v.id) {
+          if (typeEquals(operand, v)) {
             continue;
           }
-          const resolved = this.tryResolveType(v, operand);
+          const resolved = this.tryResolveType(v, this.get(operand));
           if (resolved === null) {
             return null;
           }
           operands.push(resolved);
         }
-        return {kind: 'Phi', operands};
+        if (operands.length === 0) {
+          // All operands were `v`
+          return null;
+        }
+        return {
+          kind: 'Phi',
+          operands,
+        };
       }
       case 'Type': {
         const substitution = this.get(type);
@@ -859,12 +877,19 @@ function tryUnionTypes(ty1: Type, ty2: Type): Type | null {
   } else if (ty2.kind === 'Object' && ty2.shapeId === BuiltInMixedReadonlyId) {
     readonlyType = ty2;
     otherType = ty1;
+  } else if (ty1.kind === 'Object' && ty1.shapeId === BuiltInUseFragmentId) {
+    readonlyType = ty1;
+    otherType = ty2;
+  } else if (ty2.kind === 'Object' && ty2.shapeId === BuiltInUseFragmentId) {
+    readonlyType = ty2;
+    otherType = ty1;
   } else {
     return null;
   }
   if (otherType.kind === 'Primitive') {
     /**
      * Union(Primitive | MixedReadonly) = MixedReadonly
+     * Union(Primitive | UseFragment) = UseFragment
      *
      * For example, `data ?? null` could return `data`, the fact that RHS
      * is a primitive doesn't guarantee the result is a primitive.
@@ -876,6 +901,7 @@ function tryUnionTypes(ty1: Type, ty2: Type): Type | null {
   ) {
     /**
      * Union(Array | MixedReadonly) = Array
+     * Union(Array | UseFragment) = Array
      *
      * In practice this pattern means the result is always an array. Given
      * that this behavior requires opting-in to the mixedreadonly type

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

@@ -15,6 +15,7 @@ import {
   IdentifierId,
   isSetStateType,
   isUseEffectHookType,
+  isUseFragmentType,
   Place,
   CallExpression,
   Instruction,
@@ -27,8 +28,14 @@ import {
 import {eachInstructionLValue, eachInstructionOperand} from '../HIR/visitors';
 import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables';
 import {assertExhaustive} from '../Utils/utils';
+import {printInstruction} from '../HIR/PrintHIR';
 
-type TypeOfValue = 'ignored' | 'fromProps' | 'fromState' | 'fromPropsAndState';
+type TypeOfValue =
+  | 'ignored'
+  | 'fromProps'
+  | 'fromState'
+  | 'fromPropsAndState'
+  | 'fromRelay';
 
 type DerivationMetadata = {
   typeOfValue: TypeOfValue;
@@ -299,6 +306,9 @@ function joinValue(
   if (lvalueType === 'ignored') return valueType;
   if (valueType === 'ignored') return lvalueType;
   if (lvalueType === valueType) return lvalueType;
+  if (lvalueType === 'fromRelay' || valueType === 'fromRelay') {
+    return 'fromRelay';
+  }
   return 'fromPropsAndState';
 }
 
@@ -397,6 +407,15 @@ function recordInstructionDerivations(
         true,
       );
       return;
+    } else if (isUseFragmentType(lvalue.identifier)) {
+      typeOfValue = 'fromRelay';
+      context.derivationCache.addDerivationEntry(
+        lvalue,
+        new Set(),
+        typeOfValue,
+        true,
+      );
+      return;
     }
   } else if (value.kind === 'ArrayExpression') {
     context.candidateDependencies.set(lvalue.identifier.id, value);
@@ -511,6 +530,13 @@ type TreeNode = {
   children: Array<TreeNode>;
 };
 
+type DataFlowInfo = {
+  rootSources: string;
+  trees: Array<string>;
+  propsArr: Array<string>;
+  stateArr: Array<string>;
+};
+
 function buildTreeNode(
   sourceId: IdentifierId,
   context: ValidationContext,
@@ -628,6 +654,51 @@ function renderTree(
   return result;
 }
 
+/**
+ * Builds the data flow information including trees and source categorization
+ */
+function buildDataFlowInfo(
+  sourceIds: Set<IdentifierId>,
+  context: ValidationContext,
+): DataFlowInfo {
+  const propsSet = new Set<string>();
+  const stateSet = new Set<string>();
+
+  const rootNodesMap = new Map<string, TreeNode>();
+  for (const id of sourceIds) {
+    const nodes = buildTreeNode(id, context);
+    for (const node of nodes) {
+      if (!rootNodesMap.has(node.name)) {
+        rootNodesMap.set(node.name, node);
+      }
+    }
+  }
+  const rootNodes = Array.from(rootNodesMap.values());
+
+  const trees = rootNodes.map((node, index) =>
+    renderTree(node, '', index === rootNodes.length - 1, propsSet, stateSet),
+  );
+
+  const propsArr = Array.from(propsSet);
+  const stateArr = Array.from(stateSet);
+
+  let rootSources = '';
+  if (propsArr.length > 0) {
+    rootSources += `Props: [${propsArr.join(', ')}]`;
+  }
+  if (stateArr.length > 0) {
+    if (rootSources) rootSources += '\n';
+    rootSources += `State: [${stateArr.join(', ')}]`;
+  }
+
+  return {
+    rootSources,
+    trees,
+    propsArr,
+    stateArr,
+  };
+}
+
 function getFnLocalDeps(
   fn: FunctionExpression | undefined,
 ): Set<IdentifierId> | undefined {
@@ -744,6 +815,26 @@ function validateEffect(
         );
 
         if (argMetadata !== undefined) {
+          // Check if the value being set is derived from useFragment (Relay)
+          if (argMetadata.typeOfValue === 'fromRelay') {
+            console.log('relay');
+            context.errors.pushDiagnostic(
+              CompilerDiagnostic.create({
+                description:
+                  'Deriving value from Relay store. Instead of holding it in a local state, use a client schema extension to store the value directly in the Relay store.',
+                category: ErrorCategory.EffectDerivationsOfState,
+                reason:
+                  'Avoid storing Relay data in React state. Use client schema extensions instead.',
+              }).withDetails({
+                kind: 'error',
+                loc: instr.value.callee.loc,
+                message:
+                  'setState with Relay-derived value. Use client schema instead.',
+              }),
+            );
+            continue;
+          }
+
           effectDerivedSetStateCalls.push({
             value: instr.value,
             id: instr.value.callee.identifier.id,
@@ -792,47 +883,16 @@ function validateEffect(
       effectSetStateUsages.get(rootSetStateCall)!.size ===
         context.setStateUsages.get(rootSetStateCall)!.size - 1
     ) {
-      const propsSet = new Set<string>();
-      const stateSet = new Set<string>();
-
-      const rootNodesMap = new Map<string, TreeNode>();
-      for (const id of derivedSetStateCall.sourceIds) {
-        const nodes = buildTreeNode(id, context);
-        for (const node of nodes) {
-          if (!rootNodesMap.has(node.name)) {
-            rootNodesMap.set(node.name, node);
-          }
-        }
-      }
-      const rootNodes = Array.from(rootNodesMap.values());
-
-      const trees = rootNodes.map((node, index) =>
-        renderTree(
-          node,
-          '',
-          index === rootNodes.length - 1,
-          propsSet,
-          stateSet,
-        ),
-      );
-
       for (const dep of derivedSetStateCall.sourceIds) {
         if (cleanUpFunctionDeps !== undefined && cleanUpFunctionDeps.has(dep)) {
           return;
         }
       }
 
-      const propsArr = Array.from(propsSet);
-      const stateArr = Array.from(stateSet);
-
-      let rootSources = '';
-      if (propsArr.length > 0) {
-        rootSources += `Props: [${propsArr.join(', ')}]`;
-      }
-      if (stateArr.length > 0) {
-        if (rootSources) rootSources += '\n';
-        rootSources += `State: [${stateArr.join(', ')}]`;
-      }
+      const {rootSources, trees} = buildDataFlowInfo(
+        derivedSetStateCall.sourceIds,
+        context,
+      );
 
       const description = `Using an effect triggers an additional render which can hurt performance and user experience, potentially briefly showing stale values to the user
 
@@ -857,6 +917,80 @@ See: https://react.dev/learn/you-might-not-need-an-effect#updating-state-based-o
           message: 'This should be computed during render, not in an effect',
         }),
       );
+    } else if (
+      rootSetStateCall !== null &&
+      effectSetStateUsages.has(rootSetStateCall) &&
+      context.setStateUsages.has(rootSetStateCall) &&
+      effectSetStateUsages.get(rootSetStateCall)!.size <
+        context.setStateUsages.get(rootSetStateCall)!.size
+    ) {
+      for (const dep of derivedSetStateCall.sourceIds) {
+        if (cleanUpFunctionDeps !== undefined && cleanUpFunctionDeps.has(dep)) {
+          return;
+        }
+      }
+
+      const {rootSources, trees} = buildDataFlowInfo(
+        derivedSetStateCall.sourceIds,
+        context,
+      );
+
+      // Find setState calls outside the effect
+      const allSetStateUsages = context.setStateUsages.get(rootSetStateCall)!;
+      const effectUsages = effectSetStateUsages.get(rootSetStateCall)!;
+      const outsideEffectUsages: Array<SourceLocation> = [];
+
+      for (const usage of allSetStateUsages) {
+        if (!effectUsages.has(usage)) {
+          outsideEffectUsages.push(usage);
+        }
+      }
+
+      const description = `Using an effect triggers an additional render which can hurt performance and user experience, potentially briefly showing stale values to the user
+
+This setState call is setting a derived value that depends on the following reactive sources:
+
+${rootSources}
+
+Data Flow Tree:
+${trees.join('\n')}
+
+This state is also being set outside of the effect. Consider hoisting the state up to a parent component and making this a controlled component.
+
+See: https://react.dev/learn/sharing-state-between-components`;
+
+      const diagnosticDetails: Array<{
+        kind: 'error';
+        loc: SourceLocation;
+        message: string;
+      }> = [
+        {
+          kind: 'error',
+          loc: derivedSetStateCall.value.callee.loc,
+          message: 'setState in effect',
+        },
+      ];
+
+      for (const usage of outsideEffectUsages) {
+        diagnosticDetails.push({
+          kind: 'error',
+          loc: usage,
+          message: 'setState outside effect',
+        });
+      }
+
+      let diagnostic = CompilerDiagnostic.create({
+        description: description,
+        category: ErrorCategory.EffectDerivationsOfState,
+        reason:
+          'Consider hoisting state to parent and making this a controlled component',
+      });
+
+      for (const detail of diagnosticDetails) {
+        diagnostic = diagnostic.withDetails(detail);
+      }
+
+      context.errors.pushDiagnostic(diagnostic);
     }
   }
 }

packages/shared/ReactVersion.js

@@ -1,15 +1 @@
-/**
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- *
- * This source code is licensed under the MIT license found in the
- * LICENSE file in the root directory of this source tree.
- */
-
-// TODO: this is special because it gets imported during build.
-//
-// It exists as a placeholder so that DevTools can support work tag changes between releases.
-// When we next publish a release, update the matching TODO in backend/renderer.js
-// TODO: This module is used both by the release scripts and to expose a version
-// at runtime. We should instead inject the version number as part of the build
-// process, and use the ReactVersions.js module as the single source of truth.
-export default '19.3.0';
+export default '19.3.0-canary--';