refactor(ivy): abstract .d.ts file transformations (#34235)

This commit refactors the way the compiler transforms .d.ts files during
ngtsc builds. Previously the `IvyCompilation` kept track of a
`DtsFileTransformer` for each input file. Now, any number of
`DtsTransform` operations that need to be applied to a .d.ts file are
collected in the `DtsTransformRegistry`. These are then ran using a
single `DtsTransformer` so that multiple transforms can be applied
efficiently.

PR Close #34235

packages/compiler-cli/src/ngtsc/program.ts

@@ -32,6 +32,7 @@ import {ComponentScopeReader, CompoundComponentScopeReader, LocalModuleScopeRegi
 import {FactoryGenerator, FactoryInfo, GeneratedShimsHostWrapper, ShimGenerator, SummaryGenerator, TypeCheckShimGenerator, generatedFactoryTransform} from './shims';
 import {ivySwitchTransform} from './switch';
 import {IvyCompilation, declarationTransformFactory, ivyTransformFactory} from './transform';
+import {DtsTransformRegistry} from './transform';
 import {aliasTransformFactory} from './transform/src/alias';
 import {TypeCheckContext, TypeCheckingConfig, typeCheckFilePath} from './typecheck';
 import {normalizeSeparators} from './util/src/path';
@@ -68,8 +69,8 @@ export class NgtscProgram implements api.Program {
   private perfTracker: PerfTracker|null = null;
   private incrementalDriver: IncrementalDriver;
   private typeCheckFilePath: AbsoluteFsPath;
-
   private modifiedResourceFiles: Set<string>|null;
+  private dtsTransforms: DtsTransformRegistry|null = null;
 
   constructor(
       rootNames: ReadonlyArray<string>, private options: api.CompilerOptions,
@@ -369,9 +370,12 @@ export class NgtscProgram implements api.Program {
       aliasTransformFactory(compilation.exportStatements) as ts.TransformerFactory<ts.SourceFile>,
       this.defaultImportTracker.importPreservingTransformer(),
     ];
-    const afterDeclarationsTransforms = [
+
-      declarationTransformFactory(compilation),
+    const afterDeclarationsTransforms: ts.TransformerFactory<ts.Bundle|ts.SourceFile>[] = [];
-    ];
+    if (this.dtsTransforms !== null) {
+      afterDeclarationsTransforms.push(
+          declarationTransformFactory(this.dtsTransforms, this.importRewriter));
+    }
 
     // Only add aliasing re-exports to the .d.ts output if the `AliasingHost` requests it.
     if (this.aliasingHost !== null && this.aliasingHost.aliasExportsInDts) {
@@ -617,6 +621,8 @@ export class NgtscProgram implements api.Program {
 
     this.routeAnalyzer = new NgModuleRouteAnalyzer(this.moduleResolver, evaluator);
 
+    this.dtsTransforms = new DtsTransformRegistry();
+
     // Set up the IvyCompilation, which manages state for the Ivy transformer.
     const handlers = [
       new ComponentDecoratorHandler(
@@ -645,7 +651,7 @@ export class NgtscProgram implements api.Program {
     return new IvyCompilation(
         handlers, this.reflector, this.importRewriter, this.incrementalDriver, this.perfRecorder,
         this.sourceToFactorySymbols, scopeRegistry,
-        this.options.compileNonExportedClasses !== false);
+        this.options.compileNonExportedClasses !== false, this.dtsTransforms);
   }
 
   private get reflector(): TypeScriptReflectionHost {

packages/compiler-cli/src/ngtsc/transform/index.ts

@@ -8,5 +8,5 @@
 
 export * from './src/api';
 export {IvyCompilation} from './src/compilation';
-export {DtsFileTransformer, declarationTransformFactory} from './src/declaration';
+export {declarationTransformFactory, DtsTransformRegistry, IvyDeclarationDtsTransform} from './src/declaration';
 export {ivyTransformFactory} from './src/transform';

packages/compiler-cli/src/ngtsc/transform/src/api.ts

@@ -12,6 +12,7 @@ import * as ts from 'typescript';
 import {Reexport} from '../../imports';
 import {IndexingContext} from '../../indexer';
 import {ClassDeclaration, Decorator} from '../../reflection';
+import {ImportManager} from '../../translator';
 import {TypeCheckContext} from '../../typecheck';
 
 export enum HandlerPrecedence {
@@ -156,3 +157,12 @@ export interface ResolveResult {
   reexports?: Reexport[];
   diagnostics?: ts.Diagnostic[];
 }
+
+export interface DtsTransform {
+  transformClassElement?(element: ts.ClassElement, imports: ImportManager): ts.ClassElement;
+  transformFunctionDeclaration?
+      (element: ts.FunctionDeclaration, imports: ImportManager): ts.FunctionDeclaration;
+  transformClass?
+      (clazz: ts.ClassDeclaration, elements: ReadonlyArray<ts.ClassElement>,
+       imports: ImportManager): ts.ClassDeclaration;
+}

packages/compiler-cli/src/ngtsc/transform/src/compilation.ts

@@ -14,13 +14,13 @@ import {ImportRewriter} from '../../imports';
 import {IncrementalDriver} from '../../incremental';
 import {IndexingContext} from '../../indexer';
 import {PerfRecorder} from '../../perf';
-import {ClassDeclaration, ReflectionHost, isNamedClassDeclaration, reflectNameOfDeclaration} from '../../reflection';
+import {ClassDeclaration, ReflectionHost, isNamedClassDeclaration} from '../../reflection';
 import {LocalModuleScopeRegistry} from '../../scope';
 import {TypeCheckContext} from '../../typecheck';
 import {getSourceFile, isExported} from '../../util/src/typescript';
 
 import {AnalysisOutput, CompileResult, DecoratorHandler, DetectResult, HandlerPrecedence} from './api';
-import {DtsFileTransformer} from './declaration';
+import {DtsTransformRegistry} from './declaration';
 
 const EMPTY_ARRAY: any = [];
 
@@ -54,19 +54,9 @@ export class IvyCompilation {
    */
   private ivyClasses = new Map<ClassDeclaration, IvyClass>();
 
-  /**
-   * Tracks factory information which needs to be generated.
-   */
-
-  /**
-   * Tracks the `DtsFileTransformer`s for each TS file that needs .d.ts transformations.
-   */
-  private dtsMap = new Map<string, DtsFileTransformer>();
-
   private reexportMap = new Map<string, Map<string, [string, string]>>();
   private _diagnostics: ts.Diagnostic[] = [];
 
-
   /**
    * @param handlers array of `DecoratorHandler`s which will be executed against each class in the
    * program
@@ -80,9 +70,8 @@ export class IvyCompilation {
       private handlers: DecoratorHandler<any, any>[], private reflector: ReflectionHost,
       private importRewriter: ImportRewriter, private incrementalDriver: IncrementalDriver,
       private perf: PerfRecorder, private sourceToFactorySymbols: Map<string, Set<string>>|null,
-      private scopeRegistry: LocalModuleScopeRegistry, private compileNonExportedClasses: boolean) {
+      private scopeRegistry: LocalModuleScopeRegistry, private compileNonExportedClasses: boolean,
-  }
+      private dtsTransforms: DtsTransformRegistry) {}
-
 
   get exportStatements(): Map<string, Map<string, [string, string]>> { return this.reexportMap; }
 
@@ -392,9 +381,8 @@ export class IvyCompilation {
 
     // Look up the .d.ts transformer for the input file and record that at least one field was
     // generated, which will allow the .d.ts to be transformed later.
-    const fileName = original.getSourceFile().fileName;
+    this.dtsTransforms.getIvyDeclarationTransform(original.getSourceFile())
-    const dtsTransformer = this.getDtsTransformer(fileName);
+        .addFields(original, res);
-    dtsTransformer.recordStaticField(reflectNameOfDeclaration(node) !, res);
 
     // Return the instruction to the transformer so the fields will be added.
     return res.length > 0 ? res : undefined;
@@ -424,30 +412,5 @@ export class IvyCompilation {
     return decorators;
   }
 
-  /**
-   * Process a declaration file and return a transformed version that incorporates the changes
-   * made to the source file.
-   */
-  transformedDtsFor(file: ts.SourceFile, context: ts.TransformationContext): ts.SourceFile {
-    // No need to transform if it's not a declarations file, or if no changes have been requested
-    // to the input file.
-    // Due to the way TypeScript afterDeclarations transformers work, the SourceFile path is the
-    // same as the original .ts.
-    // The only way we know it's actually a declaration file is via the isDeclarationFile property.
-    if (!file.isDeclarationFile || !this.dtsMap.has(file.fileName)) {
-      return file;
-    }
-
-    // Return the transformed source.
-    return this.dtsMap.get(file.fileName) !.transform(file, context);
-  }
-
   get diagnostics(): ReadonlyArray<ts.Diagnostic> { return this._diagnostics; }
-
-  private getDtsTransformer(tsFileName: string): DtsFileTransformer {
-    if (!this.dtsMap.has(tsFileName)) {
-      this.dtsMap.set(tsFileName, new DtsFileTransformer(this.importRewriter));
-    }
-    return this.dtsMap.get(tsFileName) !;
-  }
 }

packages/compiler-cli/src/ngtsc/transform/src/declaration.ts

@@ -6,26 +6,67 @@
  * found in the LICENSE file at https://angular.io/license
  */
 
+import {Type} from '@angular/compiler';
 import * as ts from 'typescript';
 
 import {ImportRewriter} from '../../imports';
+import {ClassDeclaration} from '../../reflection';
 import {ImportManager, translateType} from '../../translator';
 
-import {CompileResult} from './api';
+import {DtsTransform} from './api';
-import {IvyCompilation} from './compilation';
 import {addImports} from './utils';
 
+/**
+ * Keeps track of `DtsTransform`s per source file, so that it is known which source files need to
+ * have their declaration file transformed.
+ */
+export class DtsTransformRegistry {
+  private ivyDeclarationTransforms = new Map<string, IvyDeclarationDtsTransform>();
 
+  getIvyDeclarationTransform(sf: ts.SourceFile): IvyDeclarationDtsTransform {
+    if (!this.ivyDeclarationTransforms.has(sf.fileName)) {
+      this.ivyDeclarationTransforms.set(sf.fileName, new IvyDeclarationDtsTransform());
+    }
+    return this.ivyDeclarationTransforms.get(sf.fileName) !;
+  }
 
-export function declarationTransformFactory(compilation: IvyCompilation):
+  /**
-    ts.TransformerFactory<ts.Bundle|ts.SourceFile> {
+   * Gets the dts transforms to be applied for the given source file, or `null` if no transform is
+   * necessary.
+   */
+  getAllTransforms(sf: ts.SourceFile): DtsTransform[]|null {
+    // No need to transform if it's not a declarations file, or if no changes have been requested
+    // to the input file. Due to the way TypeScript afterDeclarations transformers work, the
+    // `ts.SourceFile` path is the same as the original .ts. The only way we know it's actually a
+    // declaration file is via the `isDeclarationFile` property.
+    if (!sf.isDeclarationFile) {
+      return null;
+    }
+
+    let transforms: DtsTransform[]|null = null;
+    if (this.ivyDeclarationTransforms.has(sf.fileName)) {
+      transforms = [];
+      transforms.push(this.ivyDeclarationTransforms.get(sf.fileName) !);
+    }
+    return transforms;
+  }
+}
+
+export function declarationTransformFactory(
+    transformRegistry: DtsTransformRegistry, importRewriter: ImportRewriter,
+    importPrefix?: string): ts.TransformerFactory<ts.Bundle|ts.SourceFile> {
   return (context: ts.TransformationContext) => {
+    const transformer = new DtsTransformer(context, importRewriter, importPrefix);
     return (fileOrBundle) => {
       if (ts.isBundle(fileOrBundle)) {
         // Only attempt to transform source files.
         return fileOrBundle;
       }
-      return compilation.transformedDtsFor(fileOrBundle, context);
+      const transforms = transformRegistry.getAllTransforms(fileOrBundle);
+      if (transforms === null) {
+        return fileOrBundle;
+      }
+      return transformer.transform(fileOrBundle, transforms);
     };
   };
 }
@@ -33,47 +74,140 @@ export function declarationTransformFactory(compilation: IvyCompilation):
 /**
  * Processes .d.ts file text and adds static field declarations, with types.
  */
-export class DtsFileTransformer {
+class DtsTransformer {
-  private ivyFields = new Map<string, CompileResult[]>();
+  constructor(
-  private imports: ImportManager;
+      private ctx: ts.TransformationContext, private importRewriter: ImportRewriter,
-
+      private importPrefix?: string) {}
-  constructor(private importRewriter: ImportRewriter, importPrefix?: string) {
-    this.imports = new ImportManager(importRewriter, importPrefix);
-  }
-
-  /**
-   * Track that a static field was added to the code for a class.
-   */
-  recordStaticField(name: string, decls: CompileResult[]): void { this.ivyFields.set(name, decls); }
 
   /**
    * Transform the declaration file and add any declarations which were recorded.
    */
-  transform(file: ts.SourceFile, context: ts.TransformationContext): ts.SourceFile {
+  transform(sf: ts.SourceFile, transforms: DtsTransform[]): ts.SourceFile {
+    const imports = new ImportManager(this.importRewriter, this.importPrefix);
+
     const visitor: ts.Visitor = (node: ts.Node): ts.VisitResult<ts.Node> => {
-      // This class declaration needs to have fields added to it.
+      if (ts.isClassDeclaration(node)) {
-      if (ts.isClassDeclaration(node) && node.name !== undefined &&
+        return this.transformClassDeclaration(node, transforms, imports);
-          this.ivyFields.has(node.name.text)) {
+      } else if (ts.isFunctionDeclaration(node)) {
-        const decls = this.ivyFields.get(node.name.text) !;
+        return this.transformFunctionDeclaration(node, transforms, imports);
-        const newMembers = decls.map(decl => {
+      } else {
+        // Otherwise return node as is.
+        return ts.visitEachChild(node, visitor, this.ctx);
+      }
+    };
+
+    // Recursively scan through the AST and process all nodes as desired.
+    sf = ts.visitNode(sf, visitor);
+
+    // Add new imports for this file.
+    return addImports(imports, sf);
+  }
+
+  private transformClassDeclaration(
+      clazz: ts.ClassDeclaration, transforms: DtsTransform[],
+      imports: ImportManager): ts.ClassDeclaration {
+    let elements: ts.ClassElement[]|ReadonlyArray<ts.ClassElement> = clazz.members;
+    let elementsChanged = false;
+
+    for (const transform of transforms) {
+      if (transform.transformClassElement !== undefined) {
+        for (let i = 0; i < elements.length; i++) {
+          const res = transform.transformClassElement(elements[i], imports);
+          if (res !== elements[i]) {
+            if (!elementsChanged) {
+              elements = [...elements];
+              elementsChanged = true;
+            }
+            (elements as ts.ClassElement[])[i] = res;
+          }
+        }
+      }
+    }
+
+    let newClazz: ts.ClassDeclaration = clazz;
+
+    for (const transform of transforms) {
+      if (transform.transformClass !== undefined) {
+        // If no DtsTransform has changed the class yet, then the (possibly mutated) elements have
+        // not yet been incorporated. Otherwise, `newClazz.members` holds the latest class members.
+        const inputMembers = (clazz === newClazz ? elements : newClazz.members);
+
+        newClazz = transform.transformClass(newClazz, inputMembers, imports);
+      }
+    }
+
+    // If some elements have been transformed but the class itself has not been transformed, create
+    // an updated class declaration with the updated elements.
+    if (elementsChanged && clazz === newClazz) {
+      newClazz = ts.updateClassDeclaration(
+          /* node */ clazz,
+          /* decorators */ clazz.decorators,
+          /* modifiers */ clazz.modifiers,
+          /* name */ clazz.name,
+          /* typeParameters */ clazz.typeParameters,
+          /* heritageClauses */ clazz.heritageClauses,
+          /* members */ elements);
+    }
+
+    return newClazz;
+  }
+
+  private transformFunctionDeclaration(
+      declaration: ts.FunctionDeclaration, transforms: DtsTransform[],
+      imports: ImportManager): ts.FunctionDeclaration {
+    let newDecl = declaration;
+
+    for (const transform of transforms) {
+      if (transform.transformFunctionDeclaration !== undefined) {
+        newDecl = transform.transformFunctionDeclaration(newDecl, imports);
+      }
+    }
+
+    return newDecl;
+  }
+}
+
+export interface IvyDeclarationField {
+  name: string;
+  type: Type;
+}
+
+export class IvyDeclarationDtsTransform implements DtsTransform {
+  private declarationFields = new Map<ClassDeclaration, IvyDeclarationField[]>();
+
+  addFields(decl: ClassDeclaration, fields: IvyDeclarationField[]): void {
+    this.declarationFields.set(decl, fields);
+  }
+
+  transformClass(
+      clazz: ts.ClassDeclaration, members: ReadonlyArray<ts.ClassElement>,
+      imports: ImportManager): ts.ClassDeclaration {
+    const original = ts.getOriginalNode(clazz) as ClassDeclaration;
+
+    if (!this.declarationFields.has(original)) {
+      return clazz;
+    }
+    const fields = this.declarationFields.get(original) !;
+
+    const newMembers = fields.map(decl => {
       const modifiers = [ts.createModifier(ts.SyntaxKind.StaticKeyword)];
-          const typeRef = translateType(decl.type, this.imports);
+      const typeRef = translateType(decl.type, imports);
-          return ts.createProperty(undefined, modifiers, decl.name, undefined, typeRef, undefined);
+      return ts.createProperty(
+          /* decorators */ undefined,
+          /* modifiers */ modifiers,
+          /* name */ decl.name,
+          /* questionOrExclamationToken */ undefined,
+          /* type */ typeRef,
+          /* initializer */ undefined);
     });
 
     return ts.updateClassDeclaration(
-            node, node.decorators, node.modifiers, node.name, node.typeParameters,
+        /* node */ clazz,
-            node.heritageClauses, [...node.members, ...newMembers]);
+        /* decorators */ clazz.decorators,
-      }
+        /* modifiers */ clazz.modifiers,
-
+        /* name */ clazz.name,
-      // Otherwise return node as is.
+        /* typeParameters */ clazz.typeParameters,
-      return ts.visitEachChild(node, visitor, context);
+        /* heritageClauses */ clazz.heritageClauses,
-    };
+        /* members */[...members, ...newMembers]);
-
-    // Recursively scan through the AST and add all class members needed.
-    const sf = ts.visitNode(file, visitor);
-
-    // Add new imports for this file.
-    return addImports(this.imports, sf);
   }
 }