[compiler][wip] Track and propagate function effects of arrays and objects

[ghstack-poisoned]

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',
@@ -276,6 +276,8 @@ class InferenceState {
    */
   #variables: Map<IdentifierId, Set<InstructionValue>>;
 
+  #functionEffects: Map<InstructionValue, Array<FunctionEffect>> = new Map();
+
   constructor(
     env: Environment,
     values: Map<InstructionValue, AbstractValue>,
@@ -302,6 +304,19 @@ class InferenceState {
     this.#values.set(value, kind);
   }
 
+  setFunctionEffects(value: InstructionValue, effect: Array<FunctionEffect>): void {
+    CompilerError.invariant(value.kind !== 'LoadLocal', {
+      reason:
+        'Expected all top-level identifiers to be defined as variables, not values',
+      description: null,
+      loc: value.loc,
+      suggestions: null,
+    });
+    const curEffects = this.#functionEffects.get(value) ?? [];
+    curEffects.push(...effect);
+    this.#functionEffects.set(value, curEffects);
+  }
+
   values(place: Place): Array<InstructionValue> {
     const values = this.#variables.get(place.identifier.id);
     CompilerError.invariant(values != null, {
@@ -400,50 +415,57 @@ class InferenceState {
     }
 
     // Propagate effects of function expressions to the outer (ie current) effect context
+    const dependentEffects = [];
     for (const value of values) {
       if (
         (value.kind === 'FunctionExpression' ||
-          value.kind === 'ObjectMethod') &&
-        value.loweredFunc.func.effects != null
+          value.kind === 'ObjectMethod')
       ) {
-        for (const effect of value.loweredFunc.func.effects) {
-          if (
-            effect.kind === 'GlobalMutation' ||
-            effect.kind === 'ReactMutation'
-          ) {
-            // Known effects are always propagated upwards
-            functionEffects.push(effect);
-          } else {
-            /**
-             * Contextual effects need to be replayed against the current inference
-             * state, which may know more about the value to which the effect applied.
-             * The main cases are:
-             * 1. The mutated context value is _still_ a context value in the current scope,
-             *    so we have to continue propagating the original context mutation.
-             * 2. The mutated context value is a mutable value in the current scope,
-             *    so the context mutation was fine and we can skip propagating the effect.
-             * 3. The mutated context value  is an immutable value in the current scope,
-             *    resulting in a non-ContextMutation FunctionEffect. We propagate that new,
-             *    more detailed effect to the current function context.
-             */
-            for (const place of effect.places) {
-              if (this.isDefined(place)) {
-                const replayedEffect = this.reference(
-                  {...place, loc: effect.loc},
-                  effect.effect,
-                  reason,
-                );
-                if (replayedEffect != null) {
-                  if (replayedEffect.kind === 'ContextMutation') {
-                    // Case 1, still a context value so propagate the original effect
-                    functionEffects.push(effect);
-                  } else {
-                    // Case 3, immutable value so propagate the more precise effect
-                    functionEffects.push(replayedEffect);
-                  }
-                } // else case 2, local mutable value so this effect was fine
+        dependentEffects.push(...value.loweredFunc.func.effects ?? []);
+      }
+      const knownEffects = this.#functionEffects.get(value);
+      if (knownEffects != null) {
+        dependentEffects.push(...knownEffects)
+      }
+    }
+
+
+    for (const effect of dependentEffects) {
+      if (
+        effect.kind === 'GlobalMutation' ||
+        effect.kind === 'ReactMutation'
+      ) {
+        // Known effects are always propagated upwards
+        functionEffects.push(effect);
+      } else {
+        /**
+         * Contextual effects need to be replayed against the current inference
+         * state, which may know more about the value to which the effect applied.
+         * The main cases are:
+         * 1. The mutated context value is _still_ a context value in the current scope,
+         *    so we have to continue propagating the original context mutation.
+         * 2. The mutated context value is a mutable value in the current scope,
+         *    so the context mutation was fine and we can skip propagating the effect.
+         * 3. The mutated context value  is an immutable value in the current scope,
+         *    resulting in a non-ContextMutation FunctionEffect. We propagate that new,
+         *    more detailed effect to the current function context.
+         */
+        for (const place of effect.places) {
+          if (this.isDefined(place)) {
+            const replayedEffect = this.reference(
+              {...place, loc: effect.loc},
+              effect.effect,
+              reason,
+            );
+            if (replayedEffect != null) {
+              if (replayedEffect.kind === 'ContextMutation') {
+                // Case 1, still a context value so propagate the original effect
+                functionEffects.push(effect);
+              } else {
+                // Case 3, immutable value so propagate the more precise effect
+                functionEffects.push(replayedEffect);
               }
-            }
+            } // else case 2, local mutable value so this effect was fine
           }
         }
       }
@@ -995,8 +1017,22 @@ function inferBlock(
               context: new Set(),
             };
         effect = {kind: Effect.Capture, reason: ValueReason.Other};
-        lvalueEffect = Effect.Store;
-        break;
+
+        const arrEffects: Array<FunctionEffect> = [];
+        for (const operand of eachInstructionOperand(instr)) {
+          state.referenceAndRecordEffects(
+            operand,
+            effect.kind,
+            effect.reason,
+            arrEffects,
+          );
+        }
+
+        state.initialize(instrValue, valueKind);
+        state.define(instr.lvalue, instrValue);
+        instr.lvalue.effect = Effect.Store;
+        state.setFunctionEffects(instrValue, arrEffects);
+        continue;
       }
       case 'NewExpression': {
         /**
@@ -1039,6 +1075,7 @@ function inferBlock(
         continue;
       }
       case 'ObjectExpression': {
+        const objEffects: Array<FunctionEffect> = [];
         valueKind = hasContextRefOperand(state, instrValue)
           ? {
               kind: ValueKind.Context,
@@ -1060,7 +1097,7 @@ function inferBlock(
                   property.key.name,
                   Effect.Freeze,
                   ValueReason.Other,
-                  functionEffects,
+                  objEffects,
                 );
               }
               // Object construction captures but does not modify the key/property values
@@ -1068,7 +1105,7 @@ function inferBlock(
                 property.place,
                 Effect.Capture,
                 ValueReason.Other,
-                functionEffects,
+                objEffects,
               );
               break;
             }
@@ -1078,7 +1115,7 @@ function inferBlock(
                 property.place,
                 Effect.Capture,
                 ValueReason.Other,
-                functionEffects,
+                objEffects,
               );
               break;
             }
@@ -1092,6 +1129,7 @@ function inferBlock(
         }
 
         state.initialize(instrValue, valueKind);
+        state.setFunctionEffects(instrValue, objEffects);
         state.define(instr.lvalue, instrValue);
         instr.lvalue.effect = Effect.Store;
         continue;
@@ -1151,7 +1189,7 @@ function inferBlock(
             );
             functionEffects.push(
               ...propEffects.filter(
-                propEffect => propEffect.kind !== 'GlobalMutation',
+                effect => !isEffectSafeOutsideRender(effect),
               ),
             );
           }
@@ -1330,7 +1368,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 +1394,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 +1492,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 +1522,7 @@ function inferBlock(
            */
           functionEffects.push(
             ...argumentEffects.filter(
-              argEffect =>
-                !isUseEffect || i !== 0 || argEffect.kind !== 'GlobalMutation',
+              argEffect => !isHook || !isEffectSafeOutsideRender(argEffect),
             ),
           );
           hasCaptureArgument ||= place.effect === Effect.Capture;
@@ -1549,14 +1585,16 @@ function inferBlock(
         break;
       }
       case 'PropertyLoad': {
+        const loadEffects: Array<FunctionEffect> = [];
         state.referenceAndRecordEffects(
           instrValue.object,
           Effect.Read,
           ValueReason.Other,
-          functionEffects,
+          loadEffects,
         );
         const lvalue = instr.lvalue;
         lvalue.effect = Effect.ConditionallyMutate;
+        state.setFunctionEffects(instrValue, loadEffects);
         state.initialize(instrValue, state.kind(instrValue.object));
         state.define(lvalue, instrValue);
         continue;
@@ -1613,22 +1651,24 @@ function inferBlock(
         continue;
       }
       case 'ComputedLoad': {
+        const loadEffects: Array<FunctionEffect> = [];
         state.referenceAndRecordEffects(
           instrValue.object,
           Effect.Read,
           ValueReason.Other,
-          functionEffects,
+          loadEffects,
         );
         state.referenceAndRecordEffects(
           instrValue.property,
           Effect.Read,
           ValueReason.Other,
-          functionEffects,
+          loadEffects,
         );
         const lvalue = instr.lvalue;
         lvalue.effect = Effect.ConditionallyMutate;
         state.initialize(instrValue, state.kind(instrValue.object));
         state.define(lvalue, instrValue);
+        state.setFunctionEffects(instrValue, loadEffects);
         continue;
       }
       case 'Await': {
@@ -2010,11 +2050,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 +2172,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.object-capture-global-mutation.expect.md

@@ -1,33 +0,0 @@
-
-## Input
-
-```javascript
-function Foo() {
-  const x = () => {
-    window.href = 'foo';
-  };
-  const y = {x};
-  return <Bar y={y} />;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Foo,
-  params: [],
-};
-
-```
-
-
-## Error
-
-```
-  1 | function Foo() {
-  2 |   const x = () => {
-> 3 |     window.href = 'foo';
-    |     ^^^^^^ InvalidReact: Writing to a variable defined outside a component or hook is not allowed. Consider using an effect (3:3)
-  4 |   };
-  5 |   const y = {x};
-  6 |   return <Bar y={y} />;
-```
-          
-      

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/error.reassign-global-obj-call.expect.md

@@ -0,0 +1,35 @@
+
+## Input
+
+```javascript
+let b = 1;
+
+export default function useMyHook() {
+  const fn = () => {
+    b = 2;
+  };
+  const obj = { fn };
+  obj.fn();
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMyHook,
+  params: [],
+};
+
+```
+
+
+## Error
+
+```
+  3 | export default function useMyHook() {
+  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 |   const obj = { fn };
+  8 |   obj.fn();
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.reassign-global-obj-call.js

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

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/object-capture-global-mutation.expect.md

@@ -0,0 +1,49 @@
+
+## Input
+
+```javascript
+function Foo() {
+  const x = () => {
+    window.href = 'foo';
+  };
+  const y = {x};
+  return <Bar y={y} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+function Foo() {
+  const $ = _c(1);
+  const x = _temp;
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    const y = { x };
+    t0 = <Bar y={y} />;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+function _temp() {
+  window.href = "foo";
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: exception) Bar is not defined

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-obj.expect.md

@@ -0,0 +1,51 @@
+
+## Input
+
+```javascript
+let b = 1;
+
+export default function useMyHook() {
+  const fn = () => {
+    b = 2;
+  };
+  return { fn };
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMyHook,
+  params: [],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+let b = 1;
+
+export default function useMyHook() {
+  const $ = _c(1);
+  const fn = _temp;
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = { fn };
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  return t0;
+}
+function _temp() {
+  b = 2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useMyHook,
+  params: [],
+};
+
+```
+      
+### Eval output
+(kind: ok) {"fn":"[[ function params=0 ]]"}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/reassign-global-return-obj.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: [],
+};

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: [],
+};