fix(compiler): switch to 'referencedFiles' for shim generation (#36211)

Shim generation was built on a lie.

Shims are files added to the program which aren't original files authored by
the user, but files authored effectively by the compiler. These fall into
two categories: files which will be generated (like the .ngfactory shims we
generate for View Engine compatibility) as well as files used internally in
compilation (like the __ng_typecheck__.ts file).

Previously, shim generation was driven by the `rootFiles` passed to the
compiler as input. These are effectively the `files` listed in the
`tsconfig.json`. Each shim generator (e.g. the `FactoryGenerator`) would
examine the `rootFiles` and produce a list of shim file names which it would
be responsible for generating. These names would then be added to the
`rootFiles` when the program was created.

The fatal flaw here is that `rootFiles` does not always account for all of
the files in the program. In fact, it's quite rare that it does. Users don't
typically specify every file directly in `files`. Instead, they rely on
TypeScript, during program creation, starting with a few root files and
transitively discovering all of the files in the program.

This happens, however, during `ts.createProgram`, which is too late to add
new files to the `rootFiles` list.

As a result, shim generation was only including shims for files actually
listed in the `tsconfig.json` file, and not for the transitive set of files
in the user's program as it should.

This commit completely rewrites shim generation to use a different technique
for adding files to the program, inspired by View Engine's shim generator.
In this new technique, as the program is being created and `ts.SourceFile`s
are being requested from the `NgCompilerHost`, shims for those files are
generated and a reference to them is patched onto the original file's
`ts.SourceFile.referencedFiles`. This causes TS to think that the original
file references the shim, and causes the shim to be included in the program.
The original `referencedFiles` array is saved and restored after program
creation, hiding this little hack from the rest of the system.

The new shim generation engine differentiates between two kinds of shims:
top-level shims (such as the flat module entrypoint file and
__ng_typecheck__.ts) and per-file shims such as ngfactory or ngsummary
files. The former are included via `rootFiles` as before, the latter are
included via the `referencedFiles` of their corresponding original files.

As a result of this change, shims are now correctly generated for all files
in the program, not just the ones named in `tsconfig.json`.

A few mitigating factors prevented this bug from being realized until now:

* in g3, `files` does include the transitive closure of files in the program
* in CLI apps, shims are not really used

This change also makes use of a novel technique for associating information
with source files: the use of an `NgExtension` `Symbol` to patch the
information directly onto the AST object. This is used in several
circumstances:

* For shims, metadata about a `ts.SourceFile`'s status as a shim and its
  origins are held in the extension data.
* For original files, the original `referencedFiles` are stashed in the
  extension data for later restoration.

The main benefit of this technique is a lot less bookkeeping around `Map`s
of `ts.SourceFile`s to various kinds of data, which need to be tracked/
invalidated as part of incremental builds.

This technique is based on designs used internally in the TypeScript
compiler and is serving as a prototype of this design in ngtsc. If it works
well, it could have benefits across the rest of the compiler.

PR Close #36211

packages/compiler-cli/src/ngtsc/core/BUILD.bazel

@@ -29,6 +29,7 @@ ts_library(
         "//packages/compiler-cli/src/ngtsc/routing",
         "//packages/compiler-cli/src/ngtsc/scope",
         "//packages/compiler-cli/src/ngtsc/shims",
+        "//packages/compiler-cli/src/ngtsc/shims:api",
         "//packages/compiler-cli/src/ngtsc/switch",
         "//packages/compiler-cli/src/ngtsc/transform",
         "//packages/compiler-cli/src/ngtsc/typecheck",

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

@@ -145,13 +145,10 @@ export class NgCompiler {
     }
     setIncrementalDriver(tsProgram, this.incrementalDriver);
 
-    this.ignoreForDiagnostics = new Set([
-      this.typeCheckFile,
-      ...host.factoryFiles.map(fileName => getSourceFileOrError(tsProgram, fileName)),
-      ...host.summaryFiles.map(fileName => getSourceFileOrError(tsProgram, fileName)),
-    ]);
+    this.ignoreForDiagnostics =
+        new Set(tsProgram.getSourceFiles().filter(sf => this.host.isShim(sf)));
 
-    this.ignoreForEmit = new Set([this.typeCheckFile]);
+    this.ignoreForEmit = this.host.ignoreForEmit;
   }
 
   /**

packages/compiler-cli/src/ngtsc/core/src/host.ts

@@ -11,10 +11,11 @@ import * as ts from 'typescript';
 import {ErrorCode, ngErrorCode} from '../../diagnostics';
 import {findFlatIndexEntryPoint, FlatIndexGenerator} from '../../entry_point';
 import {AbsoluteFsPath, resolve} from '../../file_system';
-import {FactoryGenerator, FactoryTracker, ShimGenerator, SummaryGenerator, TypeCheckShimGenerator} from '../../shims';
-import {typeCheckFilePath} from '../../typecheck';
+import {FactoryGenerator, FactoryTracker, isShim, ShimAdapter, ShimReferenceTagger, SummaryGenerator} from '../../shims';
+import {PerFileShimGenerator, TopLevelShimGenerator} from '../../shims/api';
+import {typeCheckFilePath, TypeCheckShimGenerator} from '../../typecheck';
 import {normalizeSeparators} from '../../util/src/path';
-import {getRootDirs} from '../../util/src/typescript';
+import {getRootDirs, isDtsPath, isNonDeclarationTsPath} from '../../util/src/typescript';
 import {ExtendedTsCompilerHost, NgCompilerOptions, UnifiedModulesHost} from '../api';
 
 // A persistent source of bugs in CompilerHost delegation has been the addition by TS of new,
@@ -96,34 +97,48 @@ export class NgCompilerHost extends DelegatingCompilerHost implements
   readonly inputFiles: ReadonlyArray<string>;
   readonly rootDirs: ReadonlyArray<AbsoluteFsPath>;
   readonly typeCheckFile: AbsoluteFsPath;
-  readonly factoryFiles: AbsoluteFsPath[];
-  readonly summaryFiles: AbsoluteFsPath[];
+
 
   constructor(
       delegate: ExtendedTsCompilerHost, inputFiles: ReadonlyArray<string>,
-      rootDirs: ReadonlyArray<AbsoluteFsPath>, private shims: ShimGenerator[],
-      entryPoint: AbsoluteFsPath|null, typeCheckFile: AbsoluteFsPath,
-      factoryFiles: AbsoluteFsPath[], summaryFiles: AbsoluteFsPath[],
-      factoryTracker: FactoryTracker|null, diagnostics: ts.Diagnostic[]) {
+      rootDirs: ReadonlyArray<AbsoluteFsPath>, private shimAdapter: ShimAdapter,
+      private shimTagger: ShimReferenceTagger, entryPoint: AbsoluteFsPath|null,
+      typeCheckFile: AbsoluteFsPath, factoryTracker: FactoryTracker|null,
+      diagnostics: ts.Diagnostic[]) {
     super(delegate);
 
     this.factoryTracker = factoryTracker;
     this.entryPoint = entryPoint;
     this.typeCheckFile = typeCheckFile;
-    this.factoryFiles = factoryFiles;
-    this.summaryFiles = summaryFiles;
     this.diagnostics = diagnostics;
-    this.inputFiles = inputFiles;
+    this.inputFiles = [...inputFiles, ...shimAdapter.extraInputFiles];
     this.rootDirs = rootDirs;
   }
 
+  /**
+   * Retrieves a set of `ts.SourceFile`s which should not be emitted as JS files.
+   *
+   * Available after this host is used to create a `ts.Program` (which causes all the files in the
+   * program to be enumerated).
+   */
+  get ignoreForEmit(): Set<ts.SourceFile> {
+    return this.shimAdapter.ignoreForEmit;
+  }
+
+  /**
+   * Performs cleanup that needs to happen after a `ts.Program` has been created using this host.
+   */
+  postProgramCreationCleanup(): void {
+    this.shimTagger.finalize();
+  }
+
   /**
    * Create an `NgCompilerHost` from a delegate host, an array of input filenames, and the full set
    * of TypeScript and Angular compiler options.
    */
   static wrap(
-      delegate: ts.CompilerHost, inputFiles: ReadonlyArray<string>,
-      options: NgCompilerOptions): NgCompilerHost {
+      delegate: ts.CompilerHost, inputFiles: ReadonlyArray<string>, options: NgCompilerOptions,
+      oldProgram: ts.Program|null): NgCompilerHost {
     // TODO(alxhub): remove the fallback to allowEmptyCodegenFiles after verifying that the rest of
     // our build tooling is no longer relying on it.
     const allowEmptyCodegenFiles = options.allowEmptyCodegenFiles || false;
@@ -135,54 +150,41 @@ export class NgCompilerHost extends DelegatingCompilerHost implements
         options.generateNgSummaryShims :
         allowEmptyCodegenFiles;
 
-    let rootFiles = [...inputFiles];
-    let normalizedInputFiles = inputFiles.map(n => resolve(n));
 
-    const generators: ShimGenerator[] = [];
-    let summaryGenerator: SummaryGenerator|null = null;
-    let summaryFiles: AbsoluteFsPath[];
+    const topLevelShimGenerators: TopLevelShimGenerator[] = [];
+    const perFileShimGenerators: PerFileShimGenerator[] = [];
 
     if (shouldGenerateSummaryShims) {
       // Summary generation.
-      summaryGenerator = SummaryGenerator.forRootFiles(normalizedInputFiles);
-      generators.push(summaryGenerator);
-      summaryFiles = summaryGenerator.getSummaryFileNames();
-    } else {
-      summaryFiles = [];
+      perFileShimGenerators.push(new SummaryGenerator());
     }
 
     let factoryTracker: FactoryTracker|null = null;
-    let factoryFiles: AbsoluteFsPath[];
     if (shouldGenerateFactoryShims) {
-      // Factory generation.
-      const factoryGenerator = FactoryGenerator.forRootFiles(normalizedInputFiles);
-      const factoryFileMap = factoryGenerator.factoryFileMap;
+      const factoryGenerator = new FactoryGenerator();
+      perFileShimGenerators.push(factoryGenerator);
 
-      factoryFiles = Array.from(factoryFileMap.keys());
-      rootFiles.push(...factoryFiles);
-      generators.push(factoryGenerator);
-
-      factoryTracker = new FactoryTracker(factoryGenerator);
-    } else {
-      factoryFiles = [];
+      factoryTracker = factoryGenerator;
     }
 
-    // Done separately to preserve the order of factory files before summary files in rootFiles.
-    // TODO(alxhub): validate that this is necessary.
-    rootFiles.push(...summaryFiles);
-
-
     const rootDirs = getRootDirs(delegate, options as ts.CompilerOptions);
 
     const typeCheckFile = typeCheckFilePath(rootDirs);
-    generators.push(new TypeCheckShimGenerator(typeCheckFile));
-    rootFiles.push(typeCheckFile);
+    topLevelShimGenerators.push(new TypeCheckShimGenerator(typeCheckFile));
 
     let diagnostics: ts.Diagnostic[] = [];
 
+    const normalizedTsInputFiles: AbsoluteFsPath[] = [];
+    for (const inputFile of inputFiles) {
+      if (!isNonDeclarationTsPath(inputFile)) {
+        continue;
+      }
+      normalizedTsInputFiles.push(resolve(inputFile));
+    }
+
     let entryPoint: AbsoluteFsPath|null = null;
     if (options.flatModuleOutFile != null && options.flatModuleOutFile !== '') {
-      entryPoint = findFlatIndexEntryPoint(normalizedInputFiles);
+      entryPoint = findFlatIndexEntryPoint(normalizedTsInputFiles);
       if (entryPoint === null) {
         // This error message talks specifically about having a single .ts file in "files". However
         // the actual logic is a bit more permissive. If a single file exists, that will be taken,
@@ -206,37 +208,49 @@ export class NgCompilerHost extends DelegatingCompilerHost implements
         const flatModuleOutFile = normalizeSeparators(options.flatModuleOutFile);
         const flatIndexGenerator =
             new FlatIndexGenerator(entryPoint, flatModuleOutFile, flatModuleId);
-        generators.push(flatIndexGenerator);
-        rootFiles.push(flatIndexGenerator.flatIndexPath);
+        topLevelShimGenerators.push(flatIndexGenerator);
       }
     }
 
+    const shimAdapter = new ShimAdapter(
+        delegate, normalizedTsInputFiles, topLevelShimGenerators, perFileShimGenerators,
+        oldProgram);
+    const shimTagger =
+        new ShimReferenceTagger(perFileShimGenerators.map(gen => gen.extensionPrefix));
     return new NgCompilerHost(
-        delegate, rootFiles, rootDirs, generators, entryPoint, typeCheckFile, factoryFiles,
-        summaryFiles, factoryTracker, diagnostics);
+        delegate, inputFiles, rootDirs, shimAdapter, shimTagger, entryPoint, typeCheckFile,
+        factoryTracker, diagnostics);
+  }
+
+  /**
+   * Check whether the given `ts.SourceFile` is a shim file.
+   *
+   * If this returns false, the file is user-provided.
+   */
+  isShim(sf: ts.SourceFile): boolean {
+    return isShim(sf);
   }
 
   getSourceFile(
       fileName: string, languageVersion: ts.ScriptTarget,
       onError?: ((message: string) => void)|undefined,
       shouldCreateNewSourceFile?: boolean|undefined): ts.SourceFile|undefined {
-    for (let i = 0; i < this.shims.length; i++) {
-      const generator = this.shims[i];
-      // TypeScript internal paths are guaranteed to be POSIX-like absolute file paths.
-      const absoluteFsPath = resolve(fileName);
-      if (generator.recognize(absoluteFsPath)) {
-        const readFile = (originalFile: string) => {
-          return this.delegate.getSourceFile(
-                     originalFile, languageVersion, onError, shouldCreateNewSourceFile) ||
-              null;
-        };
-
-        return generator.generate(absoluteFsPath, readFile) || undefined;
-      }
+    // Is this a previously known shim?
+    const shimSf = this.shimAdapter.maybeGenerate(resolve(fileName));
+    if (shimSf !== null) {
+      // Yes, so return it.
+      return shimSf;
+    }
+
+    // No, so it's a file which might need shims (or a file which doesn't exist).
+    const sf =
+        this.delegate.getSourceFile(fileName, languageVersion, onError, shouldCreateNewSourceFile);
+    if (sf === undefined) {
+      return undefined;
     }
 
-    return this.delegate.getSourceFile(
-        fileName, languageVersion, onError, shouldCreateNewSourceFile);
+    this.shimTagger.tag(sf);
+    return sf;
   }
 
   fileExists(fileName: string): boolean {
@@ -245,8 +259,10 @@ export class NgCompilerHost extends DelegatingCompilerHost implements
     //  2) at least one of the shim generators recognizes it
     // Note that we can pass the file name as branded absolute fs path because TypeScript
     // internally only passes POSIX-like paths.
+    //
+    // Also note that the `maybeGenerate` check below checks for both `null` and `undefined`.
     return this.delegate.fileExists(fileName) ||
-        this.shims.some(shim => shim.recognize(resolve(fileName)));
+        this.shimAdapter.maybeGenerate(resolve(fileName)) != null;
   }
 
   get unifiedModulesHost(): UnifiedModulesHost|null {

packages/compiler-cli/src/ngtsc/core/test/compiler_test.ts

@@ -43,7 +43,7 @@ runInEachFileSystem(() => {
         strictTemplates: true,
       };
       const baseHost = new NgtscCompilerHost(getFileSystem(), options);
-      const host = NgCompilerHost.wrap(baseHost, [COMPONENT], options);
+      const host = NgCompilerHost.wrap(baseHost, [COMPONENT], options, /* oldProgram */ null);
       const program = ts.createProgram({host, options, rootNames: host.inputFiles});
       const compiler = new NgCompiler(host, options, program);
 

packages/compiler-cli/src/ngtsc/core/test/host_spec.ts

@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright Google Inc. 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 {absoluteFrom as _, FileSystem, getFileSystem, NgtscCompilerHost} from '../../file_system';
+import {runInEachFileSystem} from '../../file_system/testing';
+import {NgCompilerOptions} from '../api';
+import {NgCompilerHost} from '../src/host';
+
+runInEachFileSystem(() => {
+  let fs!: FileSystem;
+  beforeEach(() => {
+    fs = getFileSystem();
+    fs.ensureDir(_('/'));
+  });
+
+  describe('NgCompilerHost', () => {
+    it('should add factory shims for all root files', () => {
+      // This test verifies that per-file shims are guaranteed to be generated for each file in the
+      // "root" input files, regardless of whether `referencedFiles`-based shimming is enabled. As
+      // it turns out, setting `noResolve` in TypeScript's options disables this kind of shimming,
+      // so `NgCompilerHost` needs to ensure at least the root files still get shims directly.
+
+      // File is named index.ts to trigger flat module entrypoint logic
+      // (which adds a top-level shim).
+      const fileName = _('/index.ts');
+      fs.writeFile(fileName, `export class Test {}`);
+
+      const options: NgCompilerOptions = {
+        // Using noResolve means that TS will not follow `referencedFiles`, which is how shims are
+        // normally included in the program.
+        noResolve: true,
+        rootDir: _('/'),
+
+        // Both a top-level and a per-file shim are enabled.
+        flatModuleOutFile: './entry',
+        generateNgFactoryShims: true,
+      };
+      const baseHost = new NgtscCompilerHost(fs, options);
+      const host = NgCompilerHost.wrap(baseHost, [fileName], options, null);
+
+      // A shim file should be included for the index.ts ngfactory, but not entry.ts since that
+      // itself is a shim.
+      expect(host.inputFiles).toContain(_('/index.ngfactory.ts'));
+      expect(host.inputFiles).not.toContain(_('/entry.ngfactory.ts'));
+    });
+  });
+});

packages/compiler-cli/src/ngtsc/entry_point/src/generator.ts

@@ -11,11 +11,12 @@
 import * as ts from 'typescript';
 
 import {AbsoluteFsPath, dirname, join} from '../../file_system';
-import {ShimGenerator} from '../../shims';
+import {TopLevelShimGenerator} from '../../shims';
 import {relativePathBetween} from '../../util/src/path';
 
-export class FlatIndexGenerator implements ShimGenerator {
+export class FlatIndexGenerator implements TopLevelShimGenerator {
   readonly flatIndexPath: string;
+  readonly shouldEmit = true;
 
   constructor(
       readonly entryPoint: AbsoluteFsPath, relativeFlatIndexPath: string,
@@ -24,11 +25,7 @@ export class FlatIndexGenerator implements ShimGenerator {
         join(dirname(entryPoint), relativeFlatIndexPath).replace(/\.js$/, '') + '.ts';
   }
 
-  recognize(fileName: string): boolean {
-    return fileName === this.flatIndexPath;
-  }
-
-  generate(): ts.SourceFile {
+  makeTopLevelShim(): ts.SourceFile {
     const relativeEntryPoint = relativePathBetween(this.flatIndexPath, this.entryPoint);
     const contents = `/**
  * Generated bundle index. Do not edit.

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

@@ -66,12 +66,14 @@ export class NgtscProgram implements api.Program {
     }
     this.closureCompilerEnabled = !!options.annotateForClosureCompiler;
 
-    this.host = NgCompilerHost.wrap(delegateHost, rootNames, options);
-
     const reuseProgram = oldProgram && oldProgram.reuseTsProgram;
+    this.host = NgCompilerHost.wrap(delegateHost, rootNames, options, reuseProgram ?? null);
+
     this.tsProgram = ts.createProgram(this.host.inputFiles, options, this.host, reuseProgram);
     this.reuseTsProgram = this.tsProgram;
 
+    this.host.postProgramCreationCleanup();
+
     // Create the NgCompiler which will drive the rest of the compilation.
     this.compiler =
         new NgCompiler(this.host, options, this.tsProgram, reuseProgram, this.perfRecorder);

packages/compiler-cli/src/ngtsc/shims/BUILD.bazel

@@ -2,12 +2,22 @@ load("//tools:defaults.bzl", "ts_library")
 
 package(default_visibility = ["//visibility:public"])
 
+ts_library(
+    name = "api",
+    srcs = ["api.ts"],
+    deps = [
+        "//packages/compiler-cli/src/ngtsc/file_system",
+        "@npm//typescript",
+    ],
+)
+
 ts_library(
     name = "shims",
     srcs = ["index.ts"] + glob([
         "src/**/*.ts",
     ]),
     deps = [
+        ":api",
         "//packages/compiler",
         "//packages/compiler-cli/src/ngtsc/file_system",
         "//packages/compiler-cli/src/ngtsc/imports",