[compiler] Fix VariableDeclarator source location (facebook#35348)

Putting up facebook#35129 again
Reverted in facebook#35346 after breaking
main before security patch

This change impacts output formatting in a lot of snaps, so is very
sensitive to additions in main to the fixtures resulting in broken tests
after merging, so we should try merge quickly after rebasing or do a
fast follow to the merge with a snap update.

Author: nathanmarks

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

@@ -4026,6 +4026,7 @@ function lowerAssignment(
           pattern: {
             kind: 'ArrayPattern',
             items,
+            loc: lvalue.node.loc ?? GeneratedSource,
           },
         },
         value,
@@ -4203,6 +4204,7 @@ function lowerAssignment(
           pattern: {
             kind: 'ObjectPattern',
             properties,
+            loc: lvalue.node.loc ?? GeneratedSource,
           },
         },
         value,

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

@@ -694,11 +694,13 @@ export type SpreadPattern = {
 export type ArrayPattern = {
   kind: 'ArrayPattern';
   items: Array<Place | SpreadPattern | Hole>;
+  loc: SourceLocation;
 };
 
 export type ObjectPattern = {
   kind: 'ObjectPattern';
   properties: Array<ObjectProperty | SpreadPattern>;
+  loc: SourceLocation;
 };
 
 export type ObjectPropertyKey =

compiler/packages/babel-plugin-react-compiler/src/Optimization/OutlineJsx.ts

@@ -515,6 +515,7 @@ function emitDestructureProps(
         pattern: {
           kind: 'ObjectPattern',
           properties,
+          loc: GeneratedSource,
         },
         kind: InstructionKind.Let,
       },

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

@@ -702,7 +702,7 @@ function codegenReactiveScope(
     outputComments.push(name.name);
     if (!cx.hasDeclared(identifier)) {
       statements.push(
-        t.variableDeclaration('let', [t.variableDeclarator(name)]),
+        t.variableDeclaration('let', [createVariableDeclarator(name, null)]),
       );
     }
     cacheLoads.push({name, index, value: wrapCacheDep(cx, name)});
@@ -1387,7 +1387,7 @@ function codegenInstructionNullable(
           suggestions: null,
         });
         return createVariableDeclaration(instr.loc, 'const', [
-          t.variableDeclarator(codegenLValue(cx, lvalue), value),
+          createVariableDeclarator(codegenLValue(cx, lvalue), value),
         ]);
       }
       case InstructionKind.Function: {
@@ -1451,7 +1451,7 @@ function codegenInstructionNullable(
           suggestions: null,
         });
         return createVariableDeclaration(instr.loc, 'let', [
-          t.variableDeclarator(codegenLValue(cx, lvalue), value),
+          createVariableDeclarator(codegenLValue(cx, lvalue), value),
         ]);
       }
       case InstructionKind.Reassign: {
@@ -1691,6 +1691,9 @@ function withLoc<T extends (...args: Array<any>) => t.Node>(
   };
 }
 
+const createIdentifier = withLoc(t.identifier);
+const createArrayPattern = withLoc(t.arrayPattern);
+const createObjectPattern = withLoc(t.objectPattern);
 const createBinaryExpression = withLoc(t.binaryExpression);
 const createExpressionStatement = withLoc(t.expressionStatement);
 const _createLabelledStatement = withLoc(t.labeledStatement);
@@ -1722,6 +1725,31 @@ const createTryStatement = withLoc(t.tryStatement);
 const createBreakStatement = withLoc(t.breakStatement);
 const createContinueStatement = withLoc(t.continueStatement);
 
