Update

[ghstack-poisoned]

compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts

@@ -343,7 +343,7 @@ function* runWithEnvironment(
   });
   assertTerminalSuccessorsExist(hir);
   assertTerminalPredsExist(hir);
-  if (env.config.enablePropagateDepsInHIR) {
+  if (env.config.enablePropagateDepsInHIR !== 'disabled') {
     propagateScopeDependenciesHIR(hir);
     yield log({
       kind: 'hir',
@@ -378,7 +378,7 @@ function* runWithEnvironment(
   });
   assertScopeInstructionsWithinScopes(reactiveFunction);
 
-  if (!env.config.enablePropagateDepsInHIR) {
+  if (env.config.enablePropagateDepsInHIR === 'disabled') {
     propagateScopeDependencies(reactiveFunction);
     yield log({
       kind: 'reactive',

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,10 +225,44 @@ class Tree {
   }
 }
 
-function collectPropertyLoadsInBlocks(
+function pushPropertyLoadNode(
+  node: PropertyLoadNode,
+  instrId: InstructionId,
+  knownImmutableIdentifiers: Set<IdentifierId>,
+  result: Set<PropertyLoadNode>,
+): void {
+  const object = node.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 =
+    object.mutableRange.end > object.mutableRange.start + 1 &&
+    object.scope != null &&
+    inRange({id: instrId}, object.scope.range);
+  if (
+    !isMutableAtInstr ||
+    knownImmutableIdentifiers.has(node.fullPath.identifier.id)
+  ) {
+    let curr: PropertyLoadNode | null = node;
+    while (curr != null) {
+      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
@@ -238,53 +271,75 @@ function collectPropertyLoadsInBlocks(
    * We track known immutable identifiers to reduce regressions (as PropagateScopeDeps
    * is being rewritten to HIR).
    */
-  const knownImmutableIdentifiers = new Set<Identifier>();
+  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);
+        knownImmutableIdentifiers.add(p.identifier.id);
       }
     }
   }
-  const tree = new Tree();
+  /**
+   * Known non-null objects such as functional component props can be safely
+   * read from any block.
+   */
+  const knownNonNullIdentifiers = new Set<PropertyLoadNode>();
+  if (
+    fn.env.config.enablePropagateDepsInHIR === 'enabled_with_optimizations' &&
+    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>();
+    const assumedNonNullObjects = new Set<PropertyLoadNode>(
+      knownNonNullIdentifiers,
+    );
     for (const instr of block.instructions) {
       if (instr.value.kind === 'PropertyLoad') {
-        const property = getProperty(
-          instr.value.object,
-          instr.value.property,
-          temporaries,
+        const source = temporaries.get(instr.value.object.identifier.id) ?? {
+          identifier: instr.value.object.identifier,
+          path: [],
+        };
+        const propertyNode = tree.getPropertyLoadNode(source);
+        pushPropertyLoadNode(
+          propertyNode,
+          instr.id,
+          knownImmutableIdentifiers,
+          assumedNonNullObjects,
         );
-        const propertyNode = tree.getPropertyLoadNode(property);
-        const object = propertyNode.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 =
-          object.mutableRange.end > object.mutableRange.start + 1 &&
-          object.scope != null &&
-          inRange(instr, object.scope.range);
-        if (
-          !isMutableAtInstr ||
-          knownImmutableIdentifiers.has(propertyNode.fullPath.identifier)
-        ) {
-          let curr = propertyNode.parent;
-          while (curr != null) {
-            assumedNonNullObjects.add(curr);
-            curr = curr.parent;
-          }
+      } else if (
+        instr.value.kind === 'Destructure' &&
+        fn.env.config.enablePropagateDepsInHIR === 'enabled_with_optimizations'
+      ) {
+        const source = instr.value.value.identifier.id;
+        const sourceNode = temporaries.get(source);
+        if (sourceNode != null) {
+          pushPropertyLoadNode(
+            tree.getPropertyLoadNode(sourceNode),
+            instr.id,
+            knownImmutableIdentifiers,
+            assumedNonNullObjects,
+          );
+        }
+      } else if (
+        instr.value.kind === 'ComputedLoad' &&
+        fn.env.config.enablePropagateDepsInHIR === 'enabled_with_optimizations'
+      ) {
+        const source = instr.value.object.identifier.id;
+        const sourceNode = temporaries.get(source);
+        if (sourceNode != null) {
+          pushPropertyLoadNode(
+            tree.getPropertyLoadNode(sourceNode),
+            instr.id,
+            knownImmutableIdentifiers,
+            assumedNonNullObjects,
+          );
         }
       }
-      // TODO handle destructuring
     }
 
     nodes.set(block.id, {
@@ -449,10 +504,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/HIR/Environment.ts

@@ -230,7 +230,9 @@ const EnvironmentConfigSchema = z.object({
    */
   enableUseTypeAnnotations: z.boolean().default(false),
 
-  enablePropagateDepsInHIR: z.boolean().default(false),
+  enablePropagateDepsInHIR: z
+    .enum(['disabled', 'enabled_baseline', 'enabled_with_optimizations'])
+    .default('disabled'),
 
   /**
    * Enables inference of optional dependency chains. Without this flag

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:enabled_with_optimizations
+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:enabled_with_optimizations
+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:enabled_with_optimizations
+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}],
+};