TryStatement: handle mutation of try values via catch param

compiler/packages/babel-plugin-react-forget/src/Inference/InferMutableRanges.ts

@@ -11,6 +11,7 @@ import { inferAliasForPhis } from "./InferAliasForPhis";
 import { inferAliasForStores } from "./InferAliasForStores";
 import { inferMutableLifetimes } from "./InferMutableLifetimes";
 import { inferMutableRangesForAlias } from "./InferMutableRangesForAlias";
+import { inferTryCatchAliases } from "./InferTryCatchAliases";
 
 export function inferMutableRanges(ir: HIRFunction): void {
   // Infer mutable ranges for non fields
@@ -18,6 +19,10 @@ export function inferMutableRanges(ir: HIRFunction): void {
 
   // Calculate aliases
   const aliases = inferAliases(ir);
+  // Calculate aliases for try/catch, where any value created
+  // in the try block could be aliased to the catch param
+  inferTryCatchAliases(ir, aliases);
+
   // Eagerly canonicalize so that if nothing changes we can bail out
   // after a single iteration
   let prevAliases: Map<Identifier, Identifier> = aliases.canonicalize();

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

@@ -14,6 +14,7 @@ import {
   Effect,
   HIRFunction,
   IdentifierId,
+  InstructionKind,
   InstructionValue,
   isMutableEffect,
   isObjectType,
@@ -924,7 +925,13 @@ function inferBlock(
           loc: instrValue.loc,
           value: undefined,
         };
-        state.initialize(value, ValueKind.Immutable);
+        state.initialize(
+          value,
+          // Catch params may be aliased to mutable values
+          instrValue.lvalue.kind === InstructionKind.Catch
+            ? ValueKind.Mutable
+            : ValueKind.Immutable
+        );
         state.define(instrValue.lvalue.place, value);
         continue;
       }

compiler/packages/babel-plugin-react-forget/src/Inference/InferTryCatchAliases.ts

@@ -0,0 +1,45 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { BlockId, HIRFunction, Identifier } from "../HIR";
+import DisjointSet from "../Utils/DisjointSet";
+
+/**
+ * Any values created within a try/catch block could be aliased to the try handler.
+ * Our lowering ensures that every instruction within a try block will be lowered into a
+ * basic block ending in a maybe-throw terminal that points to its catch block, so we can
+ * iterate such blocks and alias their instruction lvalues to the handler's param (if present).
+ */
+export function inferTryCatchAliases(
+  fn: HIRFunction,
+  aliases: DisjointSet<Identifier>
+): void {
+  const handlerParams: Map<BlockId, Identifier> = new Map();
+  for (const [_, block] of fn.body.blocks) {
+    if (
+      block.terminal.kind === "try" &&
+      block.terminal.handlerBinding !== null
+    ) {
+      handlerParams.set(
+        block.terminal.handler,
+        block.terminal.handlerBinding.identifier
+      );
+    } else if (block.terminal.kind === "maybe-throw") {
+      const handlerParam = handlerParams.get(block.terminal.handler);
+      if (handlerParam === undefined) {
+        // There's no catch clause param, nothing to alias to so
+        // skip this block
+        continue;
+      }
+      // Otherwise alias all values created in this block to the
+      // catch clause param
+      for (const instr of block.instructions) {
+        aliases.union([handlerParam, instr.lvalue.identifier]);
+      }
+    }
+  }
+}

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-alias-try-values.expect.md

@@ -0,0 +1,62 @@
+
+## Input
+
+```javascript
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  let y;
+  let x = [];
+  try {
+    // throws x
+    throwInput(x);
+  } catch (e) {
+    // e = x
+    y = e; // y = x
+  }
+  y.push(null);
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  const $ = useMemoCache(1);
+  let x;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    let y;
+    x = [];
+    try {
+      throwInput(x);
+    } catch (t27) {
+      const e = t27;
+
+      y = e;
+    }
+
+    y.push(null);
+    $[0] = x;
+  } else {
+    x = $[0];
+  }
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-alias-try-values.js

@@ -0,0 +1,20 @@
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  let y;
+  let x = [];
+  try {
+    // throws x
+    throwInput(x);
+  } catch (e) {
+    // e = x
+    y = e; // y = x
+  }
+  y.push(null);
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-try-value-modified-in-catch-escaping.expect.md

@@ -0,0 +1,63 @@
+
+## Input
+
+```javascript
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  let x;
+  try {
+    const y = [];
+    y.push(props.y);
+    throwInput(y);
+  } catch (e) {
+    e.push(props.e);
+    x = e;
+  }
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ y: "foo", e: "bar" }],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react";
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  const $ = useMemoCache(3);
+  let x;
+  const c_0 = $[0] !== props.y;
+  const c_1 = $[1] !== props.e;
+  if (c_0 || c_1) {
+    try {
+      const y = [];
+      y.push(props.y);
+      throwInput(y);
+    } catch (t30) {
+      const e = t30;
+      e.push(props.e);
+      x = e;
+    }
+    $[0] = props.y;
+    $[1] = props.e;
+    $[2] = x;
+  } else {
+    x = $[2];
+  }
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ y: "foo", e: "bar" }],
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-try-value-modified-in-catch-escaping.js

@@ -0,0 +1,19 @@
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  let x;
+  try {
+    const y = [];
+    y.push(props.y);
+    throwInput(y);
+  } catch (e) {
+    e.push(props.e);
+    x = e;
+  }
+  return x;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ y: "foo", e: "bar" }],
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-try-value-modified-in-catch.expect.md

@@ -0,0 +1,50 @@
+
+## Input
+
+```javascript
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  try {
+    const y = [];
+    y.push(props.y);
+    throwInput(y);
+  } catch (e) {
+    e.push(props.e);
+    return e;
+  }
+  return null;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ y: "foo", e: "bar" }],
+};
+
+```
+
+## Code
+
+```javascript
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  try {
+    const y = [];
+    y.push(props.y);
+    throwInput(y);
+  } catch (t25) {
+    const e = t25;
+    e.push(props.e);
+    return e;
+  }
+  return null;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ y: "foo", e: "bar" }],
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-try-value-modified-in-catch.js

@@ -0,0 +1,18 @@
+const { throwInput } = require("shared-runtime");
+
+function Component(props) {
+  try {
+    const y = [];
+    y.push(props.y);
+    throwInput(y);
+  } catch (e) {
+    e.push(props.e);
+    return e;
+  }
+  return null;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ y: "foo", e: "bar" }],
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/try-catch-with-catch-param.expect.md

@@ -27,25 +27,17 @@ export const FIXTURE_ENTRYPOINT = {
 ## Code
 
 ```javascript
-import { unstable_useMemoCache as useMemoCache } from "react";
 const { throwInput } = require("shared-runtime");
 
 function Component(props) {
-  const $ = useMemoCache(1);
-  let x;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    x = [];
-    try {
-      throwInput(x);
-    } catch (t22) {
-      const e = t22;
+  const x = [];
+  try {
+    throwInput(x);
+  } catch (t22) {
+    const e = t22;
 
-      e.push(null);
-      return e;
-    }
-    $[0] = x;
-  } else {
-    x = $[0];
+    e.push(null);
+    return e;
   }
   return x;
 }