[compiler] Track locations for dependencies (facebook#35794)

Tracks locations for reactive scope dependencies, both on the deps and
portions of the path. The immediate need for this is a non-public
experiment where we're exploring type-directed compilation, and
sometimes look up the types of expressions by location. We need to
preserve locations accurately for that to work, including the locations
of the deps.

## Test Plan

Locations for dependencies are not easy to test: i manually spot-checked
the new fixture to ensure that the deps look right. This is fine as
best-effort since it doesn't impact any of our core compilation logic, i
may fix forward if there are issues and will think about how to test.

Author: josephsavona

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

@@ -31,6 +31,7 @@ import {
   PropertyLiteral,
   ReactiveScopeDependency,
   ScopeId,
+  SourceLocation,
   TInstruction,
 } from './HIR';
 
@@ -244,6 +245,7 @@ class PropertyPathRegistry {
   getOrCreateIdentifier(
     identifier: Identifier,
     reactive: boolean,
+    loc: SourceLocation,
   ): PropertyPathNode {
     /**
      * Reads from a statically scoped variable are always safe in JS,
@@ -260,6 +262,7 @@ class PropertyPathRegistry {
           identifier,
           reactive,
           path: [],
+          loc,
         },
         hasOptional: false,
         parent: null,
@@ -290,6 +293,7 @@ class PropertyPathRegistry {
           identifier: parent.fullPath.identifier,
           reactive: parent.fullPath.reactive,
           path: parent.fullPath.path.concat(entry),
+          loc: entry.loc,
         },
         hasOptional: parent.hasOptional || entry.optional,
       };
@@ -304,7 +308,7 @@ class PropertyPathRegistry {
      * so all subpaths of a PropertyLoad should already exist
      * (e.g. a.b is added before a.b.c),
      */
-    let currNode = this.getOrCreateIdentifier(n.identifier, n.reactive);
+    let currNode = this.getOrCreateIdentifier(n.identifier, n.reactive, n.loc);
     if (n.path.length === 0) {
       return currNode;
     }
@@ -323,20 +327,21 @@ class PropertyPathRegistry {
 }
 
 function getMaybeNonNullInInstruction(
-  instr: InstructionValue,
+  value: InstructionValue,
   context: CollectHoistablePropertyLoadsContext,
 ): PropertyPathNode | null {
   let path: ReactiveScopeDependency | null = null;
-  if (instr.kind === 'PropertyLoad') {
-    path = context.temporaries.get(instr.object.identifier.id) ?? {
-      identifier: instr.object.identifier,
-      reactive: instr.object.reactive,
+  if (value.kind === 'PropertyLoad') {
+    path = context.temporaries.get(value.object.identifier.id) ?? {
+      identifier: value.object.identifier,
+      reactive: value.object.reactive,
       path: [],
+      loc: value.loc,
     };
-  } else if (instr.kind === 'Destructure') {
-    path = context.temporaries.get(instr.value.identifier.id) ?? null;
-  } else if (instr.kind === 'ComputedLoad') {
-    path = context.temporaries.get(instr.object.identifier.id) ?? null;
+  } else if (value.kind === 'Destructure') {
+    path = context.temporaries.get(value.value.identifier.id) ?? null;
+  } else if (value.kind === 'ComputedLoad') {
+    path = context.temporaries.get(value.object.identifier.id) ?? null;
   }
   return path != null ? context.registry.getOrCreateProperty(path) : null;
 }
@@ -393,7 +398,11 @@ function collectNonNullsInBlocks(
   ) {
     const identifier = fn.params[0].identifier;
     knownNonNullIdentifiers.add(
-      context.registry.getOrCreateIdentifier(identifier, true),
+      context.registry.getOrCreateIdentifier(
+        identifier,
+        true,
+        fn.params[0].loc,
+      ),
     );
   }
   const nodes = new Map<
@@ -468,6 +477,7 @@ function collectNonNullsInBlocks(
                 identifier: dep.root.value.identifier,
                 path: dep.path.slice(0, i),
                 reactive: dep.root.value.reactive,
+                loc: dep.loc,
               });
               assumedNonNullObjects.add(depNode);
             }
@@ -654,17 +664,23 @@ function reduceMaybeOptionalChains(
     changed = false;
 
     for (const original of optionalChainNodes) {
-      let {identifier, path: origPath, reactive} = original.fullPath;
+      let {
+        identifier,
+        path: origPath,
+        reactive,
+        loc: origLoc,
+      } = original.fullPath;
       let currNode: PropertyPathNode = registry.getOrCreateIdentifier(
         identifier,
         reactive,
+        origLoc,
       );
       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}
+            ? {property: entry.property, optional: false, loc: entry.loc}
             : entry;
         currNode = PropertyPathRegistry.getOrCreatePropertyEntry(
           currNode,

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

@@ -5,7 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {CompilerError} from '..';
+import {CompilerError, SourceLocation} from '..';
 import {assertNonNull} from './CollectHoistablePropertyLoads';
 import {
   BlockId,
@@ -169,6 +169,7 @@ function matchOptionalTestBlock(
   propertyId: IdentifierId;
   storeLocalInstr: Instruction;
   consequentGoto: BlockId;
+  propertyLoadLoc: SourceLocation;
 } | null {
   const consequentBlock = assertNonNull(blocks.get(terminal.consequent));
   if (
@@ -221,6 +222,7 @@ function matchOptionalTestBlock(
       propertyId: propertyLoad.lvalue.identifier.id,
       storeLocalInstr,
       consequentGoto: consequentBlock.terminal.block,
+      propertyLoadLoc: propertyLoad.loc,
     };
   }
   return null;
@@ -275,7 +277,11 @@ function traverseOptionalBlock(
         instrVal.kind === 'PropertyLoad' &&
         instrVal.object.identifier.id === prevInstr.lvalue.identifier.id
       ) {
-        path.push({property: instrVal.property, optional: false});
+        path.push({
+          property: instrVal.property,
+          optional: false,
+          loc: instrVal.loc,
+        });
       } else {
         return null;
       }
@@ -292,6 +298,7 @@ function traverseOptionalBlock(
       identifier: maybeTest.instructions[0].value.place.identifier,
       reactive: maybeTest.instructions[0].value.place.reactive,
       path,
+      loc: maybeTest.instructions[0].value.place.loc,
     };
     test = maybeTest.terminal;
   } else if (maybeTest.terminal.kind === 'optional') {
@@ -390,16 +397,18 @@ function traverseOptionalBlock(
       loc: optional.terminal.loc,
     },
   );
-  const load = {
+  const load: ReactiveScopeDependency = {
     identifier: baseObject.identifier,
     reactive: baseObject.reactive,
     path: [
       ...baseObject.path,
       {
         property: matchConsequentResult.property,
         optional: optional.terminal.optional,
+        loc: matchConsequentResult.propertyLoadLoc,
       },
     ],
+    loc: matchConsequentResult.propertyLoadLoc,
   };
   context.processedInstrsInOptional.add(matchConsequentResult.storeLocalInstr);
   context.processedInstrsInOptional.add(test);

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

@@ -12,6 +12,7 @@ import {
   Identifier,
   PropertyLiteral,
   ReactiveScopeDependency,
+  SourceLocation,
 } from '../HIR';
 import {printIdentifier} from '../HIR/PrintHIR';
 
@@ -36,12 +37,13 @@ export class ReactiveScopeDependencyTreeHIR {
    * duplicates when traversing the CFG.
    */
   constructor(hoistableObjects: Iterable<ReactiveScopeDependency>) {
-    for (const {path, identifier, reactive} of hoistableObjects) {
+    for (const {path, identifier, reactive, loc} of hoistableObjects) {
       let currNode = ReactiveScopeDependencyTreeHIR.#getOrCreateRoot(
         identifier,
         reactive,
         this.#hoistableObjects,
         path.length > 0 && path[0].optional ? 'Optional' : 'NonNull',
+        loc,
       );
 
       for (let i = 0; i < path.length; i++) {
@@ -62,6 +64,7 @@ export class ReactiveScopeDependencyTreeHIR {
           nextNode = {
             properties: new Map(),
             accessType,
+            loc: path[i].loc,
           };
           currNode.properties.set(path[i].property, nextNode);
         }
@@ -75,6 +78,7 @@ export class ReactiveScopeDependencyTreeHIR {
     reactive: boolean,
     roots: Map<Identifier, TreeNode<T> & {reactive: boolean}>,
     defaultAccessType: T,
+    loc: SourceLocation,
   ): TreeNode<T> {
     // roots can always be accessed unconditionally in JS
     let rootNode = roots.get(identifier);
@@ -84,6 +88,7 @@ export class ReactiveScopeDependencyTreeHIR {
         properties: new Map(),
         reactive,
         accessType: defaultAccessType,
+        loc,
       };
       roots.set(identifier, rootNode);
     } else {
@@ -102,12 +107,13 @@ export class ReactiveScopeDependencyTreeHIR {
    * safe-to-evaluate subpath
    */
   addDependency(dep: ReactiveScopeDependency): void {
-    const {identifier, reactive, path} = dep;
+    const {identifier, reactive, path, loc} = dep;
     let depCursor = ReactiveScopeDependencyTreeHIR.#getOrCreateRoot(
       identifier,
       reactive,
       this.#deps,
       PropertyAccessType.UnconditionalAccess,
+      loc,
     );
     /**
      * hoistableCursor is null if depCursor is not an object we can hoist
@@ -153,6 +159,7 @@ export class ReactiveScopeDependencyTreeHIR {
           depCursor,
           entry.property,
           accessType,
+          entry.loc,
         );
       } else if (
         hoistableCursor != null &&
@@ -163,6 +170,7 @@ export class ReactiveScopeDependencyTreeHIR {
           depCursor,
           entry.property,
           PropertyAccessType.UnconditionalAccess,
+          entry.loc,
         );
       } else {
         /**
@@ -306,6 +314,7 @@ function merge(
 type TreeNode<T extends string> = {
   properties: Map<PropertyLiteral, TreeNode<T>>;
   accessType: T;
+  loc: SourceLocation;
 };
 type HoistableNode = TreeNode<'Optional' | 'NonNull'>;
 type DependencyNode = TreeNode<PropertyAccessType>;
@@ -323,7 +332,7 @@ function collectMinimalDependenciesInSubtree(
   results: Set<ReactiveScopeDependency>,
 ): void {
   if (isDependency(node.accessType)) {
-    results.add({identifier: rootIdentifier, reactive, path});
+    results.add({identifier: rootIdentifier, reactive, path, loc: node.loc});
   } else {
     for (const [childName, childNode] of node.properties) {
       collectMinimalDependenciesInSubtree(
@@ -335,6 +344,7 @@ function collectMinimalDependenciesInSubtree(
           {
             property: childName,
             optional: isOptional(childNode.accessType),
+            loc: childNode.loc,
           },
         ],
         results,
@@ -362,12 +372,14 @@ function makeOrMergeProperty(
   node: DependencyNode,
   property: PropertyLiteral,
   accessType: PropertyAccessType,
+  loc: SourceLocation,
 ): DependencyNode {
   let child = node.properties.get(property);
   if (child == null) {
     child = {
       properties: new Map(),
       accessType,
+      loc,
     };
     node.properties.set(property, child);
   } else {

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

@@ -1639,6 +1639,7 @@ export function makePropertyLiteral(value: string | number): PropertyLiteral {
 export type DependencyPathEntry = {
   property: PropertyLiteral;
   optional: boolean;
+  loc: SourceLocation;
 };
 export type DependencyPath = Array<DependencyPathEntry>;
 export type ReactiveScopeDependency = {
@@ -1656,6 +1657,7 @@ export type ReactiveScopeDependency = {
    */
   reactive: boolean;
   path: DependencyPath;
+  loc: SourceLocation;
 };
 
 export function areEqualPaths(a: DependencyPath, b: DependencyPath): boolean {