angular · dylhunn · Dec 5, 2022 · Jan 5, 2023 · Jan 5, 2023 · Dec 20, 2022
diff --git a/packages/compiler-cli/src/ngtsc/core/src/compiler.ts b/packages/compiler-cli/src/ngtsc/core/src/compiler.ts
@@ -19,6 +19,7 @@ import {IncrementalBuildStrategy, IncrementalCompilation, IncrementalState} from
 import {SemanticSymbol} from '../../incremental/semantic_graph';
 import {generateAnalysis, IndexedComponent, IndexingContext} from '../../indexer';
 import {ComponentResources, CompoundMetadataReader, CompoundMetadataRegistry, DirectiveMeta, DtsMetadataReader, HostDirectivesResolver, LocalMetadataRegistry, MetadataReader, MetadataReaderWithIndex, PipeMeta, ResourceRegistry} from '../../metadata';
+import {NgModuleIndexImpl} from '../../metadata/src/ng_module_index';
 import {PartialEvaluator} from '../../partial_evaluator';
 import {ActivePerfRecorder, DelegatingPerfRecorder, PerfCheckpoint, PerfEvent, PerfPhase} from '../../perf';
 import {FileUpdate, ProgramDriver, UpdateMode} from '../../program_driver';
@@ -971,6 +972,7 @@ export class NgCompiler {
     const localMetaReader: MetadataReaderWithIndex = localMetaRegistry;
     const depScopeReader = new MetadataDtsModuleScopeResolver(dtsReader, aliasingHost);
     const metaReader = new CompoundMetadataReader([localMetaReader, dtsReader]);
+    const ngModuleIndex = new NgModuleIndexImpl(metaReader, localMetaReader);
     const ngModuleScopeRegistry = new LocalModuleScopeRegistry(
         localMetaReader, metaReader, depScopeReader, refEmitter, aliasingHost);
     const standaloneScopeReader =
@@ -1072,7 +1074,7 @@ export class NgCompiler {
     const templateTypeChecker = new TemplateTypeCheckerImpl(
         this.inputProgram, notifyingDriver, traitCompiler, this.getTypeCheckingConfig(), refEmitter,
         reflector, this.adapter, this.incrementalCompilation, metaReader, localMetaReader,
-        scopeReader, typeCheckScopeRegistry, this.delegatingPerfRecorder);
+        ngModuleIndex, scopeReader, typeCheckScopeRegistry, this.delegatingPerfRecorder);
 
     // Only construct the extended template checker if the configuration is valid and usable.
     const extendedTemplateChecker = this.constructionDiagnostics.length === 0 ?

diff --git a/packages/compiler-cli/src/ngtsc/metadata/src/api.ts b/packages/compiler-cli/src/ngtsc/metadata/src/api.ts
@@ -251,10 +251,17 @@ export interface MetadataReader {
 }
 
 /**
- * A MetadataReader which also allows access to the set of all known directive classes.
+ * A MetadataReader which also allows access to the set of all known trait classes.
  */
 export interface MetadataReaderWithIndex extends MetadataReader {
-  getKnown(kind: MetaKind): Iterable<ClassDeclaration>;
+  getKnown(kind: MetaKind): Array<ClassDeclaration>;
+}
+
+/**
+ * An NgModuleIndex allows access to information about traits exported by NgModules.
+ */
+export interface NgModuleIndex {
+  getNgModulesExporting(directiveOrPipe: ClassDeclaration): Array<Reference<ClassDeclaration>>;
 }
 
 /**

diff --git a/packages/compiler-cli/src/ngtsc/metadata/src/ng_module_index.ts b/packages/compiler-cli/src/ngtsc/metadata/src/ng_module_index.ts
@@ -0,0 +1,125 @@
+/**
+ * @license
+ * Copyright Google LLC All Rights Reserved.
+ *
+ * Use of this source code is governed by an MIT-style license that can be
+ * found in the LICENSE file at https://angular.io/license
+ */
+
+import {Reference} from '../../imports';
+import {ClassDeclaration} from '../../reflection';
+
+import {MetadataReader, MetadataReaderWithIndex, MetaKind, NgModuleIndex} from './api';
+
+/**
+ * An index of all NgModules that export or re-export a given trait.
+ */
+export class NgModuleIndexImpl implements NgModuleIndex {
+  constructor(private metaReader: MetadataReader, private localReader: MetadataReaderWithIndex) {}
+
+  // A map from an NgModule's Class Declaration to the "main" reference to that module, aka the one
+  // present in the reader metadata object
+  private ngModuleAuthoritativeReference = new Map<ClassDeclaration, Reference<ClassDeclaration>>();
+  // A map from a Directive/Pipe's class declaration to the class declarations of all re-exporting
+  // NgModules
+  private typeToExportingModules = new Map<ClassDeclaration, Set<ClassDeclaration>>();
+
+  private indexed = false;
+
+  private updateWith<K, V>(cache: Map<K, Set<V>>, key: K, elem: V) {
+    if (cache.has(key)) {
+      cache.get(key)!.add(elem);
+    } else {
+      const set = new Set<V>();
+      set.add(elem);
+      cache.set(key, set);
+    }
+  }
+
+  private index(): void {
+    const seenTypesWithReexports = new Map<ClassDeclaration, Set<ClassDeclaration>>();
+    const locallyDeclaredDirsAndNgModules = [
+      ...this.localReader.getKnown(MetaKind.NgModule),
+      ...this.localReader.getKnown(MetaKind.Directive),
+    ];
+    for (const decl of locallyDeclaredDirsAndNgModules) {
+      // Here it's safe to create a new Reference because these are known local types.
+      this.indexTrait(new Reference(decl), seenTypesWithReexports);
+    }
+    this.indexed = true;
+  }
+
+  private indexTrait(
+      ref: Reference<ClassDeclaration>,
+      seenTypesWithReexports: Map<ClassDeclaration, Set<ClassDeclaration>>): void {
+    if (seenTypesWithReexports.has(ref.node)) {
+      // We've processed this type before.
+      return;
+    }
+    seenTypesWithReexports.set(ref.node, new Set());
+
+    const meta =
+        this.metaReader.getDirectiveMetadata(ref) ?? this.metaReader.getNgModuleMetadata(ref);
+    if (meta === null) {
+      return;
+    }
+
+    // Component + NgModule: recurse into imports
+    if (meta.imports !== null) {
+      for (const childRef of meta.imports) {
+        this.indexTrait(childRef, seenTypesWithReexports);
+      }
+    }
+
+    if (meta.kind === MetaKind.NgModule) {
+      if (!this.ngModuleAuthoritativeReference.has(ref.node)) {
+        this.ngModuleAuthoritativeReference.set(ref.node, ref);
+      }
+
+      for (const childRef of meta.exports) {
+        this.indexTrait(childRef, seenTypesWithReexports);
+
+        const childMeta = this.metaReader.getDirectiveMetadata(childRef) ??
+            this.metaReader.getPipeMetadata(childRef) ??
+            this.metaReader.getNgModuleMetadata(childRef);
+        if (childMeta === null) {
+          continue;
+        }
+
+        switch (childMeta.kind) {
+          case MetaKind.Directive:
+          case MetaKind.Pipe:
+            this.updateWith(this.typeToExportingModules, childRef.node, ref.node);
+            this.updateWith(seenTypesWithReexports, ref.node, childRef.node);
+            break;
+          case MetaKind.NgModule:
+            if (seenTypesWithReexports.has(childRef.node)) {
+              for (const reexported of seenTypesWithReexports.get(childRef.node)!) {
+                this.updateWith(this.typeToExportingModules, reexported, ref.node);
+                this.updateWith(seenTypesWithReexports, ref.node, reexported);
+              }
+            }
+            break;
+        }
+      }
+    }
+  }
+
+  getNgModulesExporting(directiveOrPipe: ClassDeclaration): Array<Reference<ClassDeclaration>> {
+    if (!this.indexed) {
+      this.index();
+    }
+
+    if (!this.typeToExportingModules.has(directiveOrPipe)) {
+      return [];
+    }
+
+    const refs: Array<Reference<ClassDeclaration>> = [];
+    for (const ngModule of this.typeToExportingModules.get(directiveOrPipe)!) {
+      if (this.ngModuleAuthoritativeReference.has(ngModule)) {
+        refs.push(this.ngModuleAuthoritativeReference.get(ngModule)!);
+      }
+    }
+    return refs;
+  }
+}
diff --git a/packages/compiler-cli/src/ngtsc/metadata/src/registry.ts b/packages/compiler-cli/src/ngtsc/metadata/src/registry.ts
@@ -40,14 +40,14 @@ export class LocalMetadataRegistry implements MetadataRegistry, MetadataReaderWi
     this.pipes.set(meta.ref.node, meta);
   }
 
