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

ghstack-source-id: 9039083
Pull Request resolved: #30578

Author: mvitousek

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

@@ -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,15 +1097,15 @@ function inferBlock(
                   property.key.name,
                   Effect.Freeze,
                   ValueReason.Other,
-                  functionEffects,
+                  objEffects,
                 );
               }
               // Object construction captures but does not modify the key/property values
               state.referenceAndRecordEffects(
                 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;
@@ -1547,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;
@@ -1611,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': {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.object-capture-global-mutation.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.expect.md

@@ -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/error.reassign-global-obj-call.js

@@ -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