fix(compiler-cli): handle more node types when extracting dependencies (#59445)

crisbeto · kirjs · commit 3638f93acc0c · 2025-01-09T12:22:46.000-05:00
Fixes that the HMR dependency extraction logic wasn't handling some node types. Most of these are a bit edge-case-ey in component definitions, but variable initializers and arrow functions can realistically happen in factories. PR Close #59445
diff --git a/packages/compiler-cli/src/ngtsc/hmr/src/extract_dependencies.ts b/packages/compiler-cli/src/ngtsc/hmr/src/extract_dependencies.ts
@@ -210,10 +210,7 @@ class PotentialTopLevelReadsVisitor extends o.RecursiveAstVisitor {
     }
 
     // Identifier referenced at the top level. Unlikely.
-    if (
-      ts.isSourceFile(parent) ||
-      (ts.isExpressionStatement(parent) && parent.expression === node)
-    ) {
+    if (ts.isSourceFile(parent)) {
       return true;
     }
 
@@ -225,6 +222,7 @@ class PotentialTopLevelReadsVisitor extends o.RecursiveAstVisitor {
 
     // Identifier used in a nested expression is only top-level if it's the actual expression.
     if (
+      ts.isExpressionStatement(parent) ||
       ts.isPropertyAccessExpression(parent) ||
       ts.isComputedPropertyName(parent) ||
       ts.isTemplateSpan(parent) ||
@@ -235,8 +233,6 @@ class PotentialTopLevelReadsVisitor extends o.RecursiveAstVisitor {
       ts.isIfStatement(parent) ||
       ts.isDoStatement(parent) ||
       ts.isWhileStatement(parent) ||
-      ts.isForInStatement(parent) ||
-      ts.isForOfStatement(parent) ||
       ts.isSwitchStatement(parent) ||
       ts.isCaseClause(parent) ||
       ts.isThrowStatement(parent)
@@ -249,17 +245,28 @@ class PotentialTopLevelReadsVisitor extends o.RecursiveAstVisitor {
       return parent.elements.includes(node);
     }
 
-    // Identifier in a property assignment is only top level if it's the initializer.
-    if (ts.isPropertyAssignment(parent)) {
+    // If the parent is an initialized node, the identifier is
+    // at the top level if it's the initializer itself.
+    if (
+      ts.isPropertyAssignment(parent) ||
+      ts.isParameter(parent) ||
+      ts.isBindingElement(parent) ||
+      ts.isPropertyDeclaration(parent) ||
+      ts.isEnumMember(parent)
+    ) {
       return parent.initializer === node;
     }
 
+    // Identifier in a function is top level if it's either the name or the initializer.
+    if (ts.isVariableDeclaration(parent)) {
+      return parent.name === node || parent.initializer === node;
+    }
+
     // Identifier in a declaration is only top level if it's the name.
     // In shorthand assignments the name is also the value.
     if (
       ts.isClassDeclaration(parent) ||
       ts.isFunctionDeclaration(parent) ||
-      ts.isVariableDeclaration(parent) ||
       ts.isShorthandPropertyAssignment(parent)
     ) {
       return parent.name === node;
@@ -273,6 +280,20 @@ class PotentialTopLevelReadsVisitor extends o.RecursiveAstVisitor {
       return parent.left === node || parent.right === node;
     }
 
+    if (ts.isForInStatement(parent) || ts.isForOfStatement(parent)) {
+      return parent.expression === node || parent.initializer === node;
+    }
+
+    if (ts.isForStatement(parent)) {
+      return (
+        parent.condition === node || parent.initializer === node || parent.incrementor === node
+      );
+    }
+
+    if (ts.isArrowFunction(parent)) {
+      return parent.body === node;
+    }
+
     // It's unlikely that we'll run into imports/exports in this use case.
     // We handle them since it's simple and for completeness' sake.
     if (ts.isImportSpecifier(parent) || ts.isExportSpecifier(parent)) {
diff --git a/packages/compiler-cli/test/ngtsc/hmr_spec.ts b/packages/compiler-cli/test/ngtsc/hmr_spec.ts
@@ -431,5 +431,76 @@ runInEachFileSystem(() => {
         'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, providers, Component) {',
       );
     });
+
+    it('should capture variable initializer dependencies', () => {
+      enableHmr();
+      env.write(
+        'test.ts',
+        `
+          import {Component, InjectionToken} from '@angular/core';
+
+          const token = new InjectionToken<number>('TEST');
+          const value = 123;
+
+          @Component({
+            template: '',
+            providers: [{
+              provide: token,
+              useFactory: () => {
+                const v = value;
+                return v;
+              }
+            }]
+          })
+          export class Cmp {}
+        `,
+      );
+
+      env.driveMain();
+
+      const jsContents = env.getContents('test.js');
+      const hmrContents = env.driveHmr('test.ts', 'Cmp');
+
+      expect(jsContents).toContain(
+        'ɵɵreplaceMetadata(Cmp, m.default, [i0], [token, value, Component]));',
+      );
+      expect(hmrContents).toContain(
+        'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, token, value, Component) {',
+      );
+    });
+
+    it('should capture arrow function dependencies', () => {
+      enableHmr();
+      env.write(
+        'test.ts',
+        `
+          import {Component, InjectionToken} from '@angular/core';
+
+          const token = new InjectionToken<number>('TEST');
+          const value = 123;
+
+          @Component({
+            template: '',
+            providers: [{
+              provide: token,
+              useFactory: () => value
+            }]
+          })
+          export class Cmp {}
+        `,
+      );
+
+      env.driveMain();
+
+      const jsContents = env.getContents('test.js');
+      const hmrContents = env.driveHmr('test.ts', 'Cmp');
+
+      expect(jsContents).toContain(
+        'ɵɵreplaceMetadata(Cmp, m.default, [i0], [token, value, Component]));',
+      );
+      expect(hmrContents).toContain(
+        'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, token, value, Component) {',
+      );
+    });
   });
 });
