fix(compiler-cli): account for multiple generated namespace imports in HMR (#58924)

crisbeto · pkozlowski-opensource · commit f280467398c6 · 2024-11-28T10:00:56.000+01:00
The current HMR compiler assumes that there will only be one namespace import in the generated code (`@angular/core`). This is incorrect, because the compiler may need to generate additional imports in some cases (e.g. importing directives through a module). These changes adjust the compiler to capture all the namespaces in an array and pass them along. Fixes #58915. PR Close #58924
diff --git a/packages/compiler-cli/src/ngtsc/hmr/src/extract_dependencies.ts b/packages/compiler-cli/src/ngtsc/hmr/src/extract_dependencies.ts
@@ -6,27 +6,26 @@
  * found in the LICENSE file at https://angular.dev/license
  */
 
-import {R3CompiledExpression, outputAst as o} from '@angular/compiler';
+import {R3CompiledExpression, R3HmrNamespaceDependency, outputAst as o} from '@angular/compiler';
 import {DeclarationNode} from '../../reflection';
 import {CompileResult} from '../../transform';
 import ts from 'typescript';
 
 /**
- * Determines the names of the file-level locals that the HMR
- * initializer needs to capture and pass along.
+ * Determines the file-level dependencies that the HMR initializer needs to capture and pass along.
  * @param sourceFile File in which the file is being compiled.
  * @param definition Compiled component definition.
  * @param factory Compiled component factory.
  * @param classMetadata Compiled `setClassMetadata` expression, if any.
  * @param debugInfo Compiled `setClassDebugInfo` expression, if any.
  */
-export function extractHmrLocals(
+export function extractHmrDependencies(
   node: DeclarationNode,
   definition: R3CompiledExpression,
   factory: CompileResult,
   classMetadata: o.Statement | null,
   debugInfo: o.Statement | null,
-): string[] {
+): {local: string[]; external: R3HmrNamespaceDependency[]} {
   const name = ts.isClassDeclaration(node) && node.name ? node.name.text : null;
   const visitor = new PotentialTopLevelReadsVisitor();
   const sourceFile = node.getSourceFile();
@@ -44,7 +43,14 @@ export function extractHmrLocals(
   // variables inside of functions. Note that we filter out the class name since it is always
   // defined and it saves us having to repeat this logic wherever the locals are consumed.
   const availableTopLevel = getTopLevelDeclarationNames(sourceFile);
-  return Array.from(visitor.allReads).filter((r) => r !== name && availableTopLevel.has(r));
+
+  return {
+    local: Array.from(visitor.allReads).filter((r) => r !== name && availableTopLevel.has(r)),
+    external: Array.from(visitor.namespaceReads, (name, index) => ({
+      moduleName: name,
+      assignedName: `ɵhmr${index}`,
+    })),
+  };
 }
 
 /**
@@ -138,6 +144,14 @@ function trackBindingName(node: ts.BindingName, results: Set<string>): void {
  */
 class PotentialTopLevelReadsVisitor extends o.RecursiveAstVisitor {
   readonly allReads = new Set<string>();
+  readonly namespaceReads = new Set<string>();
+
+  override visitExternalExpr(ast: o.ExternalExpr, context: any) {
+    if (ast.value.moduleName !== null) {
+      this.namespaceReads.add(ast.value.moduleName);
+    }
+    super.visitExternalExpr(ast, context);
+  }
 
   override visitReadVarExpr(ast: o.ReadVarExpr, context: any) {
     this.allReads.add(ast.name);
diff --git a/packages/compiler-cli/src/ngtsc/hmr/src/metadata.ts b/packages/compiler-cli/src/ngtsc/hmr/src/metadata.ts
@@ -10,7 +10,7 @@ import {R3CompiledExpression, R3HmrMetadata, outputAst as o} from '@angular/comp
 import {DeclarationNode, ReflectionHost} from '../../reflection';
 import {getProjectRelativePath} from '../../util/src/path';
 import {CompileResult} from '../../transform';
-import {extractHmrLocals} from './extract_locals';
+import {extractHmrDependencies} from './extract_dependencies';
 import ts from 'typescript';
 
 /**
@@ -43,12 +43,13 @@ export function extractHmrMetatadata(
     getProjectRelativePath(sourceFile, rootDirs, compilerHost) ||
     compilerHost.getCanonicalFileName(sourceFile.fileName);
 
+  const dependencies = extractHmrDependencies(clazz, definition, factory, classMetadata, debugInfo);
   const meta: R3HmrMetadata = {
     type: new o.WrappedNodeExpr(clazz.name),
     className: clazz.name.text,
     filePath,
-    locals: extractHmrLocals(clazz, definition, factory, classMetadata, debugInfo),
-    coreName: '__ngCore__',
+    localDependencies: dependencies.local,
+    namespaceDependencies: dependencies.external,
   };
 
   return meta;
diff --git a/packages/compiler-cli/src/ngtsc/hmr/src/update_declaration.ts b/packages/compiler-cli/src/ngtsc/hmr/src/update_declaration.ts
@@ -27,7 +27,11 @@ export function getHmrUpdateDeclaration(
   meta: R3HmrMetadata,
   sourceFile: ts.SourceFile,
 ): ts.FunctionDeclaration {
-  const importRewriter = new HmrModuleImportRewriter(meta.coreName);
+  const namespaceSpecifiers = meta.namespaceDependencies.reduce((result, current) => {
+    result.set(current.moduleName, current.assignedName);
+    return result;
+  }, new Map<string, string>());
+  const importRewriter = new HmrModuleImportRewriter(namespaceSpecifiers);
   const importManager = new ImportManager({
     ...presetImportManagerForceNamespaceImports,
     rewriter: importRewriter,
@@ -52,12 +56,11 @@ export function getHmrUpdateDeclaration(
   );
 }
 
-/** Rewriter that replaces namespace imports to `@angular/core` with a specifier identifier. */
 class HmrModuleImportRewriter {
-  constructor(private readonly coreName: string) {}
+  constructor(private readonly lookup: Map<string, string>) {}
 
   rewriteNamespaceImportIdentifier(specifier: string, moduleName: string): string {
-    return moduleName === '@angular/core' ? this.coreName : specifier;
+    return this.lookup.has(moduleName) ? this.lookup.get(moduleName)! : specifier;
   }
 
   rewriteSymbol(symbol: string): string {
diff --git a/packages/compiler-cli/test/ngtsc/hmr_spec.ts b/packages/compiler-cli/test/ngtsc/hmr_spec.ts
@@ -102,12 +102,13 @@ runInEachFileSystem(() => {
       const jsContents = env.getContents('test.js');
       const hmrContents = env.driveHmr('test.ts', 'Cmp');
 
+      expect(jsContents).toContain(`import * as i0 from "@angular/core";`);
       expect(jsContents).toContain('function Cmp_HmrLoad(t) {');
       expect(jsContents).toContain(
         'import(/* @vite-ignore */\n"/@ng/component?c=test.ts%40Cmp&t=" + encodeURIComponent(t))',
       );
       expect(jsContents).toContain(
-        ').then(m => m.default && i0.ɵɵreplaceMetadata(Cmp, m.default, i0, ' +
+        ').then(m => m.default && i0.ɵɵreplaceMetadata(Cmp, m.default, [i0], ' +
           '[Dep, transformValue, TOKEN, Component, Inject, ViewChild, Input]));',
       );
       expect(jsContents).toContain('Cmp_HmrLoad(Date.now());');
@@ -117,12 +118,82 @@ runInEachFileSystem(() => {
       );
 
       expect(hmrContents).toContain(
-        'export default function Cmp_UpdateMetadata(Cmp, __ngCore__, Dep, transformValue, TOKEN, Component, Inject, ViewChild, Input) {',
+        'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, Dep, transformValue, TOKEN, Component, Inject, ViewChild, Input) {',
       );
+      expect(hmrContents).toContain(`const ɵhmr0 = ɵɵnamespaces[0];`);
       expect(hmrContents).toContain('Cmp.ɵfac = function Cmp_Factory');
-      expect(hmrContents).toContain('Cmp.ɵcmp = /*@__PURE__*/ __ngCore__.ɵɵdefineComponent');
-      expect(hmrContents).toContain('__ngCore__.ɵsetClassMetadata(Cmp,');
-      expect(hmrContents).toContain('__ngCore__.ɵsetClassDebugInfo(Cmp,');
+      expect(hmrContents).toContain('Cmp.ɵcmp = /*@__PURE__*/ ɵhmr0.ɵɵdefineComponent');
+      expect(hmrContents).toContain('ɵhmr0.ɵsetClassMetadata(Cmp,');
+      expect(hmrContents).toContain('ɵhmr0.ɵsetClassDebugInfo(Cmp,');
+    });
+
+    it('should generate an HMR initializer and update function for a class that depends on multiple namespaces', () => {
+      enableHmr();
+      env.write(
+        'dep.ts',
+        `
+          import {Directive, NgModule} from '@angular/core';
+
+          @Directive({
+            selector: '[dep]',
+            standalone: true,
+          })
+          export class Dep {}
+
+          @NgModule({
+            imports: [Dep],
+            exports: [Dep]
+          })
+          export class DepModule {}
+        `,
+      );
+
+      env.write(
+        'test.ts',
+        `
+          import {Component, ViewChild, Input, Inject} from '@angular/core';
+          import {DepModule} from './dep';
+
+          @Component({
+            selector: 'cmp',
+            standalone: true,
+            template: '<div dep><div>',
+            imports: [DepModule],
+          })
+          export class Cmp {}
+        `,
+      );
+
+      env.driveMain();
+
+      const jsContents = env.getContents('test.js');
+      const hmrContents = env.driveHmr('test.ts', 'Cmp');
+      expect(jsContents).toContain(`import * as i0 from "@angular/core";`);
+      expect(jsContents).toContain(`import * as i1 from "./dep";`);
+      expect(jsContents).toContain('function Cmp_HmrLoad(t) {');
+      expect(jsContents).toContain(
+        'import(/* @vite-ignore */\n"/@ng/component?c=test.ts%40Cmp&t=" + encodeURIComponent(t))',
+      );
+      expect(jsContents).toContain(
+        ').then(m => m.default && i0.ɵɵreplaceMetadata(Cmp, m.default, [i0, i1], ' +
+          '[DepModule, Component]));',
+      );
+      expect(jsContents).toContain('Cmp_HmrLoad(Date.now());');
+      expect(jsContents).toContain(
+        'import.meta.hot && import.meta.hot.on("angular:component-update", ' +
+          'd => d.id === "test.ts%40Cmp" && Cmp_HmrLoad(d.timestamp)',
+      );
+
+      expect(hmrContents).toContain(
+        'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, DepModule, Component) {',
+      );
+      expect(hmrContents).toContain(`const ɵhmr0 = ɵɵnamespaces[0];`);
+      expect(hmrContents).toContain(`const ɵhmr1 = ɵɵnamespaces[1];`);
+      expect(hmrContents).toContain('Cmp.ɵfac = function Cmp_Factory');
+      expect(hmrContents).toContain('Cmp.ɵcmp = /*@__PURE__*/ ɵhmr0.ɵɵdefineComponent');
+      expect(hmrContents).toContain('ɵhmr0.ɵsetClassMetadata(Cmp,');
+      expect(hmrContents).toContain('ɵhmr0.ɵsetClassDebugInfo(Cmp,');
+      expect(hmrContents).toContain('dependencies: [DepModule, ɵhmr1.Dep]');
     });
 
     it('should generate an HMR update function for a component that has embedded views', () => {
@@ -144,10 +215,10 @@ runInEachFileSystem(() => {
       const hmrContents = env.driveHmr('test.ts', 'Cmp');
 
       expect(hmrContents).toContain(
-        'export default function Cmp_UpdateMetadata(Cmp, __ngCore__, Component) {',
+        'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, Component) {',
       );
       expect(hmrContents).toContain('function Cmp_Conditional_0_Template(rf, ctx) {');
-      expect(hmrContents).toContain('__ngCore__.ɵɵtemplate(0, Cmp_Conditional_0_Template, 1, 0);');
+      expect(hmrContents).toContain('ɵhmr0.ɵɵtemplate(0, Cmp_Conditional_0_Template, 1, 0);');
     });
 
     it('should generate an HMR update function for a component whose definition produces variables', () => {
@@ -169,12 +240,12 @@ runInEachFileSystem(() => {
       const hmrContents = env.driveHmr('test.ts', 'Cmp');
 
       expect(hmrContents).toContain(
-        'export default function Cmp_UpdateMetadata(Cmp, __ngCore__, Component) {',
+        'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, Component) {',
       );
       expect(hmrContents).toContain('const _c0 = [[["header"]], "*"];');
       expect(hmrContents).toContain('const _c1 = ["header", "*"];');
       expect(hmrContents).toContain('ngContentSelectors: _c1');
-      expect(hmrContents).toContain('__ngCore__.ɵɵprojectionDef(_c0);');
+      expect(hmrContents).toContain('ɵhmr0.ɵɵprojectionDef(_c0);');
     });
 
     it('should not defer dependencies when HMR is enabled', () => {
@@ -219,11 +290,11 @@ runInEachFileSystem(() => {
       expect(jsContents).toContain('i0.ɵɵdefer(1, 0, Cmp_Defer_1_DepsFn);');
 
       expect(hmrContents).toContain(
-        'export default function Cmp_UpdateMetadata(Cmp, __ngCore__, Component, Dep) {',
+        'export default function Cmp_UpdateMetadata(Cmp, ɵɵnamespaces, Component, Dep) {',
       );
       expect(hmrContents).toContain('const Cmp_Defer_1_DepsFn = () => [Dep];');
       expect(hmrContents).toContain('function Cmp_Defer_0_Template(rf, ctx) {');
-      expect(hmrContents).toContain('__ngCore__.ɵɵdefer(1, 0, Cmp_Defer_1_DepsFn);');
+      expect(hmrContents).toContain('ɵhmr0.ɵɵdefer(1, 0, Cmp_Defer_1_DepsFn);');
       expect(hmrContents).not.toContain('import(');
     });
 
diff --git a/packages/compiler/src/compiler.ts b/packages/compiler/src/compiler.ts
@@ -177,6 +177,7 @@ export {
   compileHmrInitializer,
   compileHmrUpdateCallback,
   R3HmrMetadata,
+  R3HmrNamespaceDependency,
 } from './render3/r3_hmr_compiler';
 export {
   compileFactoryFunction,
diff --git a/packages/compiler/src/render3/r3_hmr_compiler.ts b/packages/compiler/src/render3/r3_hmr_compiler.ts
@@ -21,15 +21,31 @@ export interface R3HmrMetadata {
   /** File path of the component class. */
   filePath: string;
 
-  /** Name under which `@angular/core` should be referred to in the compiled HMR code. */
-  coreName: string;
+  /**
+   * When the compiler generates new imports, they get produced as namespace imports
+   * (e.g. import * as i0 from '@angular/core'). These namespaces have to be captured and passed
+   * along to the update callback.
+   */
+  namespaceDependencies: R3HmrNamespaceDependency[];
 
   /**
    * HMR update functions cannot contain imports so any locals the generated code depends on
    * (e.g. references to imports within the same file or imported symbols) have to be passed in
    * as function parameters. This array contains the names of those local symbols.
    */
-  locals: string[];
+  localDependencies: string[];
+}
+
+/** HMR dependency on a namespace import. */
+export interface R3HmrNamespaceDependency {
+  /** Module name of the import. */
+  moduleName: string;
+
+  /**
+   * Name under which to refer to the namespace inside
+   * HMR-related code. Must be a valid JS identifier.
+   */
+  assignedName: string;
 }
 
 /**
@@ -43,15 +59,18 @@ export function compileHmrInitializer(meta: R3HmrMetadata): o.Expression {
   const dataName = 'd';
   const timestampName = 't';
   const importCallbackName = `${meta.className}_HmrLoad`;
-  const locals = meta.locals.map((localName) => o.variable(localName));
+  const locals = meta.localDependencies.map((localName) => o.variable(localName));
+  const namespaces = meta.namespaceDependencies.map((dep) => {
+    return new o.ExternalExpr({moduleName: dep.moduleName, name: null});
+  });
 
   // m.default
   const defaultRead = o.variable(moduleName).prop('default');
 
-  // ɵɵreplaceMetadata(Comp, m.default, [...]);
+  // ɵɵreplaceMetadata(Comp, m.default, [...namespaces], [...locals]);
   const replaceCall = o
     .importExpr(R3.replaceMetadata)
-    .callFn([meta.type, defaultRead, new o.ExternalExpr(R3.core), o.literalArr(locals)]);
+    .callFn([meta.type, defaultRead, o.literalArr(namespaces), o.literalArr(locals)]);
 
   // (m) => m.default && ɵɵreplaceMetadata(...)
   const replaceCallback = o.arrowFn([new o.FnParam(moduleName)], defaultRead.and(replaceCall));
@@ -133,11 +152,25 @@ export function compileHmrUpdateCallback(
   constantStatements: o.Statement[],
   meta: R3HmrMetadata,
 ): o.DeclareFunctionStmt {
-  // The class name should always be first and core should be second.
-  const params = [meta.className, meta.coreName, ...meta.locals].map(
+  const namespaces = 'ɵɵnamespaces';
+  const params = [meta.className, namespaces, ...meta.localDependencies].map(
     (name) => new o.FnParam(name, o.DYNAMIC_TYPE),
   );
-  const body: o.Statement[] = [...constantStatements];
+  const body: o.Statement[] = [];
+
+  // Declare variables that read out the individual namespaces.
+  for (let i = 0; i < meta.namespaceDependencies.length; i++) {
+    body.push(
+      new o.DeclareVarStmt(
+        meta.namespaceDependencies[i].assignedName,
+        o.variable(namespaces).key(o.literal(i)),
+        o.DYNAMIC_TYPE,
+        o.StmtModifier.Final,
+      ),
+    );
+  }
+
+  body.push(...constantStatements);
 
   for (const field of definitions) {
     if (field.initializer !== null) {
diff --git a/packages/core/src/render3/hmr.ts b/packages/core/src/render3/hmr.ts
@@ -46,14 +46,14 @@ import {RendererFactory} from './interfaces/renderer';
  * Replaces the metadata of a component type and re-renders all live instances of the component.
  * @param type Class whose metadata will be replaced.
  * @param applyMetadata Callback that will apply a new set of metadata on the `type` when invoked.
- * @param environment Core runtime environment to use when applying the HMR update.
+ * @param environment Syntehtic namespace imports that need to be passed along to the callback.
  * @param locals Local symbols from the source location that have to be exposed to the callback.
  * @codeGenApi
  */
 export function ɵɵreplaceMetadata(
   type: Type<unknown>,
-  applyMetadata: (...args: [Type<unknown>, Record<string, unknown>, ...unknown[]]) => void,
-  environment: Record<string, unknown>,
+  applyMetadata: (...args: [Type<unknown>, unknown[], ...unknown[]]) => void,
+  namespaces: unknown[],
   locals: unknown[],
 ) {
   ngDevMode && assertComponentDef(type);
@@ -64,7 +64,7 @@ export function ɵɵreplaceMetadata(
   // can be functions for embedded views, the variables for the constant pool and `setClassMetadata`
   // calls. The callback allows us to keep them isolate from the rest of the app and to invoke
   // them at the right time.
-  applyMetadata.apply(null, [type, environment, ...locals]);
+  applyMetadata.apply(null, [type, namespaces, ...locals]);
 
   // If a `tView` hasn't been created yet, it means that this component hasn't been instantianted
   // before. In this case there's nothing left for us to do aside from patching it in.
diff --git a/packages/core/test/acceptance/hmr_spec.ts b/packages/core/test/acceptance/hmr_spec.ts
@@ -1967,7 +1967,7 @@ describe('hot module replacement', () => {
         (type as any)[ɵNG_COMP_DEF] = null;
         compileComponent(type, metadata);
       },
-      angularCoreEnv,
+      [angularCoreEnv],
       [],
     );
   }

Original file line number	Diff line number	Diff line change
`@@ -1967,7 +1967,7 @@ describe('hot module replacement', () => {`
`1967`	`1967`	`(type as any)[ɵNG_COMP_DEF] = null;`
`1968`	`1968`	`compileComponent(type, metadata);`
`1969`	`1969`	`},`
`1970`		`- angularCoreEnv,`
	`1970`	`+ [angularCoreEnv],`
`1971`	`1971`	`[],`
`1972`	`1972`	`);`
`1973`	`1973`	`}`