+function createVariableDeclarator(
+  id: t.LVal,
+  init?: t.Expression | null,
+): t.VariableDeclarator {
+  const node = t.variableDeclarator(id, init);
+
+  /*
+   * The variable declarator location is not preserved in HIR, however, we can use the
+   * start location of the id and the end location of the init to recreate the
+   * exact original variable declarator location.
+   *
+   * Or if init is null, we likely have a declaration without an initializer, so we can use the id.loc.end as the end location.
+   */
+  if (id.loc && (init === null || init?.loc)) {
+    node.loc = {
+      start: id.loc.start,
+      end: init?.loc?.end ?? id.loc.end,
+      filename: id.loc.filename,
+      identifierName: undefined,
+    };
+  }
+
+  return node;
+}
+
 function createHookGuard(
   guard: ExternalFunction,
   context: ProgramContext,
@@ -1829,7 +1857,7 @@ function codegenInstruction(
       );
     } else {
       return createVariableDeclaration(instr.loc, 'const', [
-        t.variableDeclarator(
+        createVariableDeclarator(
           convertIdentifier(instr.lvalue.identifier),
           expressionValue,
         ),
@@ -2756,7 +2784,7 @@ function codegenArrayPattern(
 ): t.ArrayPattern {
   const hasHoles = !pattern.items.every(e => e.kind !== 'Hole');
   if (hasHoles) {
-    const result = t.arrayPattern([]);
+    const result = createArrayPattern(pattern.loc, []);
     /*
      * Older versions of babel have a validation bug fixed by
      * https://github.com/babel/babel/pull/10917
@@ -2777,7 +2805,8 @@ function codegenArrayPattern(
     }
     return result;
   } else {
-    return t.arrayPattern(
+    return createArrayPattern(
+      pattern.loc,
       pattern.items.map(item => {
         if (item.kind === 'Hole') {
           return null;
@@ -2797,7 +2826,8 @@ function codegenLValue(
       return codegenArrayPattern(cx, pattern);
     }
     case 'ObjectPattern': {
-      return t.objectPattern(
+      return createObjectPattern(
+        pattern.loc,
         pattern.properties.map(property => {
           if (property.kind === 'ObjectProperty') {
             const key = codegenObjectPropertyKey(cx, property.key);
@@ -2916,7 +2946,7 @@ function convertIdentifier(identifier: Identifier): t.Identifier {
       suggestions: null,
     },
   );
-  return t.identifier(identifier.name.value);
+  return createIdentifier(identifier.loc, identifier.name.value);
 }
 
 function compareScopeDependency(

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidateSourceLocations.ts

@@ -27,7 +27,11 @@ import {Result} from '../Utils/Result';
 
 /**
  * Some common node types that are important for coverage tracking.
- * Based on istanbul-lib-instrument
+ * Based on istanbul-lib-instrument + some other common nodes we expect to be present in the generated AST.
+ *
+ * Note: For VariableDeclaration, VariableDeclarator, and Identifier, we enforce stricter validation
+ * that requires both the source location AND node type to match in the generated AST. This ensures
+ * that variable declarations maintain their structural integrity through compilation.
  */
 const IMPORTANT_INSTRUMENTED_TYPES = new Set([
   'ArrowFunctionExpression',
@@ -54,6 +58,14 @@ const IMPORTANT_INSTRUMENTED_TYPES = new Set([
   'LabeledStatement',
   'ConditionalExpression',
   'LogicalExpression',
+
+  /**
+   * Note: these aren't important for coverage tracking,
+   * but we still want to track them to ensure we aren't regressing them when
+   * we fix the source location tracking for other nodes.
+   */
+  'VariableDeclaration',
+  'Identifier',
 ]);
 
 /**
@@ -114,10 +126,13 @@ export function validateSourceLocations(
 ): Result<void, CompilerError> {
   const errors = new CompilerError();
 
-  // Step 1: Collect important locations from the original source
+  /*
+   * Step 1: Collect important locations from the original source
+   * Note: Multiple node types can share the same location (e.g. VariableDeclarator and Identifier)
+   */
   const importantOriginalLocations = new Map<
     string,
-    {loc: t.SourceLocation; nodeType: string}
+    {loc: t.SourceLocation; nodeTypes: Set<string>}
   >();
 
   func.traverse({
@@ -137,20 +152,31 @@ export function validateSourceLocations(
       // Collect the location if it exists
       if (node.loc) {
         const key = locationKey(node.loc);
-        importantOriginalLocations.set(key, {
-          loc: node.loc,
-          nodeType: node.type,
-        });
+        const existing = importantOriginalLocations.get(key);
+        if (existing) {
+          existing.nodeTypes.add(node.type);
+        } else {
+          importantOriginalLocations.set(key, {
+            loc: node.loc,
+            nodeTypes: new Set([node.type]),
+          });
+        }
       }
     },
   });
 
-  // Step 2: Collect all locations from the generated AST
-  const generatedLocations = new Set<string>();
+  // Step 2: Collect all locations from the generated AST with their node types
+  const generatedLocations = new Map<string, Set<string>>();
 
   function collectGeneratedLocations(node: t.Node): void {
     if (node.loc) {
-      generatedLocations.add(locationKey(node.loc));
+      const key = locationKey(node.loc);
+      const nodeTypes = generatedLocations.get(key);
+      if (nodeTypes) {
+        nodeTypes.add(node.type);
+      } else {
+        generatedLocations.set(key, new Set([node.type]));
+      }
     }
 
     // Use Babel's VISITOR_KEYS to traverse only actual node properties
@@ -183,22 +209,86 @@ export function validateSourceLocations(
     collectGeneratedLocations(outlined.fn.body);
   }
 
-  // Step 3: Validate that all important locations are preserved
-  for (const [key, {loc, nodeType}] of importantOriginalLocations) {
-    if (!generatedLocations.has(key)) {
-      errors.pushDiagnostic(
-        CompilerDiagnostic.create({
-          category: ErrorCategory.Todo,
-          reason: 'Important source location missing in generated code',
-          description:
-            `Source location for ${nodeType} is missing in the generated output. This can cause coverage instrumentation ` +
-            `to fail to track this code properly, resulting in inaccurate coverage reports.`,
-        }).withDetails({
-          kind: 'error',
-          loc,
-          message: null,
-        }),
-      );
+  /*
+   * Step 3: Validate that all important locations are preserved
+   * For certain node types, also validate that the node type matches
+   */
+  const strictNodeTypes = new Set([
+    'VariableDeclaration',
+    'VariableDeclarator',
+    'Identifier',
+  ]);
+
+  const reportMissingLocation = (
+    loc: t.SourceLocation,
+    nodeType: string,
+  ): void => {
+    errors.pushDiagnostic(
+      CompilerDiagnostic.create({
+        category: ErrorCategory.Todo,
+        reason: 'Important source location missing in generated code',
+        description:
+          `Source location for ${nodeType} is missing in the generated output. This can cause coverage instrumentation ` +
+          `to fail to track this code properly, resulting in inaccurate coverage reports.`,
+      }).withDetails({
+        kind: 'error',
+        loc,
+        message: null,
+      }),
+    );
+  };
+
+  const reportWrongNodeType = (
+    loc: t.SourceLocation,
+    expectedType: string,
+    actualTypes: Set<string>,
+  ): void => {
+    errors.pushDiagnostic(
+      CompilerDiagnostic.create({
+        category: ErrorCategory.Todo,
+        reason:
+          'Important source location has wrong node type in generated code',
+        description:
+          `Source location for ${expectedType} exists in the generated output but with wrong node type(s): ${Array.from(actualTypes).join(', ')}. ` +
+          `This can cause coverage instrumentation to fail to track this code properly, resulting in inaccurate coverage reports.`,
+      }).withDetails({
+        kind: 'error',
+        loc,
+        message: null,
+      }),
+    );
+  };
+
+  for (const [key, {loc, nodeTypes}] of importantOriginalLocations) {
+    const generatedNodeTypes = generatedLocations.get(key);
+
+    if (!generatedNodeTypes) {
+      // Location is completely missing
+      reportMissingLocation(loc, Array.from(nodeTypes).join(', '));
+    } else {
+      // Location exists, check each node type
+      for (const nodeType of nodeTypes) {
+        if (
+          strictNodeTypes.has(nodeType) &&
+          !generatedNodeTypes.has(nodeType)
+        ) {
+          /*
+           * For strict node types, the specific node type must be present
+           * Check if any generated node type is also an important original node type
+           */
+          const hasValidNodeType = Array.from(generatedNodeTypes).some(
+            genType => nodeTypes.has(genType),
+          );
+
+          if (hasValidNodeType) {
+            // At least one generated node type is valid (also in original), so this is just missing
+            reportMissingLocation(loc, nodeType);
+          } else {
+            // None of the generated node types are in original - this is wrong node type
+            reportWrongNodeType(loc, nodeType, generatedNodeTypes);
+          }
+        }
+      }
     }
   }
 

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-capture-in-method-receiver-and-mutate.expect.md

@@ -35,10 +35,8 @@ function Component() {
   let t0;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
     const a = makeObject_Primitives();
-
     const x = [];
     x.push(a);
-
     mutate(x);
     t0 = [x, a];
     $[0] = t0;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/alias-capture-in-method-receiver.expect.md

@@ -33,7 +33,6 @@ function Component() {
   if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
     const x = [];
     x.push(a);
-
     t1 = [x, a];
     $[1] = t1;
   } else {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/aliased-nested-scope-fn-expr.expect.md

@@ -85,14 +85,10 @@ function Component(t0) {
   let t1;
   if ($[0] !== prop) {
     const obj = shallowCopy(prop);
-
     const aliasedObj = identity(obj);
-
     const getId = () => obj.id;
-
     mutate(aliasedObj);
     setPropertyByKey(aliasedObj, "id", prop.id + 1);
-
     t1 = <Stringify getId={getId} shouldInvokeFns={true} />;
     $[0] = prop;
     $[1] = t1;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/aliased-nested-scope-truncated-dep.expect.md

@@ -181,12 +181,9 @@ function Component(t0) {
   if ($[0] !== prop) {
     const obj = shallowCopy(prop);
     const aliasedObj = identity(obj);
-
     const id = [obj.id];
-
     mutate(aliasedObj);
     setPropertyByKey(aliasedObj, "id", prop.id + 1);
-
     t1 = <Stringify id={id} />;
     $[0] = prop;
     $[1] = t1;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/align-scopes-within-nested-valueblock-in-array.expect.md

@@ -54,7 +54,6 @@ function Foo(t0) {
   let t1;
   if ($[0] !== cond1 || $[1] !== cond2) {
     const arr = makeArray({ a: 2 }, 2, []);
-
     t1 = cond1 ? (
       <>
         <div>{identity("foo")}</div>