[compiler][hir-rewrite] Infer non-null props, destructure source

Followup from #30894.
This adds a new flagged mode `enablePropagateScopeDepsInHIR: "enabled_with_optimizations"`, under which we infer more hoistable loads:
- it's always safe to evaluate loads from `props` (i.e. first parameter of a `component`)
- destructuring sources are safe to evaluate loads from (e.g. given `{x} = obj`, we infer that it's safe to evaluate obj.y)
- computed load sources are safe to evaluate loads from (e.g. given `arr[0]`, we can infer that it's safe to evaluate arr.length)

ghstack-source-id: 32f3bb72e9f85922825579bd785d636f4ccf724d
Pull Request resolved: https://github.com/facebook/react/pull/31033

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

@@ -8,6 +8,7 @@ import {
   HIRFunction,
   Identifier,
   IdentifierId,
+  InstructionId,
   Place,
   ReactiveScopeDependency,
   ScopeId,
@@ -66,7 +67,7 @@ export function collectHoistablePropertyLoads(
   fn: HIRFunction,
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
 ): ReadonlyMap<ScopeId, BlockInfo> {
-  const nodes = collectPropertyLoadsInBlocks(fn, temporaries);
+  const nodes = collectNonNullsInBlocks(fn, temporaries);
   propagateNonNull(fn, nodes);
 
   const nodesKeyedByScopeId = new Map<ScopeId, BlockInfo>();
@@ -165,7 +166,7 @@ type PropertyLoadNode =
 class Tree {
   roots: Map<Identifier, RootNode> = new Map();
 
-  #getOrCreateRoot(identifier: Identifier): PropertyLoadNode {
+  getOrCreateRoot(identifier: Identifier): PropertyLoadNode {
     /**
      * Reads from a statically scoped variable are always safe in JS,
      * with the exception of TDZ (not addressed by this pass).
@@ -207,17 +208,15 @@ class Tree {
   }
 
   getPropertyLoadNode(n: ReactiveScopeDependency): PropertyLoadNode {
-    CompilerError.invariant(n.path.length > 0, {
-      reason:
-        '[CollectHoistablePropertyLoads] Expected property node, found root node',
-      loc: GeneratedSource,
-    });
     /**
      * We add ReactiveScopeDependencies according to instruction ordering,
      * so all subpaths of a PropertyLoad should already exist
      * (e.g. a.b is added before a.b.c),
      */
-    let currNode = this.#getOrCreateRoot(n.identifier);
+    let currNode = this.getOrCreateRoot(n.identifier);
+    if (n.path.length === 0) {
+      return currNode;
+    }
     for (let i = 0; i < n.path.length - 1; i++) {
       currNode = assertNonNull(currNode.properties.get(n.path[i].property));
     }
@@ -226,39 +225,13 @@ class Tree {
   }
 }
 
-function collectPropertyLoadsInBlocks(
-  fn: HIRFunction,
-  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
-): ReadonlyMap<BlockId, BlockInfo> {
-  /**
-   * Due to current limitations of mutable range inference, there are edge cases in
-   * which we infer known-immutable values (e.g. props or hook params) to have a
-   * mutable range and scope.
-   * (see `destructure-array-declaration-to-context-var` fixture)
-   * We track known immutable identifiers to reduce regressions (as PropagateScopeDeps
-   * is being rewritten to HIR).
-   */
-  const knownImmutableIdentifiers = new Set<Identifier>();
-  if (fn.fnType === 'Component' || fn.fnType === 'Hook') {
-    for (const p of fn.params) {
-      if (p.kind === 'Identifier') {
-        knownImmutableIdentifiers.add(p.identifier);
-      }
-    }
-  }
-  const tree = new Tree();
-  const nodes = new Map<BlockId, BlockInfo>();
-  for (const [_, block] of fn.body.blocks) {
-    const assumedNonNullObjects = new Set<PropertyLoadNode>();
-    for (const instr of block.instructions) {
-      if (instr.value.kind === 'PropertyLoad') {
-        const property = getProperty(
-          instr.value.object,
-          instr.value.property,
-          temporaries,
-        );
-        const propertyNode = tree.getPropertyLoadNode(property);
-        const object = instr.value.object.identifier;
+function pushPropertyLoadNode(
+  loadSource: Identifier,
+  loadSourceNode: PropertyLoadNode,
+  instrId: InstructionId,
+  knownImmutableIdentifiers: Set<IdentifierId>,
+  result: Set<PropertyLoadNode>,
+): void {
   /**
    * Since this runs *after* buildReactiveScopeTerminals, identifier mutable ranges
    * are not valid with respect to current instruction id numbering.
@@ -270,21 +243,98 @@ function collectPropertyLoadsInBlocks(
    * See comment at top of function for why we track known immutable identifiers.
    */
   const isMutableAtInstr =
-          object.mutableRange.end > object.mutableRange.start + 1 &&
-          object.scope != null &&
-          inRange(instr, object.scope.range);
+    loadSource.mutableRange.end > loadSource.mutableRange.start + 1 &&
+    loadSource.scope != null &&
+    inRange({id: instrId}, loadSource.scope.range);
   if (
     !isMutableAtInstr ||
-          knownImmutableIdentifiers.has(propertyNode.fullPath.identifier)
+    knownImmutableIdentifiers.has(loadSourceNode.fullPath.identifier.id)
   ) {
-          let curr = propertyNode.parent;
+    let curr: PropertyLoadNode | null = loadSourceNode;
     while (curr != null) {
-            assumedNonNullObjects.add(curr);
+      result.add(curr);
       curr = curr.parent;
     }
   }
+}
+
+function collectNonNullsInBlocks(
+  fn: HIRFunction,
+  temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
+): ReadonlyMap<BlockId, BlockInfo> {
+  const tree = new Tree();
+  /**
+   * Due to current limitations of mutable range inference, there are edge cases in
+   * which we infer known-immutable values (e.g. props or hook params) to have a
+   * mutable range and scope.
+   * (see `destructure-array-declaration-to-context-var` fixture)
+   * We track known immutable identifiers to reduce regressions (as PropagateScopeDeps
+   * is being rewritten to HIR).
+   */
+  const knownImmutableIdentifiers = new Set<IdentifierId>();
+  if (fn.fnType === 'Component' || fn.fnType === 'Hook') {
+    for (const p of fn.params) {
+      if (p.kind === 'Identifier') {
+        knownImmutableIdentifiers.add(p.identifier.id);
+      }
+    }
+  }
+  /**
+   * Known non-null objects such as functional component props can be safely
+   * read from any block.
+   */
+  const knownNonNullIdentifiers = new Set<PropertyLoadNode>();
+  if (
+    fn.fnType === 'Component' &&
+    fn.params.length > 0 &&
+    fn.params[0].kind === 'Identifier'
+  ) {
+    const identifier = fn.params[0].identifier;
+    knownNonNullIdentifiers.add(tree.getOrCreateRoot(identifier));
+  }
+  const nodes = new Map<BlockId, BlockInfo>();
+  for (const [_, block] of fn.body.blocks) {
+    const assumedNonNullObjects = new Set<PropertyLoadNode>(
+      knownNonNullIdentifiers,
+    );
+    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: [],
+        };
+        pushPropertyLoadNode(
+          instr.value.object.identifier,
+          tree.getPropertyLoadNode(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) {
+          pushPropertyLoadNode(
+            instr.value.value.identifier,
+            tree.getPropertyLoadNode(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) {
+          pushPropertyLoadNode(
+            instr.value.object.identifier,
+            tree.getPropertyLoadNode(sourceNode),
+            instr.id,
+            knownImmutableIdentifiers,
+            assumedNonNullObjects,
+          );
+        }
       }
-      // TODO handle destructuring
     }
 
     nodes.set(block.id, {
@@ -449,10 +499,11 @@ function propagateNonNull(
   );
 
   for (const [id, node] of nodes) {
-    node.assumedNonNullObjects = Set_union(
+    const assumedNonNullObjects = Set_union(
       assertNonNull(fromEntry.get(id)),
       assertNonNull(fromExit.get(id)),
     );
+    node.assumedNonNullObjects = assumedNonNullObjects;
   }
 }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/infer-component-props-non-null.expect.md

@@ -0,0 +1,60 @@
+
+## Input
+
+```javascript
+// @enablePropagateDepsInHIR
+import {identity, Stringify} from 'shared-runtime';
+
+function Foo(props) {
+  /**
+   * props.value should be inferred as the dependency of this scope
+   * since we know that props is safe to read from (i.e. non-null)
+   * as it is arg[0] of a component function
+   */
+  const arr = [];
+  if (cond) {
+    arr.push(identity(props.value));
+  }
+  return <Stringify arr={arr} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{value: 2}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
+import { identity, Stringify } from "shared-runtime";
+
+function Foo(props) {
+  const $ = _c(2);
+  let t0;
+  if ($[0] !== props.value) {
+    const arr = [];
+    if (cond) {
+      arr.push(identity(props.value));
+    }
+
+    t0 = <Stringify arr={arr} />;
+    $[0] = props.value;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ value: 2 }],
+};
+
+```
+      
+### Eval output
+(kind: exception) cond is not defined

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/infer-component-props-non-null.tsx

@@ -0,0 +1,20 @@
+// @enablePropagateDepsInHIR
+import {identity, Stringify} from 'shared-runtime';
+
+function Foo(props) {
+  /**
+   * props.value should be inferred as the dependency of this scope
+   * since we know that props is safe to read from (i.e. non-null)
+   * as it is arg[0] of a component function
+   */
+  const arr = [];
+  if (cond) {
+    arr.push(identity(props.value));
+  }
+  return <Stringify arr={arr} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{value: 2}],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/infer-non-null-destructure.expect.md

@@ -0,0 +1,91 @@
+
+## Input
+
+```javascript
+// @enablePropagateDepsInHIR
+import {identity, useIdentity} from 'shared-runtime';
+
+function useFoo({arg, cond}: {arg: number; cond: boolean}) {
+  const maybeObj = useIdentity({value: arg});
+  const {value} = maybeObj;
+  useIdentity(null);
+  /**
+   * maybeObj.value should be inferred as the dependency of this scope
+   * since we know that maybeObj is safe to read from (i.e. non-null)
+   * due to the above destructuring instruction
+   */
+  const arr = [];
+  if (cond) {
+    arr.push(identity(maybeObj.value));
+  }
+  return {arr, value};
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{arg: 2, cond: false}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime"; // @enablePropagateDepsInHIR
+import { identity, useIdentity } from "shared-runtime";
+
+function useFoo(t0) {
+  const $ = _c(10);
+  const { arg, cond } = t0;
+  let t1;
+  if ($[0] !== arg) {
+    t1 = { value: arg };
+    $[0] = arg;
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const maybeObj = useIdentity(t1);
+  const { value } = maybeObj;
+  useIdentity(null);
+  let arr;
+  if ($[2] !== cond || $[3] !== maybeObj.value) {
+    arr = [];
+    if (cond) {
+      let t2;
+      if ($[5] !== maybeObj.value) {
+        t2 = identity(maybeObj.value);
+        $[5] = maybeObj.value;
+        $[6] = t2;
+      } else {
+        t2 = $[6];
+      }
+      arr.push(t2);
+    }
+    $[2] = cond;
+    $[3] = maybeObj.value;
+    $[4] = arr;
+  } else {
+    arr = $[4];
+  }
+  let t2;
+  if ($[7] !== arr || $[8] !== value) {
+    t2 = { arr, value };
+    $[7] = arr;
+    $[8] = value;
+    $[9] = t2;
+  } else {
+    t2 = $[9];
+  }
+  return t2;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{ arg: 2, cond: false }],
+};
+
+```
+      
+### Eval output
+(kind: ok) {"arr":[],"value":2}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/propagate-scope-deps-hir-fork/infer-non-null-destructure.ts

@@ -0,0 +1,23 @@
+// @enablePropagateDepsInHIR
+import {identity, useIdentity} from 'shared-runtime';
+
+function useFoo({arg, cond}: {arg: number; cond: boolean}) {
+  const maybeObj = useIdentity({value: arg});
+  const {value} = maybeObj;
+  useIdentity(null);
+  /**
+   * maybeObj.value should be inferred as the dependency of this scope
+   * since we know that maybeObj is safe to read from (i.e. non-null)
+   * due to the above destructuring instruction
+   */
+  const arr = [];
+  if (cond) {
+    arr.push(identity(maybeObj.value));
+  }
+  return {arr, value};
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{arg: 2, cond: false}],
+};