alxhub · Jun 19, 2020
diff --git a/packages/compiler-cli/BUILD.bazel b/packages/compiler-cli/BUILD.bazel
@@ -31,6 +31,7 @@ ts_library(
         "//packages/compiler-cli/src/ngtsc/indexer",
         "//packages/compiler-cli/src/ngtsc/perf",
         "//packages/compiler-cli/src/ngtsc/reflection",
+        "//packages/compiler-cli/src/ngtsc/shims",
         "//packages/compiler-cli/src/ngtsc/typecheck",
         "@npm//@bazel/typescript",
         "@npm//@types/node",

diff --git a/packages/compiler-cli/src/ngtsc/program.ts b/packages/compiler-cli/src/ngtsc/program.ts
@@ -17,6 +17,7 @@ import {NgCompilerOptions} from './core/api';
 import {TrackedIncrementalBuildStrategy} from './incremental';
 import {IndexedComponent} from './indexer';
 import {NOOP_PERF_RECORDER, PerfRecorder, PerfTracker} from './perf';
+import {retagAllTsFiles, untagAllTsFiles} from './shims';
 import {ReusedProgramStrategy} from './typecheck';
 
 
@@ -68,14 +69,26 @@ export class NgtscProgram implements api.Program {
     }
     this.closureCompilerEnabled = !!options.annotateForClosureCompiler;
 
-    const reuseProgram = oldProgram && oldProgram.reuseTsProgram;
+    const reuseProgram = oldProgram?.reuseTsProgram;
     this.host = NgCompilerHost.wrap(delegateHost, rootNames, options, reuseProgram ?? null);
 
+    if (reuseProgram !== undefined) {
+      // Prior to reusing the old program, restore shim tagging for all its `ts.SourceFile`s.
+      // TypeScript checks the `referencedFiles` of `ts.SourceFile`s for changes when evaluating
+      // incremental reuse of data from the old program, so it's important that these match in order
+      // to get the most benefit out of reuse.
+      retagAllTsFiles(reuseProgram);
+    }
+
     this.tsProgram = ts.createProgram(this.host.inputFiles, options, this.host, reuseProgram);
     this.reuseTsProgram = this.tsProgram;
 
     this.host.postProgramCreationCleanup();
 
+    // Shim tagging has served its purpose, and tags can now be removed from all `ts.SourceFile`s in
+    // the program.
+    untagAllTsFiles(this.tsProgram);
+
     const reusedProgramStrategy = new ReusedProgramStrategy(
         this.tsProgram, this.host, this.options, this.host.shimExtensionPrefixes);
 
@@ -93,6 +106,10 @@ export class NgtscProgram implements api.Program {
     return this.tsProgram;
   }
 
+  getReuseTsProgram(): ts.Program {
+    return this.reuseTsProgram;
+  }
+
   getTsOptionDiagnostics(cancellationToken?: ts.CancellationToken|
                          undefined): readonly ts.Diagnostic[] {
     return this.tsProgram.getOptionsDiagnostics(cancellationToken);
@@ -248,6 +265,7 @@ export class NgtscProgram implements api.Program {
       }));
       this.perfRecorder.stop(fileEmitSpan);
     }
