[compiler][hir-rewrite] Check mutability of base identifier when hoisting

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

@@ -1,6 +1,12 @@
 import {CompilerError} from '../CompilerError';
 import {inRange} from '../ReactiveScopes/InferReactiveScopeVariables';
-import {Set_intersect, Set_union, getOrInsertDefault} from '../Utils/utils';
+import {
+  Set_equal,
+  Set_filter,
+  Set_intersect,
+  Set_union,
+  getOrInsertDefault,
+} from '../Utils/utils';
 import {
   BasicBlock,
   BlockId,
@@ -9,15 +15,15 @@ import {
   HIRFunction,
   Identifier,
   IdentifierId,
-  InstructionId,
+  InstructionValue,
   ReactiveScopeDependency,
   ScopeId,
 } from './HIR';
 
 /**
- * Helper function for `PropagateScopeDependencies`.
- * Uses control flow graph analysis to determine which `Identifier`s can
- * be assumed to be non-null objects, on a per-block basis.
+ * Helper function for `PropagateScopeDependencies`. Uses control flow graph
+ * analysis to determine which `Identifier`s can be assumed to be non-null
+ * objects, on a per-block basis.
  *
  * Here is an example:
  * ```js
@@ -42,15 +48,16 @@ import {
  * }
  * ```
  *
- * Note that we currently do NOT account for mutable / declaration range
- * when doing the CFG-based traversal, producing results that are technically
+ * Note that we currently do NOT account for mutable / declaration range when
+ * doing the CFG-based traversal, producing results that are technically
  * incorrect but filtered by PropagateScopeDeps (which only takes dependencies
  * on constructed value -- i.e. a scope's dependencies must have mutable ranges
  * ending earlier than the scope start).
  *
- * Take this example, this function will infer x.foo.bar as non-nullable for bb0,
- * via the intersection of bb1 & bb2 which in turn comes from bb3. This is technically
- * incorrect bb0 is before / during x's mutable range.
+ * Take this example, this function will infer x.foo.bar as non-nullable for
+ * bb0, via the intersection of bb1 & bb2 which in turn comes from bb3. This is
+ * technically incorrect bb0 is before / during x's mutable range.
+ * ```
  *  bb0:
  *    const x = ...;
  *    if cond then bb1 else bb2
@@ -62,15 +69,30 @@ import {
  *    goto bb3:
  *  bb3:
  *    x.foo.bar
+ * ```
+ *
+ * @param fn
+ * @param temporaries sidemap of identifier -> baseObject.a.b paths. Does not
+ * contain optional chains.
+ * @param hoistableFromOptionals sidemap of optionalBlock -> baseObject?.a
+ * optional paths for which it's safe to evaluate non-optional loads (see
+ * CollectOptionalChainDependencies).
+ * @returns
  */
 export function collectHoistablePropertyLoads(
   fn: HIRFunction,
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
+  hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
 ): ReadonlyMap<ScopeId, BlockInfo> {
   const registry = new PropertyPathRegistry();
 
-  const nodes = collectNonNullsInBlocks(fn, temporaries, registry);
-  propagateNonNull(fn, nodes);
+  const nodes = collectNonNullsInBlocks(
+    fn,
+    temporaries,
+    hoistableFromOptionals,
+    registry,
+  );
+  propagateNonNull(fn, nodes, registry);
 
   const nodesKeyedByScopeId = new Map<ScopeId, BlockInfo>();
   for (const [_, block] of fn.body.blocks) {
@@ -96,17 +118,21 @@ export type BlockInfo = {
  */
 type RootNode = {
   properties: Map<string, PropertyPathNode>;
+  optionalProperties: Map<string, PropertyPathNode>;
   parent: null;
   // Recorded to make later computations simpler
   fullPath: ReactiveScopeDependency;
+  hasOptional: boolean;
   root: IdentifierId;
 };
 
 type PropertyPathNode =
   | {
       properties: Map<string, PropertyPathNode>;
+      optionalProperties: Map<string, PropertyPathNode>;
       parent: PropertyPathNode;
       fullPath: ReactiveScopeDependency;
+      hasOptional: boolean;
     }
   | RootNode;
 
@@ -124,10 +150,12 @@ class PropertyPathRegistry {
       rootNode = {
         root: identifier.id,
         properties: new Map(),
+        optionalProperties: new Map(),
         fullPath: {
           identifier,
           path: [],
         },
+        hasOptional: false,
         parent: null,
       };
       this.roots.set(identifier.id, rootNode);
@@ -139,23 +167,20 @@ class PropertyPathRegistry {
     parent: PropertyPathNode,
     entry: DependencyPathEntry,
   ): PropertyPathNode {
-    if (entry.optional) {
-      CompilerError.throwTodo({
-        reason: 'handle optional nodes',
-        loc: GeneratedSource,
-      });
-    }
-    let child = parent.properties.get(entry.property);
+    const map = entry.optional ? parent.optionalProperties : parent.properties;
+    let child = map.get(entry.property);
     if (child == null) {
       child = {
         properties: new Map(),
+        optionalProperties: new Map(),
         parent: parent,
         fullPath: {
           identifier: parent.fullPath.identifier,
           path: parent.fullPath.path.concat(entry),
         },
+        hasOptional: parent.hasOptional || entry.optional,
       };
-      parent.properties.set(entry.property, child);
+      map.set(entry.property, child);
     }
     return child;
   }
@@ -184,38 +209,29 @@ 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(
   fn: HIRFunction,
   temporaries: ReadonlyMap<IdentifierId, ReactiveScopeDependency>,
+  hoistableFromOptionals: ReadonlyMap<BlockId, ReactiveScopeDependency>,
   registry: PropertyPathRegistry,
 ): ReadonlyMap<BlockId, BlockInfo> {
   /**
@@ -252,57 +268,62 @@ function collectNonNullsInBlocks(
     const assumedNonNullObjects = new Set<PropertyPathNode>(
       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: [],
-        };
-        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,
-          );
-        }
-      }
-    }
 
     nodes.set(block.id, {
       block,
       assumedNonNullObjects,
     });
+    const maybeOptionalChain = hoistableFromOptionals.get(block.id);
+    if (maybeOptionalChain != null) {
+      assumedNonNullObjects.add(
+        registry.getOrCreateProperty(maybeOptionalChain),
+      );
+      continue;
+    }
+    for (const instr of block.instructions) {
+      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);
+        }
+      }
+    }
   }
   return nodes;
 }
 
 function propagateNonNull(
   fn: HIRFunction,
   nodes: ReadonlyMap<BlockId, BlockInfo>,
+  registry: PropertyPathRegistry,
 ): void {
   const blockSuccessors = new Map<BlockId, Set<BlockId>>();
   const terminalPreds = new Set<BlockId>();
@@ -388,10 +409,17 @@ function propagateNonNull(
 
     const prevObjects = assertNonNull(nodes.get(nodeId)).assumedNonNullObjects;
     const mergedObjects = Set_union(prevObjects, neighborAccesses);
+    reduceMaybeOptionalChains(mergedObjects, registry);
 
     assertNonNull(nodes.get(nodeId)).assumedNonNullObjects = mergedObjects;
     traversalState.set(nodeId, 'done');
-    changed ||= prevObjects.size !== mergedObjects.size;
+    /**
+     * Note that it's not sufficient to compare set sizes since
+     * reduceMaybeOptionalChains may replace optional-chain loads with
+     * unconditional loads. This could in turn change `assumedNonNullObjects` of
+     * downstream blocks and backedges.
+     */
+    changed ||= !Set_equal(prevObjects, mergedObjects);
     return changed;
   }
   const traversalState = new Map<BlockId, 'done' | 'active'>();
@@ -440,3 +468,50 @@ export function assertNonNull<T extends NonNullable<U>, U>(
   });
   return value;
 }
+
+/**
+ * Any two optional chains with different operations . vs ?. but the same set of
+ * property strings paths de-duplicates.
+ *
+ * Intuitively: given <base>?.b, we know <base> to be either hoistable or not.
+ * If unconditional reads from <base> are hoistable, we can replace all
+ * <base>?.PROPERTY_STRING subpaths with <base>.PROPERTY_STRING
+ */
+function reduceMaybeOptionalChains(
+  nodes: Set<PropertyPathNode>,
+  registry: PropertyPathRegistry,
+): void {
+  let optionalChainNodes = Set_filter(nodes, n => n.hasOptional);
+  if (optionalChainNodes.size === 0) {
+    return;
+  }
+  let changed: boolean;
+  do {
+    changed = false;
+
+    for (const original of optionalChainNodes) {
+      let {identifier, path: origPath} = original.fullPath;
+      let currNode: PropertyPathNode =
+        registry.getOrCreateIdentifier(identifier);
+      for (let i = 0; i < origPath.length; i++) {
+        const entry = origPath[i];
+        // If the base is known to be non-null, replace with a non-optional load
+        const nextEntry: DependencyPathEntry =
+          entry.optional && nodes.has(currNode)
+            ? {property: entry.property, optional: false}
+            : entry;
+        currNode = PropertyPathRegistry.getOrCreatePropertyEntry(
+          currNode,
+          nextEntry,
+        );
+      }
+      if (currNode !== original) {
+        changed = true;
+        optionalChainNodes.delete(original);
+        optionalChainNodes.add(currNode);
+        nodes.delete(original);
+        nodes.add(currNode);
+      }
+    }
+  } while (changed);
+}