[compiler][hir-rewrite] Check mutability of base identifier when hoisting (#31032)

mofeiZ · web-flow · commit edacbde73f50 · 2024-10-03T13:55:59.000-04:00
Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom): * #31066 * __->__ #31032 Prior to this PR, we check whether the property load source (e.g. the evaluation of `<base>` in `<base>.property`) is mutable + scoped to determine whether the property load itself is eligible for hoisting. This changes to check the base identifier of the load. - This is needed for the next PR #31066. We want to evaluate whether the base identifier is mutable within the context of the *outermost function*. This is because all LoadLocals and PropertyLoads within a nested function declaration have mutable-ranges within the context of the function, but the base identifier is a context variable. - A side effect is that we no longer infer loads from props / other function arguments as mutable in edge cases (e.g. props escaping out of try-blocks or being assigned to context variables)
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/CollectHoistablePropertyLoads.ts
@@ -15,7 +15,7 @@ import {
   HIRFunction,
   Identifier,
   IdentifierId,
-  InstructionId,
+  InstructionValue,
   ReactiveScopeDependency,
   ScopeId,
 } from './HIR';
@@ -209,33 +209,23 @@ class PropertyPathRegistry {
   }
 }
 
-function addNonNullPropertyPath(
-  source: Identifier,
-  sourceNode: PropertyPathNode,
-  instrId: InstructionId,
-  knownImmutableIdentifiers: Set<IdentifierId>,
-  result: Set<PropertyPathNode>,
-): void {
-  /**
-   * Since this runs *after* buildReactiveScopeTerminals, identifier mutable ranges
-   * are not valid with respect to current instruction id numbering.
-   * We use attached reactive scope ranges as a proxy for mutable range, but this
-   * is an overestimate as (1) scope ranges merge and align to form valid program
-   * blocks and (2) passes like MemoizeFbtAndMacroOperands may assign scopes to
-   * non-mutable identifiers.
-   *
-   * See comment at top of function for why we track known immutable identifiers.
-   */
-  const isMutableAtInstr =
-    source.mutableRange.end > source.mutableRange.start + 1 &&
-    source.scope != null &&
-    inRange({id: instrId}, source.scope.range);
-  if (
-    !isMutableAtInstr ||
-    knownImmutableIdentifiers.has(sourceNode.fullPath.identifier.id)
-  ) {
-    result.add(sourceNode);
+function getMaybeNonNullInInstruction(
+  instr: InstructionValue,
+  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
+  registry: PropertyPathRegistry,
+): PropertyPathNode | null {
+  let path = null;
+  if (instr.kind === 'PropertyLoad') {
+    path = temporaries.get(instr.object.identifier.id) ?? {
+      identifier: instr.object.identifier,
+      path: [],
+    };
+  } else if (instr.kind === 'Destructure') {
+    path = temporaries.get(instr.value.identifier.id) ?? null;
+  } else if (instr.kind === 'ComputedLoad') {
+    path = temporaries.get(instr.object.identifier.id) ?? null;
   }
+  return path != null ? registry.getOrCreateProperty(path) : null;
 }
 
 function collectNonNullsInBlocks(
@@ -286,41 +276,38 @@ function collectNonNullsInBlocks(
       );
     }
     for (const instr of block.instructions) {
-      if (instr.value.kind === 'PropertyLoad') {
-        const source = temporaries.get(instr.value.object.identifier.id) ?? {
-          identifier: instr.value.object.identifier,
-          path: [],
-        };
-        addNonNullPropertyPath(
-          instr.value.object.identifier,
-          registry.getOrCreateProperty(source),
-          instr.id,
-          knownImmutableIdentifiers,
-          assumedNonNullObjects,
-        );
-      } else if (instr.value.kind === 'Destructure') {
-        const source = instr.value.value.identifier.id;
-        const sourceNode = temporaries.get(source);
-        if (sourceNode != null) {
-          addNonNullPropertyPath(
-            instr.value.value.identifier,
-            registry.getOrCreateProperty(sourceNode),
-            instr.id,
-            knownImmutableIdentifiers,
-            assumedNonNullObjects,
-          );
-        }
-      } else if (instr.value.kind === 'ComputedLoad') {
-        const source = instr.value.object.identifier.id;
-        const sourceNode = temporaries.get(source);
-        if (sourceNode != null) {
-          addNonNullPropertyPath(
-            instr.value.object.identifier,
-            registry.getOrCreateProperty(sourceNode),
-            instr.id,
-            knownImmutableIdentifiers,
-            assumedNonNullObjects,
+      const maybeNonNull = getMaybeNonNullInInstruction(
+        instr.value,
+        temporaries,
+        registry,
+      );
+      if (maybeNonNull != null) {
+        const baseIdentifier = maybeNonNull.fullPath.identifier;
+        /**
+         * Since this runs *after* buildReactiveScopeTerminals, identifier mutable ranges
+         * are not valid with respect to current instruction id numbering.
+         * We use attached reactive scope ranges as a proxy for mutable range, but this
+         * is an overestimate as (1) scope ranges merge and align to form valid program
+         * blocks and (2) passes like MemoizeFbtAndMacroOperands may assign scopes to
+         * non-mutable identifiers.
+         *
+         * See comment at top of function for why we track known immutable identifiers.
+         */
+        const isMutableAtInstr =
+          baseIdentifier.mutableRange.end >
+            baseIdentifier.mutableRange.start + 1 &&
+          baseIdentifier.scope != null &&
+          inRange(
+            {
+              id: instr.id,
+            },
+            baseIdentifier.scope.range,
           );
+        if (
+          !isMutableAtInstr ||
+          knownImmutableIdentifiers.has(baseIdentifier.id)
+        ) {
+          assumedNonNullObjects.add(maybeNonNull);
         }
       }
     }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-try-catch-maybe-null-dependency.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-try-catch-maybe-null-dependency.expect.md
@@ -0,0 +1,65 @@
+
+## Input
+
+```javascript
+import {identity} from 'shared-runtime';
+
+/**
+ * Not safe to hoist read of maybeNullObject.value.inner outside of the
+ * try-catch block, as that might throw
+ */
+function useFoo(maybeNullObject: {value: {inner: number}} | null) {
+  const y = [];
+  try {
+    y.push(identity(maybeNullObject.value.inner));
+  } catch {
+    y.push('null');
+  }
+
+  return y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, {value: 2}, {value: 3}, null],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { identity } from "shared-runtime";
+
+/**
+ * Not safe to hoist read of maybeNullObject.value.inner outside of the
+ * try-catch block, as that might throw
+ */
+function useFoo(maybeNullObject) {
+  const $ = _c(2);
+  let y;
+  if ($[0] !== maybeNullObject.value.inner) {
+    y = [];
+    try {
+      y.push(identity(maybeNullObject.value.inner));
+    } catch {
+      y.push("null");
+    }
+    $[0] = maybeNullObject.value.inner;
+    $[1] = y;
+  } else {
+    y = $[1];
+  }
+  return y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, { value: 2 }, { value: 3 }, null],
+};
+
+```
+      
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-try-catch-maybe-null-dependency.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-try-catch-maybe-null-dependency.ts
@@ -0,0 +1,22 @@
+import {identity} from 'shared-runtime';
+
+/**
+ * Not safe to hoist read of maybeNullObject.value.inner outside of the
+ * try-catch block, as that might throw
+ */
+function useFoo(maybeNullObject: {value: {inner: number}} | null) {
+  const y = [];
+  try {
+    y.push(identity(maybeNullObject.value.inner));
+  } catch {
+    y.push('null');
+  }
+
+  return y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, {value: 2}, {value: 3}, null],
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/try-catch-maybe-null-dependency.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/try-catch-maybe-null-dependency.expect.md
@@ -0,0 +1,79 @@
+
+## Input
+
+```javascript
+// @enablePropagateDepsInHIR
+import {identity} from 'shared-runtime';
+
+/**
+ * Not safe to hoist read of maybeNullObject.value.inner outside of the
+ * try-catch block, as that might throw
+ */
+function useFoo(maybeNullObject: {value: {inner: number}} | null) {
+  const y = [];
+  try {
+    y.push(identity(maybeNullObject.value.inner));
+  } catch {
+    y.push('null');
+  }
+
+  return y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, {value: 2}, {value: 3}, null],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
+import { identity } from "shared-runtime";
+
+/**
+ * Not safe to hoist read of maybeNullObject.value.inner outside of the
+ * try-catch block, as that might throw
+ */
+function useFoo(maybeNullObject) {
+  const $ = _c(4);
+  let y;
+  if ($[0] !== maybeNullObject) {
+    y = [];
+    try {
+      let t0;
+      if ($[2] !== maybeNullObject.value.inner) {
+        t0 = identity(maybeNullObject.value.inner);
+        $[2] = maybeNullObject.value.inner;
+        $[3] = t0;
+      } else {
+        t0 = $[3];
+      }
+      y.push(t0);
+    } catch {
+      y.push("null");
+    }
+    $[0] = maybeNullObject;
+    $[1] = y;
+  } else {
+    y = $[1];
+  }
+  return y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, { value: 2 }, { value: 3 }, null],
+};
+
+```
+      
+### Eval output
+(kind: ok) ["null"]
+[null]
+[null]
+["null"]
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/try-catch-maybe-null-dependency.ts b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/try-catch-maybe-null-dependency.ts
@@ -0,0 +1,23 @@
+// @enablePropagateDepsInHIR
+import {identity} from 'shared-runtime';
+
+/**
+ * Not safe to hoist read of maybeNullObject.value.inner outside of the
+ * try-catch block, as that might throw
+ */
+function useFoo(maybeNullObject: {value: {inner: number}} | null) {
+  const y = [];
+  try {
+    y.push(identity(maybeNullObject.value.inner));
+  } catch {
+    y.push('null');
+  }
+
+  return y;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [null],
+  sequentialRenders: [null, {value: 2}, {value: 3}, null],
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/try-catch-try-value-modified-in-catch-escaping.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/try-catch-try-value-modified-in-catch-escaping.expect.md
@@ -32,9 +32,9 @@ import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
 const { throwInput } = require("shared-runtime");
 
 function Component(props) {
-  const $ = _c(2);
+  const $ = _c(3);
   let x;
-  if ($[0] !== props) {
+  if ($[0] !== props.y || $[1] !== props.e) {
     try {
       const y = [];
       y.push(props.y);
@@ -44,10 +44,11 @@ function Component(props) {
       e.push(props.e);
       x = e;
     }
-    $[0] = props;
-    $[1] = x;
+    $[0] = props.y;
+    $[1] = props.e;
+    $[2] = x;
   } else {
-    x = $[1];
+    x = $[2];
   }
   return x;
 }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/try-catch-try-value-modified-in-catch.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/try-catch-try-value-modified-in-catch.expect.md
@@ -31,9 +31,9 @@ import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
 const { throwInput } = require("shared-runtime");
 
 function Component(props) {
-  const $ = _c(2);
+  const $ = _c(3);
   let t0;
-  if ($[0] !== props) {
+  if ($[0] !== props.y || $[1] !== props.e) {
     t0 = Symbol.for("react.early_return_sentinel");
     bb0: {
       try {
@@ -47,10 +47,11 @@ function Component(props) {
         break bb0;
       }
     }
-    $[0] = props;
-    $[1] = t0;
+    $[0] = props.y;
+    $[1] = props.e;
+    $[2] = t0;
   } else {
-    t0 = $[1];
+    t0 = $[2];
   }
   if (t0 !== Symbol.for("react.early_return_sentinel")) {
     return t0;
diff --git a/compiler/packages/snap/src/SproutTodoFilter.ts b/compiler/packages/snap/src/SproutTodoFilter.ts
@@ -478,6 +478,7 @@ const skipFilter = new Set([
   'fbt/bug-fbt-plural-multiple-function-calls',
   'fbt/bug-fbt-plural-multiple-mixed-call-tag',
   'bug-invalid-hoisting-functionexpr',
+  'bug-try-catch-maybe-null-dependency',
   'reduce-reactive-deps/bug-merge-uncond-optional-chain-and-cond',
   'original-reactive-scopes-fork/bug-nonmutating-capture-in-unsplittable-memo-block',
   'original-reactive-scopes-fork/bug-hoisted-declaration-with-scope',
