[compiler] fixes for effects capturing context vars, hoisted functions

ghstack-source-id: 7528f90
Pull Request resolved: #33469

Author: josephsavona

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

@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {Effect, ValueKind, ValueReason} from './HIR';
+import {Effect, makeIdentifierId, ValueKind, ValueReason} from './HIR';
 import {
   BUILTIN_SHAPES,
   BuiltInArrayId,
@@ -32,6 +32,7 @@ import {
   addFunction,
   addHook,
   addObject,
+  signatureArgument,
 } from './ObjectShape';
 import {BuiltInType, ObjectType, PolyType} from './Types';
 import {TypeConfig} from './TypeSchema';
@@ -642,6 +643,41 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
         calleeEffect: Effect.Read,
         hookKind: 'useEffect',
         returnValueKind: ValueKind.Frozen,
+        aliasing: {
+          receiver: makeIdentifierId(0),
+          params: [],
+          rest: makeIdentifierId(1),
+          returns: makeIdentifierId(2),
+          temporaries: [signatureArgument(3)],
+          effects: [
+            // Freezes the function and deps
+            {
+              kind: 'Freeze',
+              value: signatureArgument(1),
+              reason: ValueReason.Effect,
+            },
+            // Internally creates an effect object that captures the function and deps
+            {
+              kind: 'Create',
+              into: signatureArgument(3),
+              value: ValueKind.Frozen,
+              reason: ValueReason.KnownReturnSignature,
+            },
+            // The effect stores the function and dependencies
+            {
+              kind: 'Capture',
+              from: signatureArgument(1),
+              into: signatureArgument(3),
+            },
+            // Returns undefined
+            {
+              kind: 'Create',
+              into: signatureArgument(2),
+              value: ValueKind.Primitive,
+              reason: ValueReason.KnownReturnSignature,
+            },
+          ],
+        },
       },
       BuiltInUseEffectHookId,
     ),

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

@@ -1396,6 +1396,11 @@ export enum ValueReason {
    */
   JsxCaptured = 'jsx-captured',
 
+  /**
+   * Passed to an effect
+   */
+  Effect = 'effect',
+
   /**
    * Return value of a function with known frozen return value, e.g. `useState`.
    */

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

@@ -357,6 +357,8 @@ export function getWriteErrorReason(abstractValue: AbstractValue): string {
     return "Mutating a value returned from 'useState()', which should not be mutated. Use the setter function to update instead";
   } else if (abstractValue.reason.has(ValueReason.ReducerState)) {
     return "Mutating a value returned from 'useReducer()', which should not be mutated. Use the dispatch function to update instead";
+  } else if (abstractValue.reason.has(ValueReason.Effect)) {
+    return 'Updating a value used previously in an effect function or as an effect dependency is not allowed. Consider moving the mutation before calling useEffect()';
   } else {
     return 'This mutates a variable that React considers immutable';
   }

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

@@ -66,6 +66,8 @@ import {getWriteErrorReason} from './InferFunctionEffects';
 import prettyFormat from 'pretty-format';
 import {createTemporaryPlace} from '../HIR/HIRBuilder';
 
+const DEBUG = false;
+
 export function inferMutationAliasingEffects(
   fn: HIRFunction,
   {isFunctionExpression}: {isFunctionExpression: boolean} = {
@@ -849,7 +851,7 @@ function applyEffect(
       ) {
         const value = state.kind(effect.value);
         if (DEBUG) {
-          console.log(printAliasingEffect(effect));
+          console.log(`invalid mutation: ${printAliasingEffect(effect)}`);
           console.log(prettyFormat(state.debugAbstractValue(value)));
         }
 
@@ -891,8 +893,6 @@ function applyEffect(
   }
 }
 
-const DEBUG = false;
-
 class InferenceState {
   env: Environment;
   #isFunctionExpression: boolean;

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

@@ -381,6 +381,39 @@ export function inferMutationAliasingRanges(
         const effect = operandEffects.get(operand.identifier.id) ?? Effect.Read;
         operand.effect = effect;
       }
+
+      /**
+       * This case is targeted at hoisted functions like:
+       *
+       * ```
+       * x();
+       * function x() { ... }
+       * ```
+       *
+       * Which turns into:
+       *
+       * t0 = DeclareContext HoistedFunction x
+       * t1 = LoadContext x
+       * t2 = CallExpression t1 ( )
+       * t3 = FunctionExpression ...
+       * t4 = StoreContext Function x = t3
+       *
+       * If the function had captured mutable values, it would already have its
+       * range extended to include the StoreContext. But if the function doesn't
+       * capture any mutable values its range won't have been extended yet. We
+       * want to ensure that the value is memoized along with the context variable,
+       * not independently of it (bc of the way we do codegen for hoisted functions).
+       * So here we check for StoreContext rvalues and if they haven't already had
+       * their range extended to at least this instruction, we extend it.
+       */
+      if (
+        instr.value.kind === 'StoreContext' &&
+        instr.value.value.identifier.mutableRange.end <= instr.id
+      ) {
+        instr.value.value.identifier.mutableRange.end = makeInstructionId(
+          instr.id + 1,
+        );
+      }
     }
     if (block.terminal.kind === 'return') {
       block.terminal.value.effect = isFunctionExpression