[compiler] Allow global mutation effects in arguments passed to hooks and in return values

compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts

@@ -26,6 +26,7 @@ import {
   Type,
   ValueKind,
   ValueReason,
+  getHookKind,
   isArrayType,
   isMutableEffect,
   isObjectType,
@@ -48,7 +49,6 @@ import {
   eachTerminalSuccessor,
 } from '../HIR/visitors';
 import {assertExhaustive} from '../Utils/utils';
-import {isEffectHook} from '../Validation/ValidateMemoizedEffectDependencies';
 
 const UndefinedValue: InstructionValue = {
   kind: 'Primitive',
@@ -1151,7 +1151,7 @@ function inferBlock(
             );
             functionEffects.push(
               ...propEffects.filter(
-                propEffect => propEffect.kind !== 'GlobalMutation',
+                effect => !isEffectSafeOutsideRender(effect),
               ),
             );
           }
@@ -1330,7 +1330,7 @@ function inferBlock(
                 context: new Set(),
               };
         let hasCaptureArgument = false;
-        let isUseEffect = isEffectHook(instrValue.callee.identifier);
+        let isHook = getHookKind(env, instrValue.callee.identifier) != null;
         for (let i = 0; i < instrValue.args.length; i++) {
           const argumentEffects: Array<FunctionEffect> = [];
           const arg = instrValue.args[i];
@@ -1356,8 +1356,7 @@ function inferBlock(
            */
           functionEffects.push(
             ...argumentEffects.filter(
-              argEffect =>
-                !isUseEffect || i !== 0 || argEffect.kind !== 'GlobalMutation',
+              argEffect => !isHook || !isEffectSafeOutsideRender(argEffect),
             ),
           );
           hasCaptureArgument ||= place.effect === Effect.Capture;
@@ -1455,7 +1454,7 @@ function inferBlock(
         const effects =
           signature !== null ? getFunctionEffects(instrValue, signature) : null;
         let hasCaptureArgument = false;
-        let isUseEffect = isEffectHook(instrValue.property.identifier);
+        let isHook = getHookKind(env, instrValue.property.identifier) != null;
         for (let i = 0; i < instrValue.args.length; i++) {
           const argumentEffects: Array<FunctionEffect> = [];
           const arg = instrValue.args[i];
@@ -1485,8 +1484,7 @@ function inferBlock(
            */
           functionEffects.push(
             ...argumentEffects.filter(
-              argEffect =>
-                !isUseEffect || i !== 0 || argEffect.kind !== 'GlobalMutation',
+              argEffect => !isHook || !isEffectSafeOutsideRender(argEffect),
             ),
           );
           hasCaptureArgument ||= place.effect === Effect.Capture;
@@ -2010,11 +2008,15 @@ function inferBlock(
     } else {
       effect = Effect.Read;
     }
+    const propEffects: Array<FunctionEffect> = [];
     state.referenceAndRecordEffects(
       operand,
       effect,
       ValueReason.Other,
-      functionEffects,
+      propEffects,
+    );
+    functionEffects.push(
+      ...propEffects.filter(effect => !isEffectSafeOutsideRender(effect)),
     );
   }
 }
@@ -2128,6 +2130,10 @@ function areArgumentsImmutableAndNonMutating(
   return true;
 }
 
+function isEffectSafeOutsideRender(effect: FunctionEffect): boolean {
+  return effect.kind === 'GlobalMutation';
+}
+
 function getWriteErrorReason(abstractValue: AbstractValue): string {
   if (abstractValue.reason.has(ValueReason.Global)) {
     return 'Writing to a variable defined outside a component or hook is not allowed. Consider using an effect';

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.reassign-global-fn-arg.expect.md

@@ -0,0 +1,36 @@
+
+## Input
+
+```javascript
+let b = 1;
+
+export default function MyApp() {
+  const fn = () => {
+    b = 2;
+  };
+  return foo(fn);
+}
+
+function foo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};
+
+```
+
+
+## Error
+
+```
+  3 | export default function MyApp() {
+  4 |   const fn = () => {
+> 5 |     b = 2;
+    |     ^ InvalidReact: Unexpected reassignment of a variable which was defined outside of the component. Components and hooks should be pure and side-effect free, but variable reassignment is a form of side-effect. If this variable is used in rendering, use useState instead. (https://react.dev/reference/rules/components-and-hooks-must-be-pure#side-effects-must-run-outside-of-render) (5:5)
+  6 |   };
+  7 |   return foo(fn);
+  8 | }
+```
+
+

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.reassign-global-fn-arg.js

@@ -0,0 +1,15 @@
+let b = 1;
+
+export default function MyApp() {
+  const fn = () => {
+    b = 2;
+  };
+  return foo(fn);
+}
+
+function foo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reassign-global-hook-arg.expect.md

@@ -0,0 +1,46 @@
+
+## Input
+
+```javascript
+let b = 1;
+
+export default function MyApp() {
+  const fn = () => {
+    b = 2;
+  };
+  return useFoo(fn);
+}
+
+function useFoo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+let b = 1;
+
+export default function MyApp() {
+  const fn = _temp;
+  return useFoo(fn);
+}
+function _temp() {
+  b = 2;
+}
+
+function useFoo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};
+
+```
+
+### Eval output
+(kind: ok) 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reassign-global-hook-arg.js

@@ -0,0 +1,15 @@
+let b = 1;
+
+export default function MyApp() {
+  const fn = () => {
+    b = 2;
+  };
+  return useFoo(fn);
+}
+
+function useFoo(fn) {}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: MyApp,
+  params: [],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reassign-global-return.expect.md

@@ -0,0 +1,42 @@
+
+## Input
+
+```javascript
+let b = 1;
+
+export default function useMyHook() {
+  const fn = () => {
+    b = 2;
+  };
+  return fn;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMyHook,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+let b = 1;
+
+export default function useMyHook() {
+  const fn = _temp;
+  return fn;
+}
+function _temp() {
+  b = 2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMyHook,
+  params: [],
+};
+
+```
+
+### Eval output
+(kind: ok) "[[ function params=0 ]]"

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reassign-global-return.js

@@ -0,0 +1,13 @@
+let b = 1;
+
+export default function useMyHook() {
+  const fn = () => {
+    b = 2;
+  };
+  return fn;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMyHook,
+  params: [],
+};