-  getKnown(kind: MetaKind): Iterable<ClassDeclaration> {
+  getKnown(kind: MetaKind): Array<ClassDeclaration> {
     switch (kind) {
       case MetaKind.Directive:
-        return this.directives.keys();
+        return Array.from(this.directives.values()).map(v => v.ref.node);
       case MetaKind.Pipe:
-        return this.pipes.keys();
+        return Array.from(this.pipes.values()).map(v => v.ref.node);
       case MetaKind.NgModule:
-        return this.ngModules.keys();
+        return Array.from(this.ngModules.values()).map(v => v.ref.node);
     }
   }
 }

diff --git a/packages/compiler-cli/src/ngtsc/typecheck/api/checker.ts b/packages/compiler-cli/src/ngtsc/typecheck/api/checker.ts
@@ -149,8 +149,9 @@ export interface TemplateTypeChecker {
   /**
    * In the context of an Angular trait, generate potential imports for a directive.
    */
-  getPotentialImportsFor(directive: PotentialDirective, inComponent: ts.ClassDeclaration):
-      ReadonlyArray<PotentialImport>;
+  getPotentialImportsFor(
+      toImport: PotentialDirective|PotentialPipe,
+      inComponent: ts.ClassDeclaration): ReadonlyArray<PotentialImport>;
 
   /**
    * Get the primary decorator for an Angular class (such as @Component). This does not work for

diff --git a/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts b/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts
@@ -13,7 +13,7 @@ import {ErrorCode, ngErrorCode} from '../../diagnostics';
 import {absoluteFrom, absoluteFromSourceFile, AbsoluteFsPath, getSourceFileOrError} from '../../file_system';
 import {Reference, ReferenceEmitKind, ReferenceEmitter} from '../../imports';
 import {IncrementalBuild} from '../../incremental/api';
-import {DirectiveMeta, MetadataReader, MetadataReaderWithIndex, MetaKind, NgModuleMeta, PipeMeta} from '../../metadata';
+import {DirectiveMeta, MetadataReader, MetadataReaderWithIndex, MetaKind, NgModuleIndex, PipeMeta} from '../../metadata';
 import {PerfCheckpoint, PerfEvent, PerfPhase, PerfRecorder} from '../../perf';
 import {ProgramDriver, UpdateMode} from '../../program_driver';
 import {ClassDeclaration, DeclarationNode, isNamedClassDeclaration, ReflectionHost} from '../../reflection';
@@ -86,6 +86,7 @@ export class TemplateTypeCheckerImpl implements TemplateTypeChecker {
       private priorBuild: IncrementalBuild<unknown, FileTypeCheckingData>,
       private readonly metaReader: MetadataReader,
       private readonly localMetaReader: MetadataReaderWithIndex,
+      private readonly ngModuleIndex: NgModuleIndex,
       private readonly componentScopeReader: ComponentScopeReader,
       private readonly typeCheckScopeRegistry: TypeCheckScopeRegistry,
       private readonly perf: PerfRecorder) {}
@@ -687,46 +688,58 @@ export class TemplateTypeCheckerImpl implements TemplateTypeChecker {
     return scope.ngModule;
   }
 
-  getPotentialImportsFor(toImport: PotentialDirective, inContext: ts.ClassDeclaration):
-      ReadonlyArray<PotentialImport> {
-    // Look up the original reference associated with the trait's ngModule, so we don't lose the
-    // Reference context (such as identifiers). If the trait is standalone, this will be
-    // `undefined`.
-    let ngModuleRef: Reference<ClassDeclaration<DeclarationNode>>|undefined;
-    if (toImport.ngModule !== null) {
-      ngModuleRef = this.metaReader.getNgModuleMetadata(new Reference(toImport.ngModule))?.ref;
+  private emit(
+      kind: PotentialImportKind, refTo: Reference<ClassDeclaration>,
+      inContext: ts.ClassDeclaration): PotentialImport|null {
+    const emittedRef = this.refEmitter.emit(refTo, inContext.getSourceFile());
+    if (emittedRef.kind === ReferenceEmitKind.Failed) {
+      return null;
     }
-    const kind = ngModuleRef ? PotentialImportKind.NgModule : PotentialImportKind.Standalone;
-
-    // Import the ngModule if one exists. Otherwise, import the standalone trait directly.
-    const importTarget = ngModuleRef ?? toImport.ref;
+    const emitted = emittedRef.expression;
+    if (emitted instanceof WrappedNodeExpr) {
+      // An appropriate identifier is already in scope.
+      return {kind, symbolName: emitted.node.text};
+    } else if (
+        emitted instanceof ExternalExpr && emitted.value.moduleName !== null &&
+        emitted.value.name !== null) {
+      return {
+        kind,
+        moduleSpecifier: emitted.value.moduleName,
+        symbolName: emitted.value.name,
+      };
+    }
+    return null;
+  }
 
-    // Using the compiler's ReferenceEmitter, try to emit a reference to the trait.
-    // TODO(dylhunn): In the future, we can use a more sophisticated strategy for generating and
-    // ranking references, such as keeping a record of import specifiers used in existing code.
-    const emittedRef = this.refEmitter.emit(importTarget, inContext.getSourceFile());
-    if (emittedRef.kind === ReferenceEmitKind.Failed) return [];
-    const emittedExpression = emittedRef.expression;
+  getPotentialImportsFor(
+      toImport: PotentialDirective|PotentialPipe,
+      inContext: ts.ClassDeclaration): ReadonlyArray<PotentialImport> {
+    const imports: PotentialImport[] = [];
 
-    // This is not be a true import if an appropriate identifier is already in scope.
-    if (emittedExpression instanceof WrappedNodeExpr) {
-      return [{kind, symbolName: emittedExpression.node.getText()}];
+    const meta = this.metaReader.getDirectiveMetadata(toImport.ref) ??
+        this.metaReader.getPipeMetadata(toImport.ref);
+    if (meta === null) {
+      return imports;
     }
-    // Otherwise, it must be a genuine external expression.
-    if (!(emittedExpression instanceof ExternalExpr)) {
-      return [];
+
+    if (meta.isStandalone) {
+      const emitted = this.emit(PotentialImportKind.Standalone, toImport.ref, inContext);
+      if (emitted !== null) {
+        imports.push(emitted);
+      }
     }
 
-    if (emittedExpression.value.moduleName === null || emittedExpression.value.name === null)
-      return [];
+    const exportingNgModules = this.ngModuleIndex.getNgModulesExporting(meta.ref.node);
+    if (exportingNgModules !== null) {
+      for (const exporter of exportingNgModules) {
+        const emittedRef = this.emit(PotentialImportKind.Standalone, exporter, inContext);
+        if (emittedRef !== null) {
+          imports.push(emittedRef);
+        }
+      }
+    }
 
-    // Extract and return the TS module and identifier names.
-    const preferredImport: PotentialImport = {
-      kind: ngModuleRef ? PotentialImportKind.NgModule : PotentialImportKind.Standalone,
-      moduleSpecifier: emittedExpression.value.moduleName,
-      symbolName: emittedExpression.value.name,
-    };
-    return [preferredImport];
+    return imports;
   }
 
   private getScopeData(component: ts.ClassDeclaration): ScopeData|null {

diff --git a/packages/compiler-cli/src/ngtsc/typecheck/testing/index.ts b/packages/compiler-cli/src/ngtsc/typecheck/testing/index.ts
@@ -13,7 +13,7 @@ import {absoluteFrom, AbsoluteFsPath, getSourceFileOrError, LogicalFileSystem} f
 import {TestFile} from '../../file_system/testing';
 import {AbsoluteModuleStrategy, LocalIdentifierStrategy, LogicalProjectStrategy, ModuleResolver, Reexport, Reference, ReferenceEmitter, RelativePathStrategy} from '../../imports';
 import {NOOP_INCREMENTAL_BUILD} from '../../incremental';
-import {ClassPropertyMapping, CompoundMetadataReader, DirectiveMeta, HostDirectivesResolver, MatchSource, MetadataReader, MetadataReaderWithIndex, MetaKind} from '../../metadata';
+import {ClassPropertyMapping, CompoundMetadataReader, DirectiveMeta, HostDirectivesResolver, MatchSource, MetadataReader, MetadataReaderWithIndex, MetaKind, NgModuleIndex, NgModuleMeta, PipeMeta} from '../../metadata';
 import {NOOP_PERF_RECORDER} from '../../perf';
 import {TsCreateProgramDriver} from '../../program_driver';
 import {ClassDeclaration, isNamedClassDeclaration, TypeScriptReflectionHost} from '../../reflection';
@@ -550,14 +550,15 @@ export function setup(targets: TypeCheckingTarget[], overrides: {
   };
 
   const fakeMetadataReader = getFakeMetadataReader(fakeMetadataRegistry);
+  const fakeNgModuleIndex = getFakeNgModuleIndex(fakeMetadataRegistry);
   const typeCheckScopeRegistry = new TypeCheckScopeRegistry(
       fakeScopeReader, new CompoundMetadataReader([fakeMetadataReader]),
       new HostDirectivesResolver(fakeMetadataReader));
 
   const templateTypeChecker = new TemplateTypeCheckerImpl(
       program, programStrategy, checkAdapter, fullConfig, emitter, reflectionHost, host,
-      NOOP_INCREMENTAL_BUILD, fakeMetadataReader, fakeMetadataReader, fakeScopeReader,
-      typeCheckScopeRegistry, NOOP_PERF_RECORDER);
+      NOOP_INCREMENTAL_BUILD, fakeMetadataReader, fakeMetadataReader, fakeNgModuleIndex,
+      fakeScopeReader, typeCheckScopeRegistry, NOOP_PERF_RECORDER);
   return {
     templateTypeChecker,
     program,
@@ -577,17 +578,24 @@ function getFakeMetadataReader(fakeMetadataRegistry: Map<any, DirectiveMeta|null
         null {
           return fakeMetadataRegistry.get(node.debugName) ?? null;
         },
-    getKnown(kind: MetaKind): Iterable<ClassDeclaration> {
+    getKnown(kind: MetaKind): Array<ClassDeclaration> {
       switch (kind) {
-        case MetaKind.Directive:
-          return fakeMetadataRegistry.keys();
+        // TODO: This is not needed for these ngtsc tests, but may be wanted in the future.
         default:
           return [];
       }
     }
   } as MetadataReaderWithIndex;
 }
 
+function getFakeNgModuleIndex(fakeMetadataRegistry: Map<any, DirectiveMeta|null>): NgModuleIndex {
+  return {
+    getNgModulesExporting(trait: ClassDeclaration): Array<Reference<ClassDeclaration>> {
+      return [];
+    }
+  } as NgModuleIndex;
+}
+
 type DeclarationResolver = (decl: TestDeclaration) => ClassDeclaration<ts.ClassDeclaration>;
 
 function prepareDeclarations(