fix(compiler-cli): handle transformed classes when generating HMR code (angular#60298)

We had several places where we were trying to get the source file of a class for which we're generating HMR-related code. These calls will fail if the class was transformed so we have to get its source file through the original node.

Fixes angular#60287.

PR Close angular#60298

Author: crisbeto

packages/compiler-cli/src/ngtsc/annotations/component/src/handler.ts

@@ -1896,7 +1896,7 @@ export class ComponentDecoratorHandler
     const res = compileResults(fac, def, classMetadata, 'ɵcmp', null, null, debugInfo, null);
     return hmrMeta === null || res.length === 0
       ? null
-      : getHmrUpdateDeclaration(res, pool.statements, hmrMeta, node.getSourceFile());
+      : getHmrUpdateDeclaration(res, pool.statements, hmrMeta, node);
   }
 
   /**

packages/compiler-cli/src/ngtsc/hmr/src/extract_dependencies.ts

@@ -42,7 +42,7 @@ export function extractHmrDependencies(
 } | null {
   const name = ts.isClassDeclaration(node) && node.name ? node.name.text : null;
   const visitor = new PotentialTopLevelReadsVisitor();
-  const sourceFile = node.getSourceFile();
+  const sourceFile = ts.getOriginalNode(node).getSourceFile();
 
   // Visit all of the compiled expressions to look for potential
   // local references that would have to be retained.

packages/compiler-cli/src/ngtsc/hmr/src/metadata.ts

@@ -47,7 +47,7 @@ export function extractHmrMetatadata(
     return null;
   }
 
-  const sourceFile = clazz.getSourceFile();
+  const sourceFile = ts.getOriginalNode(clazz).getSourceFile();
   const filePath =
     getProjectRelativePath(sourceFile.fileName, rootDirs, compilerHost) ||
     compilerHost.getCanonicalFileName(sourceFile.fileName);

packages/compiler-cli/src/ngtsc/hmr/src/update_declaration.ts

@@ -13,19 +13,20 @@ import {
   presetImportManagerForceNamespaceImports,
   translateStatement,
 } from '../../translator';
+import {ClassDeclaration} from '../../reflection';
 import ts from 'typescript';
 
 /**
  * Gets the declaration for the function that replaces the metadata of a class during HMR.
  * @param compilationResults Code generated for the class during compilation.
  * @param meta HMR metadata about the class.
- * @param sourceFile File in which the class is defined.
+ * @param declaration Class for which the update declaration is being generated.
  */
 export function getHmrUpdateDeclaration(
   compilationResults: CompileResult[],
   constantStatements: o.Statement[],
   meta: R3HmrMetadata,
-  sourceFile: ts.SourceFile,
+  declaration: ClassDeclaration,
 ): ts.FunctionDeclaration {
   const namespaceSpecifiers = meta.namespaceDependencies.reduce((result, current) => {
     result.set(current.moduleName, current.assignedName);
@@ -37,6 +38,7 @@ export function getHmrUpdateDeclaration(
     rewriter: importRewriter,
   });
   const callback = compileHmrUpdateCallback(compilationResults, constantStatements, meta);
+  const sourceFile = ts.getOriginalNode(declaration).getSourceFile();
   const node = translateStatement(sourceFile, callback, importManager) as ts.FunctionDeclaration;
 
   // The output AST doesn't support modifiers so we have to emit to

packages/compiler-cli/test/ngtsc/hmr_spec.ts

@@ -9,6 +9,10 @@
 import {runInEachFileSystem} from '../../src/ngtsc/file_system/testing';
 import {loadStandardTestFiles} from '../../src/ngtsc/testing';
 import {NgtscTestEnvironment} from './env';
+import {CompilerOptions} from '@angular/compiler-cli/src/transformers/api';
+import {createCompilerHost} from '@angular/compiler-cli/src/transformers/compiler_host';
+import {NgtscProgram} from '@angular/compiler-cli/src/ngtsc/program';
+import ts from 'typescript';
 
 runInEachFileSystem(() => {
   describe('HMR code generation', () => {
@@ -901,5 +905,71 @@ runInEachFileSystem(() => {
         'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, token, Foo, Component) {',
       );
     });
+
+    it('should generate HMR code for a transformed class', () => {
+      env.write(
+        'test.ts',
+        `
+          import {Component} from '@angular/core';
+
+          @Component({template: ''})
+          export class Cmp {}
+        `,
+      );
+
+      const options: CompilerOptions = {
+        target: ts.ScriptTarget.Latest,
+        module: ts.ModuleKind.ESNext,
+        _enableHmr: true,
+      };
+
+      const program = new NgtscProgram(['/test.ts'], options, createCompilerHost({options}));
+      const transformers = program.compiler.prepareEmit().transformers;
+
+      transformers.before!.unshift((ctx) => (sourceFile) => {
+        const visitor = (node: ts.Node) => {
+          if (ts.isClassDeclaration(node) && node.name?.getText() === 'Cmp') {
+            const newMember = ctx.factory.createPropertyDeclaration(
+              undefined,
+              'newProp',
+              undefined,
+              undefined,
+              ctx.factory.createNumericLiteral(123),
+            );
+
+            return ctx.factory.updateClassDeclaration(
+              node,
+              node.modifiers,
+              node.name,
+              node.typeParameters,
+              node.heritageClauses,
+              [newMember, ...node.members],
+            );
+          }
+          return ts.visitEachChild(node, visitor, ctx);
+        };
+        return ts.visitEachChild(sourceFile, visitor, ctx);
+      });
+
+      const {diagnostics, emitSkipped} = program
+        .getTsProgram()
+        .emit(undefined, undefined, undefined, undefined, transformers);
+      const declaration = program
+        .getTsProgram()
+        .getSourceFile('/test.ts')
+        ?.statements.find(
+          (stmt) => ts.isClassDeclaration(stmt) && stmt.name?.getText() === 'Cmp',
+        ) as ts.ClassDeclaration;
+
+      const jsContents = env.getContents('/test.js');
+      const hmrContents = program.compiler.emitHmrUpdateModule(declaration);
+
+      expect(diagnostics.length).toBe(0);
+      expect(emitSkipped).toBe(false);
+
+      expect(jsContents).toContain('ɵreplaceMetadata(Cmp');
+      expect(jsContents).toContain('newProp = 123');
+      expect(hmrContents).toContain('export default function Cmp_UpdateMetadata');
+    });
   });
 });