[compiler] Use aliasing effects for impurity inference

Author: josephsavona

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

@@ -664,7 +664,7 @@ const RenderHookAliasing: (
 
 const EffectHookAliasing: AliasingSignatureConfig = {
   receiver: '@receiver',
-  params: [],
+  params: ['@fn', '@deps'],
   rest: '@rest',
   returns: '@returns',
   temporaries: ['@effect'],
@@ -675,6 +675,21 @@ const EffectHookAliasing: AliasingSignatureConfig = {
       value: '@rest',
       reason: ValueReason.Effect,
     },
+    {
+      kind: 'Freeze',
+      value: '@fn',
+      reason: ValueReason.Effect,
+    },
+    {
+      kind: 'Freeze',
+      value: '@deps',
+      reason: ValueReason.Effect,
+    },
+    // Deps are accessed during render
+    {
+      kind: 'Render',
+      place: '@deps',
+    },
     // Internally creates an effect object that captures the function and deps
     {
       kind: 'Create',
@@ -688,6 +703,11 @@ const EffectHookAliasing: AliasingSignatureConfig = {
       from: '@rest',
       into: '@effect',
     },
+    {
+      kind: 'Capture',
+      from: '@fn',
+      into: '@effect',
+    },
     // Returns undefined
     {
       kind: 'Create',
@@ -703,6 +723,39 @@ const EffectHookAliasing: AliasingSignatureConfig = {
  * now that FeatureFlag `enableTreatHooksAsFunctions` is removed we can
  * use positional params too (?)
  */
+const useEffectEvent = addHook(
+  DEFAULT_SHAPES,
+  {
+    positionalParams: [],
+    restParam: Effect.Freeze,
+    returnType: {
+      kind: 'Function',
+      return: {kind: 'Poly'},
+      shapeId: BuiltInEffectEventId,
+      isConstructor: false,
+    },
+    calleeEffect: Effect.Read,
+    hookKind: 'useEffectEvent',
+    // Frozen because it should not mutate any locally-bound values
+    returnValueKind: ValueKind.Frozen,
+    aliasing: {
+      receiver: '@receiver',
+      params: ['@value'],
+      rest: null,
+      returns: '@return',
+      temporaries: [],
+      effects: [
+        {kind: 'Assign', from: '@value', into: '@return'},
+        {
+          kind: 'Freeze',
+          value: '@value',
+          reason: ValueReason.HookCaptured,
+        },
+      ],
+    },
+  },
+  BuiltInUseEffectEventId,
+);
 const REACT_APIS: Array<[string, BuiltInType]> = [
   [
     'useContext',
@@ -915,27 +968,8 @@ const REACT_APIS: Array<[string, BuiltInType]> = [
       BuiltInFireId,
     ),
   ],
-  [
-    'useEffectEvent',
-    addHook(
-      DEFAULT_SHAPES,
-      {
-        positionalParams: [],
-        restParam: Effect.Freeze,
-        returnType: {
-          kind: 'Function',
-          return: {kind: 'Poly'},
-          shapeId: BuiltInEffectEventId,
-          isConstructor: false,
-        },
-        calleeEffect: Effect.Read,
-        hookKind: 'useEffectEvent',
-        // Frozen because it should not mutate any locally-bound values
-        returnValueKind: ValueKind.Frozen,
-      },
-      BuiltInUseEffectEventId,
-    ),
-  ],
+  ['useEffectEvent', useEffectEvent],
+  ['experimental_useEffectEvent', useEffectEvent],
   ['AUTODEPS', addObject(DEFAULT_SHAPES, BuiltInAutodepsId, [])],
 ];
 

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

@@ -1879,7 +1879,15 @@ export function isRefValueType(id: Identifier): boolean {
 }
 
 export function isUseRefType(id: Identifier): boolean {
-  return id.type.kind === 'Object' && id.type.shapeId === 'BuiltInUseRefId';
+  return isUseRefType_(id.type);
+}
+
+export function isUseRefType_(type: Type): boolean {
+  return (
+    (type.kind === 'Object' && type.shapeId === 'BuiltInUseRefId') ||
+    (type.kind === 'Phi' &&
+      type.operands.some(operand => isUseRefType_(operand)))
+  );
 }
 
 export function isUseStateType(id: Identifier): boolean {

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

@@ -29,6 +29,7 @@ import {
   isArrayType,
   isJsxOrJsxUnionType,
   isMapType,
+  isMutableEffect,
   isPrimitiveType,
   isRefOrRefValue,
   isSetType,
@@ -587,6 +588,17 @@ function inferBlock(
         }),
       );
     }
+    if (
+      context.fn.fnType === 'Component' ||
+      isJsxOrJsxUnionType(context.fn.returns.identifier.type)
+    ) {
+      terminal.effects.push(
+        context.internEffect({
+          kind: 'Render',
+          place: terminal.value,
+        }),
+      );
+    }
   }
 }
 
@@ -758,17 +770,7 @@ function applyEffect(
       break;
     }
     case 'ImmutableCapture': {
-      const kind = state.kind(effect.from).kind;
-      switch (kind) {
-        case ValueKind.Global:
-        case ValueKind.Primitive: {
-          // no-op: we don't need to track data flow for copy types
-          break;
-        }
-        default: {
-          effects.push(effect);
-        }
-      }
+      effects.push(effect);
       break;
     }
     case 'CreateFrom': {
@@ -1070,6 +1072,17 @@ function applyEffect(
             reason: new Set(fromValue.reason),
           });
           state.define(effect.into, value);
+          applyEffect(
+            context,
+            state,
+            {
+              kind: 'ImmutableCapture',
+              from: effect.from,
+              into: effect.into,
+            },
+            initialized,
+            effects,
+          );
           break;
         }
         default: {
@@ -1975,6 +1988,11 @@ function computeSignatureForInstruction(
           value: ValueKind.Primitive,
           reason: ValueReason.Other,
         });
+        effects.push({
+          kind: 'ImmutableCapture',
+          from: value.object,
+          into: lvalue,
+        });
       } else {
         effects.push({
           kind: 'CreateFrom',
@@ -2180,6 +2198,9 @@ function computeSignatureForInstruction(
         for (const prop of value.props) {
           const place =
             prop.kind === 'JsxAttribute' ? prop.place : prop.argument;
+          if (isUseRefType(place.identifier)) {
+            continue;
+          }
           if (place.identifier.type.kind === 'Function') {
             if (isJsxOrJsxUnionType(place.identifier.type.return)) {
               effects.push({
@@ -2226,6 +2247,11 @@ function computeSignatureForInstruction(
             value: ValueKind.Primitive,
             reason: ValueReason.Other,
           });
+          effects.push({
+            kind: 'ImmutableCapture',
+            from: value.value,
+            into: place,
+          });
         } else if (patternItem.kind === 'Identifier') {
           effects.push({
             kind: 'CreateFrom',
@@ -2407,15 +2433,46 @@ function computeSignatureForInstruction(
       });
       break;
     }
+    case 'BinaryExpression': {
+      effects.push({
+        kind: 'Create',
+        into: lvalue,
+        value: ValueKind.Primitive,
+        reason: ValueReason.Other,
+      });
+      effects.push({
+        kind: 'ImmutableCapture',
+        into: lvalue,
+        from: value.left,
+      });
+      effects.push({
+        kind: 'ImmutableCapture',
+        into: lvalue,
+        from: value.right,
+      });
+      break;
+    }
+    case 'UnaryExpression': {
+      effects.push({
+        kind: 'Create',
+        into: lvalue,
+        value: ValueKind.Primitive,
+        reason: ValueReason.Other,
+      });
+      effects.push({
+        kind: 'ImmutableCapture',
+        into: lvalue,
+        from: value.value,
+      });
+      break;
+    }
     case 'TaggedTemplateExpression':
-    case 'BinaryExpression':
     case 'Debugger':
     case 'JSXText':
     case 'MetaProperty':
     case 'Primitive':
     case 'RegExpLiteral':
     case 'TemplateLiteral':
-    case 'UnaryExpression':
     case 'UnsupportedNode': {
       effects.push({
         kind: 'Create',

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

@@ -224,17 +224,20 @@ export function inferMutationAliasingRanges(
         if (effect.kind === 'Alias') {
           state.assign(index++, effect.from, effect.into);
         } else {
-          CompilerError.invariant(effect.kind === 'Freeze', {
-            reason: `Unexpected '${effect.kind}' effect for MaybeThrow terminal`,
-            description: null,
-            details: [
-              {
-                kind: 'error',
-                loc: block.terminal.loc,
-                message: null,
-              },
-            ],
-          });
+          CompilerError.invariant(
+            effect.kind === 'Freeze' || effect.kind === 'Render',
+            {
+              reason: `Unexpected '${effect.kind}' effect for MaybeThrow terminal`,
+              description: null,
+              details: [
+                {
+                  kind: 'error',
+                  loc: block.terminal.loc,
+                  message: null,
+                },
+              ],
+            },
+          );
         }
       }
     }

compiler/packages/babel-plugin-react-compiler/src/Utils/utils.ts

@@ -167,6 +167,14 @@ export function Set_filter<T>(
   return result;
 }
 
+export function Set_subtract<T>(
+  source: ReadonlySet<T>,
+  other: Iterable<T>,
+): Set<T> {
+  const otherSet = other instanceof Set ? other : new Set(other);
+  return Set_filter(source, item => !otherSet.has(item));
+}
+
 export function hasNode<T>(
   input: NodePath<T | null | undefined>,
 ): input is NodePath<NonNullable<T>> {