perf(@angular/build): cache pre-declaration transformed SourceFiles during AoT emit

packages/angular/build/src/tools/angular/compilation/angular-compilation_spec.ts

@@ -8,6 +8,7 @@
 
 import ts from 'typescript';
 import type { AngularHostOptions } from '../angular-host';
+import { ensurePreDeclarationTransformers } from './aot-compilation';
 import { CompilerOptionOverrides, transformCompilerOptions } from './compiler-options';
 import { TypeScriptCompilation } from './typescript-compilation';
 import {
@@ -480,4 +481,141 @@ describe('AngularCompilation', () => {
       expect(compilerOptions.composite).toBe(false);
     });
   });
+
+  describe('ensurePreDeclarationTransformers', () => {
+    function createTestProgram(
+      fileName: string,
+      content: string,
+    ): {
+      program: ts.Program;
+      sourceFile: ts.SourceFile;
+    } {
+      const options: ts.CompilerOptions = {
+        target: ts.ScriptTarget.ES2022,
+        module: ts.ModuleKind.ES2022,
+        declaration: true,
+        noLib: true,
+      };
+      const sourceFile = ts.createSourceFile(fileName, content, ts.ScriptTarget.ES2022, true);
+      const host = ts.createCompilerHost(options);
+      host.getSourceFile = (name) => (name === fileName ? sourceFile : undefined);
+      host.fileExists = (name) => name === fileName;
+      host.readFile = (name) => (name === fileName ? content : undefined);
+      host.writeFile = () => {};
+
+      const program = ts.createProgram([fileName], options, host);
+
+      return { program, sourceFile };
+    }
+
+    it('transforms each SourceFile only once when emitting .d.ts followed by .js and reuses the cached SourceFile', () => {
+      const { program, sourceFile } = createTestProgram('/index.ts', 'export const value = 1;');
+
+      let preDeclStatePopulated = false;
+      const beforeSpy1 = jasmine
+        .createSpy('beforeTransformer1')
+        .and.callFake((sf: ts.SourceFile) => {
+          preDeclStatePopulated = true;
+          const addedStmt = ts.factory.createVariableStatement(
+            [ts.factory.createModifier(ts.SyntaxKind.ExportKeyword)],
+            ts.factory.createVariableDeclarationList(
+              [
+                ts.factory.createVariableDeclaration(
+                  'fromBefore1',
+                  undefined,
+                  undefined,
+                  ts.factory.createNumericLiteral(2),
+                ),
+              ],
+              ts.NodeFlags.Const,
+            ),
+          );
+
+          return ts.factory.updateSourceFile(sf, [...sf.statements, addedStmt]);
+        });
+
+      const beforeSpy2 = jasmine
+        .createSpy('beforeTransformer2')
+        .and.callFake((sf: ts.SourceFile) => {
+          const addedStmt = ts.factory.createVariableStatement(
+            [ts.factory.createModifier(ts.SyntaxKind.ExportKeyword)],
+            ts.factory.createVariableDeclarationList(
+              [
+                ts.factory.createVariableDeclaration(
+                  'fromBefore2',
+                  undefined,
+                  undefined,
+                  ts.factory.createNumericLiteral(3),
+                ),
+              ],
+              ts.NodeFlags.Const,
+            ),
+          );
+
+          return ts.factory.updateSourceFile(sf, [...sf.statements, addedStmt]);
+        });
+
+      let sawPopulatedPreDeclState = false;
+      const afterDeclSpy = jasmine
+        .createSpy('afterDeclarationsTransformer')
+        .and.callFake((sf: ts.SourceFile) => {
+          sawPopulatedPreDeclState = preDeclStatePopulated;
+
+          return sf;
+        });
+
+      const transformers: ts.CustomTransformers = {
+        before: [() => beforeSpy1, () => beforeSpy2],
+        afterDeclarations: [() => afterDeclSpy],
+      };
+
+      ensurePreDeclarationTransformers(transformers);
+
+      const emittedFiles = new Map<string, string>();
+      const writeFile: ts.WriteFileCallback = (fileName, text) => {
+        emittedFiles.set(fileName, text);
+      };
+
+      // Emit .d.ts first (emitOnlyDtsFiles = true)
+      program.emit(sourceFile, writeFile, undefined, true, transformers);
+
+      expect(sawPopulatedPreDeclState).toBeTrue();
+      expect(beforeSpy1).toHaveBeenCalledTimes(1);
+      expect(beforeSpy2).toHaveBeenCalledTimes(1);
+      expect(afterDeclSpy).toHaveBeenCalledTimes(1);
+      expect(emittedFiles.has('/index.d.ts')).toBeTrue();
+      expect(emittedFiles.has('/index.js')).toBeFalse();
+
+      // Emit .js second; before transformers must not execute again and cached AST must be emitted
+      program.emit(sourceFile, writeFile, undefined, 0 as unknown as boolean, transformers);
+
+      expect(beforeSpy1).toHaveBeenCalledTimes(1);
+      expect(beforeSpy2).toHaveBeenCalledTimes(1);
+      expect(emittedFiles.get('/index.js')).toContain('export const fromBefore1 = 2;');
+      expect(emittedFiles.get('/index.js')).toContain('export const fromBefore2 = 3;');
+    });
+
+    it('transforms each SourceFile only once when emitting .js and .d.ts together', () => {
+      const { program, sourceFile } = createTestProgram('/index.ts', 'export const value = 1;');
+
+      const beforeSpy = jasmine
+        .createSpy('beforeTransformer')
+        .and.callFake((sf: ts.SourceFile) => sf);
+      const afterDeclSpy = jasmine
+        .createSpy('afterDeclarationsTransformer')
+        .and.callFake((sf: ts.SourceFile) => sf);
+
+      const transformers: ts.CustomTransformers = {
+        before: [() => beforeSpy],
+        afterDeclarations: [() => afterDeclSpy],
+      };
+
+      ensurePreDeclarationTransformers(transformers);
+
+      program.emit(sourceFile, () => {}, undefined, false, transformers);
+
+      expect(beforeSpy).toHaveBeenCalledTimes(1);
+      expect(afterDeclSpy).toHaveBeenCalledTimes(1);
+    });
+  });
 });

