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

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(
+ 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.
+ * 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 =
+ loadSource.mutableRange.end > loadSource.mutableRange.start + 1 &&
+ loadSource.scope != null &&
+ inRange({id: instrId}, loadSource.scope.range);
+ if (
+ !isMutableAtInstr ||
+ knownImmutableIdentifiers.has(loadSourceNode.fullPath.identifier.id)
+ ) {
+ let curr: PropertyLoadNode | null = loadSourceNode;
+ 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,77 @@ 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: [],
+ };
+ pushPropertyLoadNode(
+ instr.value.object.identifier,
+ tree.getPropertyLoadNode(source),
+ instr.id,
+ knownImmutableIdentifiers,
+ assumedNonNullObjects,
  );
- const propertyNode = tree.getPropertyLoadNode(property);
- const object = instr.value.object.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(
+ instr.value.value.identifier,
+ 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(
+ instr.value.object.identifier,
+ tree.getPropertyLoadNode(sourceNode),
+ instr.id,
+ knownImmutableIdentifiers,
+ assumedNonNullObjects,
+ );
  }
  }
- // TODO handle destructuring
  }
 
  nodes.set(block.id, {
@@ -449,10 +506,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}],
+};