[wip][compiler] Infer optional dependencies

Updates PropagateScopeDeps and DeriveMinimalDeps to understand optional dependency paths (`a?.b`). There a few key pieces to this:

In PropagateScopeDeps we jump through some hoops to work around the awkward structure of nested OptionalExpressions. This is much easier in HIR form, but I managed to get this pretty close and i think it will be landable with further cleanup. A good chunk of this is avoiding prematurely registering a value as a dependency - there are a bunch of indirections in the ReactiveFunction structure:

```
t0 = OptionalExpression
  SequenceExpression
    t0 = Sequence
      ...
    LoadLocal t0
```

Where if at any point we call `visitOperand()` we'll prematurely register a dependency instead of declareProperty(). The other bit is that optionals can be optional=false for nested member expressions where not all the parts are actually optional (`foo.bar?.bar.call()`). And of course, parts of an optional chain can still be conditional even when optional=true (for example the `x` in `foo.bar?.[x]?.baz`). Not all of this is tested yet so there are likely bugs still.

The other bit is DeriveMinimalDeps, which is thankfully easier. We add OptionalAccess and OptionalDep and update the merge and reducing logic for these cases. There is probably still more to update though, for things like merging subtrees. There are a lot of ternaries that assume a result can be exactly one of two states (conditional/unconditional, dependency/access) and these assumptions don't hold anymore. I'd like to refactor to dependency/access separate from conditional/optional/unconditional. Also, the reducing logic isn't quite right: once a child is optional we keep inferring all the parents as optional too, losing some precision. I need to adjust the reducing logic to let children decide whether their path token is optional or not.

ghstack-source-id: 2a8817d
Pull Request resolved: #30819

Author: josephsavona

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

