[compiler] Effects for Return/MaybeThrow terminals

ghstack-source-id: a0411c7
Pull Request resolved: #33429

Author: josephsavona

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

@@ -179,6 +179,7 @@ export function lower(
       loc: GeneratedSource,
       value: lowerExpressionToTemporary(builder, body),
       id: makeInstructionId(0),
+      effects: null,
     };
     builder.terminateWithContinuation(terminal, fallthrough);
   } else if (body.isBlockStatement()) {
@@ -208,6 +209,7 @@ export function lower(
         loc: GeneratedSource,
       }),
       id: makeInstructionId(0),
+      effects: null,
     },
     null,
   );
@@ -287,6 +289,7 @@ function lowerStatement(
         loc: stmt.node.loc ?? GeneratedSource,
         value,
         id: makeInstructionId(0),
+        effects: null,
       };
       builder.terminate(terminal, 'block');
       return;

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

@@ -457,6 +457,7 @@ export type ReturnTerminal = {
   value: Place;
   id: InstructionId;
   fallthrough?: never;
+  effects: Array<AliasingEffect> | null;
 };
 
 export type GotoTerminal = {
@@ -617,6 +618,7 @@ export type MaybeThrowTerminal = {
   id: InstructionId;
   loc: SourceLocation;
   fallthrough?: never;
+  effects: Array<AliasingEffect> | null;
 };
 
 export type ReactiveScopeTerminal = {

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

@@ -165,6 +165,7 @@ export default class HIRBuilder {
           handler: exceptionHandler,
           id: makeInstructionId(0),
           loc: instruction.loc,
+          effects: null,
         },
         continuationBlock,
       );

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

@@ -222,6 +222,9 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
       value = `[${terminal.id}] Return${
         terminal.value != null ? ' ' + printPlace(terminal.value) : ''
       }`;
+      if (terminal.effects != null) {
+        value += `\n    ${terminal.effects.map(printAliasingEffect).join('\n    ')}`;
+      }
       break;
     }
     case 'goto': {
@@ -290,6 +293,9 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
     }
     case 'maybe-throw': {
       value = `[${terminal.id}] MaybeThrow continuation=bb${terminal.continuation} handler=bb${terminal.handler}`;
+      if (terminal.effects != null) {
+        value += `\n    ${terminal.effects.map(printAliasingEffect).join('\n    ')}`;
+      }
       break;
     }
     case 'scope': {

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

@@ -735,6 +735,7 @@ export function mapTerminalSuccessors(
         loc: terminal.loc,
         value: terminal.value,
         id: makeInstructionId(0),
+        effects: terminal.effects,
       };
     }
     case 'throw': {
@@ -842,6 +843,7 @@ export function mapTerminalSuccessors(
         handler,
         id: makeInstructionId(0),
         loc: terminal.loc,
+        effects: terminal.effects,
       };
     }
     case 'try': {

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

@@ -15,7 +15,6 @@ import {
   Place,
   isRefOrRefValue,
   makeInstructionId,
-  printFunction,
 } from '../HIR';
 import {deadCodeElimination} from '../Optimization';
 import {inferReactiveScopeVariables} from '../ReactiveScopes';

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

@@ -155,7 +155,7 @@ export function inferMutationAliasingEffects(
   }
   queue(fn.body.entry, initialState);
 
-  const context = new Context();
+  const context = new Context(isFunctionExpression);
 
   let count = 0;
   while (queuedStates.size !== 0) {
@@ -192,6 +192,11 @@ class Context {
   effectInstructionValueCache: Map<AliasingEffect, InstructionValue> =
     new Map();
   catchHandlers: Map<BlockId, Place> = new Map();
+  isFuctionExpression: boolean;
+
+  constructor(isFunctionExpression: boolean) {
+    this.isFuctionExpression = isFunctionExpression;
+  }
 }
 
 function inferParam(
@@ -233,6 +238,7 @@ function inferBlock(
   } else if (terminal.kind === 'maybe-throw') {
     const handlerParam = context.catchHandlers.get(terminal.handler);
     if (handlerParam != null) {
+      const effects: Array<AliasingEffect> = [];
       for (const instr of block.instructions) {
         if (
           instr.value.kind === 'CallExpression' ||
@@ -243,10 +249,33 @@ function inferBlock(
            * itself. For example, `c = a.b` can throw if `a` is nullish, but the thrown value
            * is an error object synthesized by the JS runtime. Whereas `throwsInput(x)` can
            * throw (effectively) the result of the call.
+           *
+           * TODO: call applyEffect() instead. This meant that the catch param wasn't inferred
+           * as a mutable value, though. See `try-catch-try-value-modified-in-catch-escaping.js`
+           * fixture as an example
            */
           state.appendAlias(handlerParam, instr.lvalue);
+          const kind = state.kind(instr.lvalue).kind;
+          if (kind === ValueKind.Mutable || kind == ValueKind.Context) {
+            effects.push({
+              kind: 'Alias',
+              from: instr.lvalue,
+              into: handlerParam,
+            });
+          }
         }
       }
+      terminal.effects = effects.length !== 0 ? effects : null;
+    }
+  } else if (terminal.kind === 'return') {
+    if (!context.isFuctionExpression) {
+      terminal.effects = [
+        {
+          kind: 'Freeze',
+          value: terminal.value,
+          reason: ValueReason.JsxCaptured,
+        },
+      ];
     }
   }
 }
@@ -326,7 +355,7 @@ function applySignature(
   }
 
   for (const effect of signature.effects) {
-    applyEffect(context, state, effect, instruction, aliased, effects);
+    applyEffect(context, state, effect, aliased, effects);
   }
   if (DEBUG) {
     console.log(
@@ -351,7 +380,6 @@ function applyEffect(
   context: Context,
   state: InferenceState,
   effect: AliasingEffect,
-  instruction: Instruction,
   aliased: Set<IdentifierId>,
   effects: Array<AliasingEffect>,
 ): void {
@@ -477,7 +505,6 @@ function applyEffect(
             from: capture,
             into: effect.into,
           },
-          instruction,
           aliased,
           effects,
         );
@@ -618,14 +645,7 @@ function applyEffect(
             console.log('apply function expression effects');
           }
           for (const signatureEffect of signatureEffects) {
-            applyEffect(
-              context,
-              state,
-              signatureEffect,
-              instruction,
-              aliased,
-              effects,
-            );
+            applyEffect(context, state, signatureEffect, aliased, effects);
           }
           break;
         }
@@ -645,14 +665,7 @@ function applyEffect(
           console.log('apply aliasing signature effects');
         }
         for (const signatureEffect of signatureEffects) {
-          applyEffect(
-            context,
-            state,
-            signatureEffect,
-            instruction,
-            aliased,
-            effects,
-          );
+          applyEffect(context, state, signatureEffect, aliased, effects);
         }
       } else if (effect.signature != null) {
         if (DEBUG) {
@@ -666,14 +679,7 @@ function applyEffect(
           effect.args,
         );
         for (const legacyEffect of legacyEffects) {
-          applyEffect(
-            context,
-            state,
-            legacyEffect,
-            instruction,
-            aliased,
-            effects,
-          );
+          applyEffect(context, state, legacyEffect, aliased, effects);
         }
       } else {
         if (DEBUG) {
@@ -687,7 +693,6 @@ function applyEffect(
             into: effect.into,
             value: ValueKind.Mutable,
           },
-          instruction,
           aliased,
           effects,
         );
@@ -711,7 +716,6 @@ function applyEffect(
                 kind: 'MutateTransitiveConditionally',
                 value: operand,
               },
-              instruction,
               aliased,
               effects,
             );
@@ -721,7 +725,6 @@ function applyEffect(
             state,
             // OK: recording information flow
             {kind: 'Alias', from: operand, into: effect.into},
-            instruction,
             aliased,
             effects,
           );
@@ -750,7 +753,6 @@ function applyEffect(
                 from: operand,
                 into: other,
               },
-              instruction,
               aliased,
               effects,
             );
@@ -772,7 +774,6 @@ function applyEffect(
       ) {
         const value = state.kind(effect.value);
         if (DEBUG) {
-          console.log(printInstruction(instruction));
           console.log(printAliasingEffect(effect));
           console.log(prettyFormat(state.debugAbstractValue(value)));
         }

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

@@ -71,7 +71,6 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
     kind: MutationKind;
     place: Place;
   }> = [];
-  const catchHandlers = new Map<BlockId, Place>();
 
   let index = 0;
 
@@ -157,20 +156,19 @@ export function inferMutationAliasingRanges(fn: HIRFunction): void {
       state.assign(index++, block.terminal.value, fn.returns);
     }
 
-    /**
-     * TODO: add effects to terminals so that these can be emitted by the equivalent
-     * logic in InferMutationAliasingEffects
-     */
     if (
-      block.terminal.kind === 'try' &&
-      block.terminal.handlerBinding != null
+      (block.terminal.kind === 'maybe-throw' ||
+        block.terminal.kind === 'return') &&
+      block.terminal.effects != null
     ) {
-      catchHandlers.set(block.terminal.handler, block.terminal.handlerBinding);
-    } else if (block.terminal.kind === 'maybe-throw') {
-      const handlerParam = catchHandlers.get(block.terminal.handler);
-      if (handlerParam != null) {
-        for (const instr of block.instructions) {
-          state.assign(index++, instr.lvalue, handlerParam);
+      for (const effect of block.terminal.effects) {
+        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`,
+            loc: block.terminal.loc,
+          });
         }
       }
     }

compiler/packages/babel-plugin-react-compiler/src/Optimization/LowerContextAccess.ts

@@ -241,6 +241,7 @@ function emitSelectorFn(env: Environment, keys: Array<string>): Instruction {
       kind: 'return',
       loc: GeneratedSource,
       value: arrayInstr.lvalue,
+      effects: null,
     },
     preds: new Set(),
     phis: new Set(),

compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts

@@ -356,6 +356,7 @@ function emitOutlinedFn(
       kind: 'return',
       loc: GeneratedSource,
       value: instructions.at(-1)!.lvalue,
+      effects: null,
     },
     preds: new Set(),
     phis: new Set(),