packages/angular/build/src/tools/angular/compilation/aot-compilation.ts

@@ -657,40 +657,60 @@ export class AotCompilation extends TypeScriptCompilation {
  * `afterDeclarations`, and caches the transformed `SourceFile` so a subsequent `.js` emit for the
  * same file does not re-execute `before` transformers.
  */
-function ensurePreDeclarationTransformers(transformers: ts.CustomTransformers): void {
+export function ensurePreDeclarationTransformers(transformers: ts.CustomTransformers): void {
   if (!transformers.before?.length || !transformers.afterDeclarations?.length) {
     return;
   }
 
   const beforeTransformers = transformers.before as ts.TransformerFactory<ts.SourceFile>[];
-  const transformedSourceFiles = new Set<ts.SourceFile>();
-  transformers.before = beforeTransformers.map((factory, index) =>
-    index === 0
-      ? (context) => {
-          const transform = factory(context);
+  const seenSourceFiles = new Set<ts.SourceFile>();
+  const preTransformedSourceFiles = new Map<ts.SourceFile, ts.SourceFile>();
+  const runBefore = (
+    beforeTransforms: ts.Transformer<ts.SourceFile>[],
+    sourceFile: ts.SourceFile,
+    originalSourceFile: ts.SourceFile,
+  ): ts.SourceFile => {
+    seenSourceFiles.add(originalSourceFile);
+    let current = sourceFile;
+    for (const transform of beforeTransforms) {
+      current = transform(current);
+    }
+
+    return current;
+  };
+
+  transformers.before = [
+    (context) => {
+      let beforeTransforms: ts.Transformer<ts.SourceFile>[] | undefined;
 
-          return (sourceFile) => {
-            transformedSourceFiles.add(ts.getOriginalNode(sourceFile, ts.isSourceFile));
+      return (sourceFile) => {
+        const originalSourceFile = ts.getOriginalNode(sourceFile, ts.isSourceFile);
+        const preTransformed = preTransformedSourceFiles.get(originalSourceFile);
+        if (preTransformed !== undefined) {
+          preTransformedSourceFiles.delete(originalSourceFile);
 
-            return transform(sourceFile);
-          };
+          return preTransformed;
         }
-      : factory,
-  );
+
+        beforeTransforms ??= beforeTransformers.map((factory) => factory(context));
+
+        return runBefore(beforeTransforms, sourceFile, originalSourceFile);
+      };
+    },
+  ];
   transformers.afterDeclarations = [
     (context) => {
       let beforeTransforms: ts.Transformer<ts.SourceFile>[] | undefined;
 
       return (fileOrBundle) => {
         if (!ts.isBundle(fileOrBundle)) {
           const originalSourceFile = ts.getOriginalNode(fileOrBundle, ts.isSourceFile);
-          if (!transformedSourceFiles.has(originalSourceFile)) {
-            transformedSourceFiles.add(originalSourceFile);
+          if (!seenSourceFiles.has(originalSourceFile)) {
             beforeTransforms ??= beforeTransformers.map((factory) => factory(context));
-            let current = originalSourceFile;
-            for (const transform of beforeTransforms) {
-              current = transform(current);
-            }
+            preTransformedSourceFiles.set(
+              originalSourceFile,
+              runBefore(beforeTransforms, originalSourceFile, originalSourceFile),
+            );
           }
         }