+
     this.perfRecorder.stop(emitSpan);
 
     if (this.perfTracker !== null && this.options.tracePerformance !== undefined) {

diff --git a/packages/compiler-cli/src/ngtsc/shims/index.ts b/packages/compiler-cli/src/ngtsc/shims/index.ts
@@ -9,7 +9,7 @@
 /// <reference types="node" />
 
 export {ShimAdapter} from './src/adapter';
-export {copyFileShimData, isShim} from './src/expando';
+export {copyFileShimData, isShim, retagAllTsFiles, retagTsFile, sfExtensionData, untagAllTsFiles, untagTsFile} from './src/expando';
 export {FactoryGenerator, generatedFactoryTransform} from './src/factory_generator';
 export {ShimReferenceTagger} from './src/reference_tagger';
 export {SummaryGenerator} from './src/summary_generator';
diff --git a/packages/compiler-cli/src/ngtsc/shims/src/adapter.ts b/packages/compiler-cli/src/ngtsc/shims/src/adapter.ts
@@ -12,7 +12,7 @@ import {absoluteFrom, absoluteFromSourceFile, AbsoluteFsPath} from '../../file_s
 import {isDtsPath} from '../../util/src/typescript';
 import {PerFileShimGenerator, TopLevelShimGenerator} from '../api';
 
-import {isFileShimSourceFile, isShim, NgExtension, sfExtensionData} from './expando';
+import {isFileShimSourceFile, isShim, sfExtensionData} from './expando';
 import {makeShimFileName} from './util';
 
 interface ShimGeneratorData {

diff --git a/packages/compiler-cli/src/ngtsc/shims/src/expando.ts b/packages/compiler-cli/src/ngtsc/shims/src/expando.ts
@@ -21,7 +21,16 @@ export const NgExtension = Symbol('NgExtension');
 export interface NgExtensionData {
   isTopLevelShim: boolean;
   fileShim: NgFileShimData|null;
+
+  /**
+   * The contents of the `referencedFiles` array, before modification by a `ShimReferenceTagger`.
+   */
   originalReferencedFiles: ReadonlyArray<ts.FileReference>|null;
+
+  /**
+   * The contents of the `referencedFiles` array, after modification by a `ShimReferenceTagger`.
+   */
+  taggedReferenceFiles: ReadonlyArray<ts.FileReference>|null;
 }
 
 /**
@@ -65,6 +74,7 @@ export function sfExtensionData(sf: ts.SourceFile): NgExtensionData {
     isTopLevelShim: false,
     fileShim: null,
     originalReferencedFiles: null,
+    taggedReferenceFiles: null,
   };
   extSf[NgExtension] = extension;
   return extension;
@@ -110,3 +120,53 @@ export function copyFileShimData(from: ts.SourceFile, to: ts.SourceFile): void {
   }
   sfExtensionData(to).fileShim = sfExtensionData(from).fileShim;
 }
+
+/**
+ * For those `ts.SourceFile`s in the `program` which have previously been tagged by a
+ * `ShimReferenceTagger`, restore the original `referencedFiles` array that does not have shim tags.
+ */
+export function untagAllTsFiles(program: ts.Program): void {
+  for (const sf of program.getSourceFiles()) {
+    untagTsFile(sf);
+  }
+}
+
+/**
+ * For those `ts.SourceFile`s in the `program` which have previously been tagged by a
+ * `ShimReferenceTagger`, re-apply the effects of tagging by updating the `referencedFiles` array to
+ * the tagged version produced previously.
+ */
+export function retagAllTsFiles(program: ts.Program): void {
+  for (const sf of program.getSourceFiles()) {
+    retagTsFile(sf);
+  }
+}
+
+/**
+ * Restore the original `referencedFiles` for the given `ts.SourceFile`.
+ */
+export function untagTsFile(sf: ts.SourceFile): void {
+  if (sf.isDeclarationFile || !isExtended(sf)) {
+    return;
+  }
+
+  const ext = sfExtensionData(sf);
+  if (ext.originalReferencedFiles !== null) {
+    sf.referencedFiles = ext.originalReferencedFiles as Array<ts.FileReference>;
+  }
+}
+
+/**
+ * Apply the previously tagged `referencedFiles` to the given `ts.SourceFile`, if it was previously
+ * tagged.
+ */
+export function retagTsFile(sf: ts.SourceFile): void {
+  if (sf.isDeclarationFile || !isExtended(sf)) {
+    return;
+  }
+
+  const ext = sfExtensionData(sf);
+  if (ext.taggedReferenceFiles !== null) {
+    sf.referencedFiles = ext.taggedReferenceFiles as Array<ts.FileReference>;
+  }
+}
diff --git a/packages/compiler-cli/src/ngtsc/shims/src/reference_tagger.ts b/packages/compiler-cli/src/ngtsc/shims/src/reference_tagger.ts
@@ -8,10 +8,10 @@
 
 import * as ts from 'typescript';
 
-import {absoluteFrom, absoluteFromSourceFile} from '../../file_system';
+import {absoluteFromSourceFile} from '../../file_system';
 import {isNonDeclarationTsPath} from '../../util/src/typescript';
 
-import {isExtended as isExtendedSf, isShim, NgExtension, sfExtensionData} from './expando';
+import {isShim, sfExtensionData} from './expando';
 import {makeShimFileName} from './util';
 
 /**
@@ -48,8 +48,16 @@ export class ShimReferenceTagger {
       return;
     }
 
-    sfExtensionData(sf).originalReferencedFiles = sf.referencedFiles;
-    const referencedFiles = [...sf.referencedFiles];
+    const ext = sfExtensionData(sf);
+
+    // If this file has never been tagged before, capture its `referencedFiles` in the extension
+    // data.
+    if (ext.originalReferencedFiles === null) {
+      ext.originalReferencedFiles = sf.referencedFiles;
+    }
+
+    const referencedFiles = [...ext.originalReferencedFiles];
+
 
     const sfPath = absoluteFromSourceFile(sf);
     for (const suffix of this.suffixes) {
@@ -60,26 +68,16 @@ export class ShimReferenceTagger {
       });
     }
 
+    ext.taggedReferenceFiles = referencedFiles;
     sf.referencedFiles = referencedFiles;
     this.tagged.add(sf);
   }
 
   /**
-   * Restore the original `referencedFiles` values of all tagged `ts.SourceFile`s and disable the
-   * `ShimReferenceTagger`.
+   * Disable the `ShimReferenceTagger` and free memory associated with tracking tagged files.
    */
   finalize(): void {
     this.enabled = false;
-    for (const sf of this.tagged) {
-      if (!isExtendedSf(sf)) {
-        continue;
-      }
-
-      const extensionData = sfExtensionData(sf);
-      if (extensionData.originalReferencedFiles !== null) {
-        sf.referencedFiles = extensionData.originalReferencedFiles! as ts.FileReference[];
-      }
-    }
     this.tagged.clear();
   }
 }
diff --git a/packages/compiler-cli/src/ngtsc/shims/test/reference_tagger_spec.ts b/packages/compiler-cli/src/ngtsc/shims/test/reference_tagger_spec.ts
@@ -12,6 +12,7 @@ import {absoluteFrom as _, AbsoluteFsPath, getSourceFileOrError} from '../../fil
 import {runInEachFileSystem} from '../../file_system/testing';
 import {makeProgram} from '../../testing';
 import {ShimAdapter} from '../src/adapter';
+import {retagTsFile, untagTsFile} from '../src/expando';
 import {ShimReferenceTagger} from '../src/reference_tagger';
 
 import {TestShimGenerator} from './util';
@@ -67,49 +68,65 @@ runInEachFileSystem(() => {
       expect(shimSf.referencedFiles).toEqual([]);
     });
 
-    it('should remove tags during finalization', () => {
+    it('should not tag shims after finalization', () => {
       const tagger = new ShimReferenceTagger(['test1', 'test2']);
+      tagger.finalize();
 
       const fileName = _('/file.ts');
       const sf = makeArbitrarySf(fileName);
 
-      expectReferencedFiles(sf, []);
-
       tagger.tag(sf);
-      expectReferencedFiles(sf, ['/file.test1.ts', '/file.test2.ts']);
-
-      tagger.finalize();
       expectReferencedFiles(sf, []);
     });
 
-    it('should not remove references it did not add during finalization', () => {
-      const tagger = new ShimReferenceTagger(['test1', 'test2']);
-      const fileName = _('/file.ts');
-      const libFileName = _('/lib.d.ts');
-
-      const sf = makeSf(fileName, `
-        /// <reference path="/lib.d.ts" />
-        export const UNIMPORTANT = true;
-      `);
+    it('should not overwrite original referencedFiles', () => {
+      const tagger = new ShimReferenceTagger(['test']);
 
-      expectReferencedFiles(sf, [libFileName]);
+      const fileName = _('/file.ts');
+      const sf = makeArbitrarySf(fileName);
+      sf.referencedFiles = [{
+        fileName: _('/other.ts'),
+        pos: 0,
+        end: 0,
+      }];
 
       tagger.tag(sf);
-      expectReferencedFiles(sf, ['/file.test1.ts', '/file.test2.ts', libFileName]);
-
-      tagger.finalize();
-      expectReferencedFiles(sf, [libFileName]);
+      expectReferencedFiles(sf, ['/other.ts', '/file.test.ts']);
     });
 
-    it('should not tag shims after finalization', () => {
-      const tagger = new ShimReferenceTagger(['test1', 'test2']);
-      tagger.finalize();
+    it('should always tag against the original referencedFiles', () => {
+      const tagger1 = new ShimReferenceTagger(['test1']);
+      const tagger2 = new ShimReferenceTagger(['test2']);
 
       const fileName = _('/file.ts');
       const sf = makeArbitrarySf(fileName);
 
-      tagger.tag(sf);
-      expectReferencedFiles(sf, []);
+      tagger1.tag(sf);
+      tagger2.tag(sf);
+      expectReferencedFiles(sf, ['/file.test2.ts']);
+    });
+
+    describe('tagging and untagging', () => {
+      it('should be able to untag references and retag them later', () => {
+        const tagger = new ShimReferenceTagger(['test']);
+
+        const fileName = _('/file.ts');
+        const sf = makeArbitrarySf(fileName);
+        sf.referencedFiles = [{
+          fileName: _('/other.ts'),
+          pos: 0,
+          end: 0,
+        }];
+
+        tagger.tag(sf);
+        expectReferencedFiles(sf, ['/other.ts', '/file.test.ts']);
+
+        untagTsFile(sf);
+        expectReferencedFiles(sf, ['/other.ts']);
+
+        retagTsFile(sf);
+        expectReferencedFiles(sf, ['/other.ts', '/file.test.ts']);
+      });
     });
   });
 });

diff --git a/packages/compiler-cli/src/ngtsc/testing/index.ts b/packages/compiler-cli/src/ngtsc/testing/index.ts
@@ -5,4 +5,4 @@
  * 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
  */
-export {getDeclaration, makeProgram} from './src/utils';
+export {expectCompleteReuse, getDeclaration, makeProgram} from './src/utils';
diff --git a/packages/compiler-cli/src/ngtsc/testing/src/utils.ts b/packages/compiler-cli/src/ngtsc/testing/src/utils.ts
@@ -97,6 +97,23 @@ export function walkForDeclaration(name: string, rootNode: ts.Node): ts.Declarat
   return chosenDecl;
 }
 
+const COMPLETE_REUSE_FAILURE_MESSAGE =
+    'The original program was not reused completely, even though no changes should have been made to its structure';
+
+/**
+ * Extracted from TypeScript's internal enum `StructureIsReused`.
+ */
+enum TsStructureIsReused {
+  Not = 0,
+  SafeModules = 1,
+  Completely = 2,
+}
+
+export function expectCompleteReuse(oldProgram: ts.Program): void {
+  // Assert complete reuse using TypeScript's private API.
+  expect((oldProgram as any).structureIsReused)
+      .toBe(TsStructureIsReused.Completely, COMPLETE_REUSE_FAILURE_MESSAGE);
+}
 
 function bindingNameEquals(node: ts.BindingName, name: string): boolean {
   if (ts.isIdentifier(node)) {

diff --git a/packages/compiler-cli/src/ngtsc/tsc_plugin.ts b/packages/compiler-cli/src/ngtsc/tsc_plugin.ts
@@ -13,6 +13,7 @@ import {NgCompilerOptions, UnifiedModulesHost} from './core/api';
 import {NodeJSFileSystem, setFileSystem} from './file_system';
 import {PatchedProgramIncrementalBuildStrategy} from './incremental';
 import {NOOP_PERF_RECORDER} from './perf';
+import {untagAllTsFiles} from './shims';
 import {ReusedProgramStrategy} from './typecheck/src/augmented_program';
 
 // The following is needed to fix a the chicken-and-egg issue where the sync (into g3) script will
@@ -80,6 +81,9 @@ export class NgTscPlugin implements TscPlugin {
   wrapHost(
       host: ts.CompilerHost&UnifiedModulesHost, inputFiles: readonly string[],
       options: ts.CompilerOptions): PluginCompilerHost {
+    // TODO(alxhub): Eventually the `wrapHost()` API will accept the old `ts.Program` (if one is
+    // available). When it does, its `ts.SourceFile`s need to be re-tagged to enable proper
+    // incremental compilation.
     this.options = {...this.ngOptions, ...options} as NgCompilerOptions;
     this.host = NgCompilerHost.wrap(host, inputFiles, this.options, /* oldProgram */ null);
     return this.host;
@@ -92,6 +96,8 @@ export class NgTscPlugin implements TscPlugin {
     if (this.host === null || this.options === null) {
       throw new Error('Lifecycle error: setupCompilation() before wrapHost().');
     }
+    this.host.postProgramCreationCleanup();
+    untagAllTsFiles(program);
     const typeCheckStrategy = new ReusedProgramStrategy(
         program, this.host, this.options, this.host.shimExtensionPrefixes);
     this._compiler = new NgCompiler(