@@ -224,6 +224,14 @@ const EnvironmentConfigSchema = z.object({
 
   enableReactiveScopesInHIR: z.boolean().default(true),
 
+  /**
+   * Enables inference of optional dependency chains. Without this flag
+   * a property chain such as `props?.items?.foo` will infer as a dep on
+   * just `props`. With this flag enabled, we'll infer that full path as
+   * the dependency.
+   */
+  enableOptionalDependencies: z.boolean().default(false),
+
   /*
    * Enable validation of hooks to partially check that the component honors the rules of hooks.
    * When disabled, the component is assumed to follow the rules (though the Babel plugin looks

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

@@ -191,7 +191,7 @@ export function printTerminal(terminal: Terminal): Array<string> | string {
     case 'branch': {
       value = `[${terminal.id}] Branch (${printPlace(terminal.test)}) then:bb${
         terminal.consequent
-      } else:bb${terminal.alternate}`;
+      } else:bb${terminal.alternate} fallthrough:bb${terminal.fallthrough}`;
       break;
     }
     case 'logical': {

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -1446,9 +1446,19 @@ function codegenDependency(
   dependency: ReactiveScopeDependency,
 ): t.Expression {
   let object: t.Expression = convertIdentifier(dependency.identifier);
-  if (dependency.path !== null) {
+  if (dependency.path.length !== 0) {
+    const hasOptional = dependency.path.some(path => path.optional);
     for (const path of dependency.path) {
-      object = t.memberExpression(object, t.identifier(path.property));
+      if (hasOptional) {
+        object = t.optionalMemberExpression(
+          object,
+          t.identifier(path.property),
+          false,
+          path.optional,
+        );
+      } else {
+        object = t.memberExpression(object, t.identifier(path.property));
+      }
     }
   }
   return object;

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/DeriveMinimalDependencies.ts

@@ -60,13 +60,14 @@ export class ReactiveScopeDependencyTree {
     const {path} = dep;
     let currNode = this.#getOrCreateRoot(dep.identifier);
 
-    const accessType = inConditional
-      ? PropertyAccessType.ConditionalAccess
-      : PropertyAccessType.UnconditionalAccess;
-
     for (const item of path) {
       // all properties read 'on the way' to a dependency are marked as 'access'
       let currChild = getOrMakeProperty(currNode, item.property);
+      const accessType = inConditional
+        ? PropertyAccessType.ConditionalAccess
+        : item.optional
+          ? PropertyAccessType.OptionalAccess
+          : PropertyAccessType.UnconditionalAccess;
       currChild.accessType = merge(currChild.accessType, accessType);
       currNode = currChild;
     }
@@ -77,18 +78,22 @@ export class ReactiveScopeDependencyTree {
      */
     const depType = inConditional
       ? PropertyAccessType.ConditionalDependency
-      : PropertyAccessType.UnconditionalDependency;
+      : isOptional(currNode.accessType)
+        ? PropertyAccessType.OptionalDependency
+        : PropertyAccessType.UnconditionalDependency;
 
     currNode.accessType = merge(currNode.accessType, depType);
   }
 
   deriveMinimalDependencies(): Set<ReactiveScopeDependency> {
     const results = new Set<ReactiveScopeDependency>();
     for (const [rootId, rootNode] of this.#roots.entries()) {
-      const deps = deriveMinimalDependenciesInSubtree(rootNode);
+      const deps = deriveMinimalDependenciesInSubtree(rootNode, null);
       CompilerError.invariant(
         deps.every(
-          dep => dep.accessType === PropertyAccessType.UnconditionalDependency,
+          dep =>
+            dep.accessType === PropertyAccessType.UnconditionalDependency ||
+            dep.accessType == PropertyAccessType.OptionalDependency,
         ),
         {
           reason:
@@ -173,6 +178,27 @@ export class ReactiveScopeDependencyTree {
     }
     return res.flat().join('\n');
   }
+
+  debug(): string {
+    const buf: Array<string> = [`tree() [`];
+    for (const [rootId, rootNode] of this.#roots) {
+      buf.push(`${printIdentifier(rootId)} (${rootNode.accessType}):`);
+      this.#debugImpl(buf, rootNode, 1);
+    }
+    buf.push(']');
+    return buf.length > 2 ? buf.join('\n') : buf.join('');
+  }
+
+  #debugImpl(
+    buf: Array<string>,
+    node: DependencyNode,
+    depth: number = 0,
+  ): void {
+    for (const [property, childNode] of node.properties) {
+      buf.push(`${'  '.repeat(depth)}.${property} (${childNode.accessType}):`);
+      this.#debugImpl(buf, childNode, depth + 1);
+    }
+  }
 }
 
 /*
@@ -196,8 +222,10 @@ export class ReactiveScopeDependencyTree {
  */
 enum PropertyAccessType {
   ConditionalAccess = 'ConditionalAccess',
+  OptionalAccess = 'OptionalAccess',
   UnconditionalAccess = 'UnconditionalAccess',
   ConditionalDependency = 'ConditionalDependency',
+  OptionalDependency = 'OptionalDependency',
   UnconditionalDependency = 'UnconditionalDependency',
 }
 
@@ -211,9 +239,16 @@ function isUnconditional(access: PropertyAccessType): boolean {
 function isDependency(access: PropertyAccessType): boolean {
   return (
     access === PropertyAccessType.ConditionalDependency ||
+    access === PropertyAccessType.OptionalDependency ||
     access === PropertyAccessType.UnconditionalDependency
   );
 }
+function isOptional(access: PropertyAccessType): boolean {
+  return (
+    access === PropertyAccessType.OptionalAccess ||
+    access === PropertyAccessType.OptionalDependency
+  );
+}
 
 function merge(
   access1: PropertyAccessType,
@@ -222,6 +257,7 @@ function merge(
   const resultIsUnconditional =
     isUnconditional(access1) || isUnconditional(access2);
   const resultIsDependency = isDependency(access1) || isDependency(access2);
+  const resultIsOptional = isOptional(access1) || isOptional(access2);
 
   /*
    * Straightforward merge.
@@ -237,6 +273,12 @@ function merge(
     } else {
       return PropertyAccessType.UnconditionalAccess;
     }
+  } else if (resultIsOptional) {
+    if (resultIsDependency) {
+      return PropertyAccessType.OptionalDependency;
+    } else {
+      return PropertyAccessType.OptionalAccess;
+    }
   } else {
     if (resultIsDependency) {
       return PropertyAccessType.ConditionalDependency;
@@ -256,19 +298,34 @@ type ReduceResultNode = {
   accessType: PropertyAccessType;
 };
 
-const promoteUncondResult = [
-  {
+function promoteResult(
+  accessType: PropertyAccessType,
+  path: {property: string; optional: boolean} | null,
+): Array<ReduceResultNode> {
+  const result: ReduceResultNode = {
     relativePath: [],
-    accessType: PropertyAccessType.UnconditionalDependency,
-  },
-];
+    accessType,
+  };
+  if (path !== null) {
+    result.relativePath.push(path);
+  }
+  return [result];
+}
 
-const promoteCondResult = [
-  {
-    relativePath: [],
-    accessType: PropertyAccessType.ConditionalDependency,
-  },
-];
+function prependPath(
+  results: Array<ReduceResultNode>,
+  path: {property: string; optional: boolean} | null,
+): Array<ReduceResultNode> {
+  if (path === null) {
+    return results;
+  }
+  return results.map(result => {
+    return {
+      accessType: result.accessType,
+      relativePath: [path, ...result.relativePath],
+    };
+  });
+}
 
 /*
  * Recursively calculates minimal dependencies in a subtree.
@@ -277,42 +334,76 @@ const promoteCondResult = [
  */
 function deriveMinimalDependenciesInSubtree(
   dep: DependencyNode,
+  property: string | null,
 ): Array<ReduceResultNode> {
   const results: Array<ReduceResultNode> = [];
   for (const [childName, childNode] of dep.properties) {
-    const childResult = deriveMinimalDependenciesInSubtree(childNode).map(
-      ({relativePath, accessType}) => {
-        return {
-          relativePath: [
-            {property: childName, optional: false},
-            ...relativePath,
-          ],
-          accessType,
-        };
-      },
+    const childResult = deriveMinimalDependenciesInSubtree(
+      childNode,
+      childName,
     );
     results.push(...childResult);
   }
 
   switch (dep.accessType) {
     case PropertyAccessType.UnconditionalDependency: {
-      return promoteUncondResult;
+      return promoteResult(
+        PropertyAccessType.UnconditionalDependency,
+        property !== null ? {property, optional: false} : null,
+      );
     }
     case PropertyAccessType.UnconditionalAccess: {
       if (
         results.every(
           ({accessType}) =>
-            accessType === PropertyAccessType.UnconditionalDependency,
+            accessType === PropertyAccessType.UnconditionalDependency ||
+            accessType === PropertyAccessType.OptionalDependency,
         )
       ) {
         // all children are unconditional dependencies, return them to preserve granularity
-        return results;
+        return prependPath(
+          results,
+          property !== null ? {property, optional: false} : null,
+        );
       } else {
         /*
          * at least one child is accessed conditionally, so this node needs to be promoted to
          * unconditional dependency
          */
-        return promoteUncondResult;
+        return promoteResult(
+          PropertyAccessType.UnconditionalDependency,
+          property !== null ? {property, optional: false} : null,
+        );
+      }
+    }
+    case PropertyAccessType.OptionalDependency: {
+      return promoteResult(
+        PropertyAccessType.OptionalDependency,
+        property !== null ? {property, optional: true} : null,
+      );
+    }
+    case PropertyAccessType.OptionalAccess: {
+      if (
+        results.every(
+          ({accessType}) =>
+            accessType === PropertyAccessType.UnconditionalDependency ||
+            accessType === PropertyAccessType.OptionalDependency,
+        )
+      ) {
+        // all children are unconditional dependencies, return them to preserve granularity
+        return prependPath(
+          results,
+          property !== null ? {property, optional: true} : null,
+        );
+      } else {
+        /*
+         * at least one child is accessed conditionally, so this node needs to be promoted to
+         * unconditional dependency
+         */
+        return promoteResult(
+          PropertyAccessType.OptionalDependency,
+          property !== null ? {property, optional: true} : null,
+        );
       }
     }
     case PropertyAccessType.ConditionalAccess:
@@ -328,13 +419,19 @@ function deriveMinimalDependenciesInSubtree(
          * unconditional access.
          * Truncate results of child nodes here, since we shouldn't access them anyways
          */
-        return promoteCondResult;
+        return promoteResult(
+          PropertyAccessType.ConditionalDependency,
+          property !== null ? {property, optional: true} : null,
+        );
       } else {
         /*
          * at least one child is accessed unconditionally, so this node can be promoted to
          * unconditional dependency
          */
-        return promoteUncondResult;
+        return promoteResult(
+          PropertyAccessType.UnconditionalDependency,
+          property !== null ? {property, optional: true} : null,
+        );
       }
     }
     default: {

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/MergeReactiveScopesThatInvalidateTogether.ts

@@ -8,7 +8,6 @@
 import {CompilerError} from '..';
 import {
   DeclarationId,
-  DependencyPath,
   InstructionId,
   InstructionKind,
   Place,

compiler/packages/babel-plugin-react-compiler/src/ReactiveScopes/PrintReactiveFunction.ts

@@ -113,7 +113,7 @@ export function printDependency(dependency: ReactiveScopeDependency): string {
   const identifier =
     printIdentifier(dependency.identifier) +
     printType(dependency.identifier.type);
-  return `${identifier}${dependency.path.map(token => `.${token.property}`).join('')}`;
+  return `${identifier}${dependency.path.map(token => `${token.optional ? '?.' : '.'}${token.property}`).join('')}`;
 }
 
 export function printReactiveInstructions(