perf(migrations): speed up language service lookups (#49010)

For the module pruning and bootstrap API migration steps we depend heavily upon the TypeScript `LanguageService` which ends up being slow on a large project. E.g. in some large internal projects single-file lookups were taking around 30s.

These changes introduce a wrapper around the `LanguageService` that we can use to trick it into not traversing the entire project every time.

PR Close #49010

packages/core/schematics/ng-generate/standalone-migration/index.ts

@@ -74,6 +74,7 @@ function standaloneMigration(
         skipLibCheck: true,
         skipDefaultLibCheck: true,
       });
+  const referenceLookupExcludedFiles = /node_modules|\.ngtypecheck\.ts/;
   const program = createProgram({rootNames, host, options}) as NgtscProgram;
   const printer = ts.createPrinter();
 
@@ -82,10 +83,9 @@ function standaloneMigration(
         pathToMigrate} has to be a directory. Cannot run the standalone migration.`);
   }
 
-  const sourceFiles = program.getTsProgram().getSourceFiles().filter(sourceFile => {
-    return sourceFile.fileName.startsWith(pathToMigrate) &&
-        canMigrateFile(basePath, sourceFile, program.getTsProgram());
-  });
+  const sourceFiles = program.getTsProgram().getSourceFiles().filter(
+      sourceFile => sourceFile.fileName.startsWith(pathToMigrate) &&
+          canMigrateFile(basePath, sourceFile, program.getTsProgram()));
 
   if (sourceFiles.length === 0) {
     return 0;
@@ -95,14 +95,17 @@ function standaloneMigration(
   let filesToRemove: Set<ts.SourceFile>|null = null;
 
   if (schematicOptions.mode === MigrationMode.pruneModules) {
-    const result = pruneNgModules(program, host, basePath, rootNames, sourceFiles, printer);
+    const result = pruneNgModules(
+        program, host, basePath, rootNames, sourceFiles, printer, undefined,
+        referenceLookupExcludedFiles);
     pendingChanges = result.pendingChanges;
     filesToRemove = result.filesToRemove;
   } else if (schematicOptions.mode === MigrationMode.standaloneBootstrap) {
-    pendingChanges =
-        toStandaloneBootstrap(program, host, basePath, rootNames, sourceFiles, printer);
+    pendingChanges = toStandaloneBootstrap(
+        program, host, basePath, rootNames, sourceFiles, printer, undefined,
+        referenceLookupExcludedFiles);
   } else {
-    /** MigrationMode.toStandalone */
+    // This shouldn't happen, but default to `MigrationMode.toStandalone` just in case.
     pendingChanges = toStandalone(sourceFiles, program, printer);
   }
 

packages/core/schematics/ng-generate/standalone-migration/prune-modules.ts

@@ -12,10 +12,7 @@ import ts from 'typescript';
 import {getAngularDecorators, NgDecorator} from '../../utils/ng_decorators';
 import {closestNode} from '../../utils/typescript/nodes';
 
-import {ChangeTracker, createLanguageService, findClassDeclaration, findLiteralProperty, getNodeLookup, ImportRemapper, offsetsToNodes, UniqueItemTracker} from './util';
-
-/** Mapping between a file name and spans for node references inside of it. */
-type ReferencesByFile = Map<string, [start: number, end: number][]>;
+import {ChangeTracker, findClassDeclaration, findLiteralProperty, getNodeLookup, ImportRemapper, offsetsToNodes, ReferenceResolver, UniqueItemTracker} from './util';
 
 /** Keeps track of the places from which we need to remove AST nodes. */
 interface RemovalLocations {
@@ -28,11 +25,13 @@ interface RemovalLocations {
 
 export function pruneNgModules(
     program: NgtscProgram, host: ts.CompilerHost, basePath: string, rootFileNames: string[],
-    sourceFiles: ts.SourceFile[], printer: ts.Printer, importRemapper?: ImportRemapper) {
+    sourceFiles: ts.SourceFile[], printer: ts.Printer, importRemapper?: ImportRemapper,
+    referenceLookupExcludedFiles?: RegExp) {
   const filesToRemove = new Set<ts.SourceFile>();
   const tracker = new ChangeTracker(printer, importRemapper);
   const typeChecker = program.getTsProgram().getTypeChecker();
-  const languageService = createLanguageService(program, host, rootFileNames, basePath);
+  const referenceResolver =
+      new ReferenceResolver(program, host, rootFileNames, basePath, referenceLookupExcludedFiles);
   const removalLocations: RemovalLocations = {
     arrays: new UniqueItemTracker<ts.ArrayLiteralExpression, ts.Node>(),
     imports: new UniqueItemTracker<ts.NamedImports, ts.Node>(),
@@ -43,7 +42,7 @@ export function pruneNgModules(
 
   sourceFiles.forEach(function walk(node: ts.Node) {
     if (ts.isClassDeclaration(node) && canRemoveClass(node, typeChecker)) {
-      collectRemovalLocations(node, removalLocations, languageService, program);
+      collectRemovalLocations(node, removalLocations, referenceResolver, program);
       removalLocations.classes.add(node);
     }
     node.forEachChild(walk);
@@ -73,13 +72,13 @@ export function pruneNgModules(
  * Collects all the nodes that a module needs to be removed from.
  * @param ngModule Module being removed.
  * @param removalLocations
- * @param languageService
+ * @param referenceResolver
  * @param program
  */
 function collectRemovalLocations(
     ngModule: ts.ClassDeclaration, removalLocations: RemovalLocations,
-    languageService: ts.LanguageService, program: NgtscProgram) {
-  const refsByFile = extractReferences(ngModule, languageService);
+    referenceResolver: ReferenceResolver, program: NgtscProgram) {
+  const refsByFile = referenceResolver.findReferencesInProject(ngModule.name!);
   const tsProgram = program.getTsProgram();
   const nodes = new Set<ts.Node>();
 
@@ -295,34 +294,6 @@ function canRemoveFile(sourceFile: ts.SourceFile, classesToBeRemoved: Set<ts.Cla
   return true;
 }
 
-
-/**
- * Finds all the locations in a file where a node is referenced.
- * @param node Node that is being looked up.
- * @param languageService Language service used to find the references.
- */
-function extractReferences(
-    node: ts.ClassDeclaration, languageService: ts.LanguageService): ReferencesByFile {
-  const result: ReferencesByFile = new Map();
-  const referencedSymbols =
-      languageService.findReferences(node.getSourceFile().fileName, node.name!.getStart()) || [];
-
-  for (const symbol of referencedSymbols) {
-    for (const ref of symbol.references) {
-      if (!ref.isDefinition || symbol.definition.kind === ts.ScriptElementKind.alias) {
-        if (!result.has(ref.fileName)) {
-          result.set(ref.fileName, []);
-        }
-
-        result.get(ref.fileName)!.push(
-            [ref.textSpan.start, ref.textSpan.start + ref.textSpan.length]);
-      }
-    }
-  }
-
-  return result;
-}
-
 /**
  * Gets whether an AST node contains another AST node.
  * @param parent Parent node that may contain the child.

packages/core/schematics/ng-generate/standalone-migration/standalone-bootstrap.ts

@@ -16,7 +16,7 @@ import {getAngularDecorators} from '../../utils/ng_decorators';
 import {closestNode} from '../../utils/typescript/nodes';
 
 import {convertNgModuleDeclarationToStandalone, extractDeclarationsFromModule, findTestObjectsToMigrate, migrateTestDeclarations} from './to-standalone';
-import {ChangeTracker, createLanguageService, findClassDeclaration, findLiteralProperty, getNodeLookup, getRelativeImportPath, ImportRemapper, NamedClassDeclaration, NodeLookup, offsetsToNodes, UniqueItemTracker} from './util';
+import {ChangeTracker, findClassDeclaration, findLiteralProperty, getNodeLookup, getRelativeImportPath, ImportRemapper, NamedClassDeclaration, NodeLookup, offsetsToNodes, ReferenceResolver, UniqueItemTracker} from './util';
 
 /** Information extracted from a `bootstrapModule` call necessary to migrate it. */
 interface BootstrapCallAnalysis {
@@ -34,11 +34,13 @@ interface BootstrapCallAnalysis {
 
 export function toStandaloneBootstrap(
     program: NgtscProgram, host: ts.CompilerHost, basePath: string, rootFileNames: string[],
-    sourceFiles: ts.SourceFile[], printer: ts.Printer, importRemapper?: ImportRemapper) {
+    sourceFiles: ts.SourceFile[], printer: ts.Printer, importRemapper?: ImportRemapper,
+    referenceLookupExcludedFiles?: RegExp) {
   const tracker = new ChangeTracker(printer, importRemapper);
   const typeChecker = program.getTsProgram().getTypeChecker();
   const templateTypeChecker = program.compiler.getTemplateTypeChecker();
-  const languageService = createLanguageService(program, host, rootFileNames, basePath);
+  const referenceResolver =
+      new ReferenceResolver(program, host, rootFileNames, basePath, referenceLookupExcludedFiles);
   const bootstrapCalls: BootstrapCallAnalysis[] = [];
   const testObjects: ts.ObjectLiteralExpression[] = [];
   const allDeclarations: Reference<ts.ClassDeclaration>[] = [];
@@ -62,7 +64,7 @@ export function toStandaloneBootstrap(
 
   for (const call of bootstrapCalls) {
     allDeclarations.push(...call.declarations);
-    migrateBootstrapCall(call, tracker, languageService, typeChecker, printer);
+    migrateBootstrapCall(call, tracker, referenceResolver, typeChecker, printer);
   }
 
   // The previous migrations explicitly skip over bootstrapped
@@ -133,12 +135,12 @@ function analyzeBootstrapCall(
  * Converts a `bootstrapModule` call to `bootstrapApplication`.
  * @param analysis Analysis result of the call.
  * @param tracker Tracker in which to register the changes.
- * @param languageService
+ * @param referenceResolver
  * @param typeChecker
  * @param printer
  */
 function migrateBootstrapCall(
-    analysis: BootstrapCallAnalysis, tracker: ChangeTracker, languageService: ts.LanguageService,
+    analysis: BootstrapCallAnalysis, tracker: ChangeTracker, referenceResolver: ReferenceResolver,
     typeChecker: ts.TypeChecker, printer: ts.Printer) {
   const sourceFile = analysis.call.getSourceFile();
   const moduleSourceFile = analysis.metadata.getSourceFile();
@@ -168,14 +170,14 @@ function migrateBootstrapCall(
       providersInNewCall.push(ts.factory.createSpreadElement(providers.initializer));
     }
 
-    addNodesToCopy(sourceFile, providers, nodeLookup, tracker, nodesToCopy, languageService);
+    addNodesToCopy(sourceFile, providers, nodeLookup, tracker, nodesToCopy, referenceResolver);
   }
 
   if (imports && ts.isPropertyAssignment(imports)) {
     nodeLookup = nodeLookup || getNodeLookup(moduleSourceFile);
     migrateImportsForBootstrapCall(
         sourceFile, imports, nodeLookup, moduleImportsInNewCall, providersInNewCall, tracker,
-        nodesToCopy, languageService, typeChecker);
+        nodesToCopy, referenceResolver, typeChecker);
   }
 
   if (nodesToCopy.size > 0) {
@@ -253,13 +255,13 @@ function replaceBootstrapCallExpression(
  * @param providersInNewCall Array keeping track of the providers in the new call.
  * @param tracker Tracker in which changes to files are being stored.
  * @param nodesToCopy Nodes that should be copied to the new file.
- * @param languageService
+ * @param referenceResolver
  * @param typeChecker
  */
 function migrateImportsForBootstrapCall(
     sourceFile: ts.SourceFile, imports: ts.PropertyAssignment, nodeLookup: NodeLookup,
     importsForNewCall: ts.Expression[], providersInNewCall: ts.Expression[], tracker: ChangeTracker,
-    nodesToCopy: Set<ts.Node>, languageService: ts.LanguageService,
+    nodesToCopy: Set<ts.Node>, referenceResolver: ReferenceResolver,
     typeChecker: ts.TypeChecker): void {
   if (!ts.isArrayLiteralExpression(imports.initializer)) {
     importsForNewCall.push(imports.initializer);
@@ -282,9 +284,9 @@ function migrateImportsForBootstrapCall(
             tracker.addImport(sourceFile, 'provideRouter', '@angular/router'), [],
             [element.arguments[0], ...features]));
         addNodesToCopy(
-            sourceFile, element.arguments[0], nodeLookup, tracker, nodesToCopy, languageService);
+            sourceFile, element.arguments[0], nodeLookup, tracker, nodesToCopy, referenceResolver);
         if (options) {
-          addNodesToCopy(sourceFile, options, nodeLookup, tracker, nodesToCopy, languageService);
+          addNodesToCopy(sourceFile, options, nodeLookup, tracker, nodesToCopy, referenceResolver);
         }
         continue;
       }
@@ -326,7 +328,7 @@ function migrateImportsForBootstrapCall(
         decorators.every(
             ({name}) => name !== 'Directive' && name !== 'Component' && name !== 'Pipe')) {
       importsForNewCall.push(element);
-      addNodesToCopy(sourceFile, element, nodeLookup, tracker, nodesToCopy, languageService);
+      addNodesToCopy(sourceFile, element, nodeLookup, tracker, nodesToCopy, referenceResolver);
     }
   }
 }
@@ -444,12 +446,12 @@ function getRouterModuleForRootFeatures(
  * @param nodeLookup Map used to look up nodes based on their positions in a file.
  * @param tracker Tracker in which changes to files are stored.
  * @param nodesToCopy Set that keeps track of the nodes being copied.
- * @param languageService
+ * @param referenceResolver
  */
 function addNodesToCopy(
     targetFile: ts.SourceFile, rootNode: ts.Node, nodeLookup: NodeLookup, tracker: ChangeTracker,
-    nodesToCopy: Set<ts.Node>, languageService: ts.LanguageService): void {
-  const refs = findAllSameFileReferences(rootNode, nodeLookup, languageService);
+    nodesToCopy: Set<ts.Node>, referenceResolver: ReferenceResolver): void {
+  const refs = findAllSameFileReferences(rootNode, nodeLookup, referenceResolver);
 
   for (const ref of refs) {
     const importSpecifier = closestOrSelf(ref, ts.isImportSpecifier);
@@ -504,10 +506,10 @@ function addNodesToCopy(
  * Finds all the nodes referenced within the root node in the same file.
  * @param rootNode Node from which to start looking for references.
  * @param nodeLookup Map used to look up nodes based on their positions in a file.
- * @param languageService
+ * @param referenceResolver
  */
 function findAllSameFileReferences(
-    rootNode: ts.Node, nodeLookup: NodeLookup, languageService: ts.LanguageService): Set<ts.Node> {
+    rootNode: ts.Node, nodeLookup: NodeLookup, referenceResolver: ReferenceResolver): Set<ts.Node> {
   const results = new Set<ts.Node>();
   const excludeStart = rootNode.getStart();
   const excludeEnd = rootNode.getEnd();
@@ -518,8 +520,8 @@ function findAllSameFileReferences(
       return;
     }
 
-    const refs =
-        referencesToNodeWithinSameFile(node, nodeLookup, excludeStart, excludeEnd, languageService);
+    const refs = referencesToNodeWithinSameFile(
+        node, nodeLookup, excludeStart, excludeEnd, referenceResolver);
 
     if (refs === null) {
       return;
@@ -551,38 +553,20 @@ function findAllSameFileReferences(
  * @param nodeLookup Map used to look up nodes based on their positions in a file.
  * @param excludeStart Start of a range that should be excluded from the results.
  * @param excludeEnd End of a range that should be excluded from the results.
- * @param languageService
+ * @param referenceResolver
  */
 function referencesToNodeWithinSameFile(
     node: ts.Identifier, nodeLookup: NodeLookup, excludeStart: number, excludeEnd: number,
-    languageService: ts.LanguageService): Set<ts.Node>|null {
-  const sourceFile = node.getSourceFile();
-  const fileName = sourceFile.fileName;
-  const highlights = languageService.getDocumentHighlights(fileName, node.getStart(), [fileName]);
+    referenceResolver: ReferenceResolver): Set<ts.Node>|null {
+  const offsets =
+      referenceResolver.findSameFileReferences(node, node.getSourceFile().fileName)
+          .filter(([start, end]) => isOutsideRange(excludeStart, excludeEnd, start, end));
 
-  if (highlights) {
-    const offsets: [start: number, end: number][] = [];
+  if (offsets.length > 0) {
+    const nodes = offsetsToNodes(nodeLookup, offsets, new Set());
 
-    for (const file of highlights) {
-      // We are pretty much guaranteed to only have one match from the current file since it is
-      // the only one being passed in `getDocumentHighlight`, but we check here just in case.
-      if (file.fileName === fileName) {
-        for (const {textSpan: {start, length}, kind} of file.highlightSpans) {
-          const end = start + length;
-          if (kind !== ts.HighlightSpanKind.none &&
-              isOutsideRange(excludeStart, excludeEnd, start, end)) {
-            offsets.push([start, end]);
-          }
-        }
-      }
-    }
-
-    if (offsets.length > 0) {
-      const nodes = offsetsToNodes(nodeLookup, offsets, new Set());
-
-      if (nodes.size > 0) {
-        return nodes;
-      }
+    if (nodes.size > 0) {
+      return nodes;
     }
   }
 

packages/core/schematics/ng-generate/standalone-migration/util.ts

@@ -24,6 +24,12 @@ export type NamedClassDeclaration = ts.ClassDeclaration&{name: ts.Identifier};
 /** Function that can be used to remap a generated import. */
 export type ImportRemapper = (moduleName: string) => string;
 
+/** Text span of an AST node. */
+export type ReferenceSpan = [start: number, end: number];
+
+/** Mapping between a file name and spans for node references inside of it. */
+export type ReferencesByFile = Map<string, ReferenceSpan[]>;
+
 /** Change that needs to be applied to a file. */
 interface PendingChange {
   /** Index at which to start changing the file. */
@@ -172,30 +178,114 @@ export class UniqueItemTracker<K, V> {
   }
 }
 
+/** Resolves references to nodes. */
+export class ReferenceResolver {
+  private _languageService: ts.LanguageService|undefined;
 
-/**
- * Creates a TypeScript language service.
- * @param program Program used to analyze the project.
- * @param host Compiler host used to interact with the file system.
- * @param rootFileNames Root files of the project.
- * @param basePath Root path of the project.
- */
-export function createLanguageService(
-    program: NgtscProgram, host: ts.CompilerHost, rootFileNames: string[],
-    basePath: string): ts.LanguageService {
-  return ts.createLanguageService({
-    getCompilationSettings: () => program.getTsProgram().getCompilerOptions(),
-    getScriptFileNames: () => rootFileNames,
-    getScriptVersion: () => '0',  // The files won't change so we can return the same version.
-    getScriptSnapshot: (fileName: string) => {
-      const content = host.readFile(fileName);
-      return content ? ts.ScriptSnapshot.fromString(content) : undefined;
-    },
-    getCurrentDirectory: () => basePath,
-    getDefaultLibFileName: options => ts.getDefaultLibFilePath(options),
-    readFile: (path: string) => host.readFile(path),
-    fileExists: (path: string) => host.fileExists(path)
-  });
+  /**
+   * If set, allows the language service to *only* read a specific file.
+   * Used to speed up single-file lookups.
+   */
+  private _tempOnlyFile: string|null = null;
+
+  constructor(
+      private _program: NgtscProgram, private _host: ts.CompilerHost,
+      private _rootFileNames: string[], private _basePath: string,
+      private _excludedFiles?: RegExp) {}
+
+  /** Finds all references to a node within the entire project. */
+  findReferencesInProject(node: ts.Node): ReferencesByFile {
+    const referencedSymbols =
+        this._getLanguageService().findReferences(node.getSourceFile().fileName, node.getStart()) ||
+        [];
+    const results: ReferencesByFile = new Map();
+
+    for (const symbol of referencedSymbols) {
+      for (const ref of symbol.references) {
+        if (!ref.isDefinition || symbol.definition.kind === ts.ScriptElementKind.alias) {
+          if (!results.has(ref.fileName)) {
+            results.set(ref.fileName, []);
+          }
+
+          results.get(ref.fileName)!.push(
+              [ref.textSpan.start, ref.textSpan.start + ref.textSpan.length]);
+        }
+      }
+    }
+
+    return results;
+  }
+
+  /** Finds all references to a node within a single file. */
+  findSameFileReferences(node: ts.Node, fileName: string): ReferenceSpan[] {
+    // Even though we're only passing in a single file into `getDocumentHighlights`, the language
+    // service ends up traversing the entire project. Prevent it from reading any files aside from
+    // the one we're interested in by intercepting it at the compiler host level.
+    // This is an order of magnitude faster on a large project.
+    this._tempOnlyFile = fileName;
+    const results: ReferenceSpan[] = [];
+    const highlights =
+        this._getLanguageService().getDocumentHighlights(fileName, node.getStart(), [fileName]);
+
+    if (highlights) {
+      for (const file of highlights) {
+        // We are pretty much guaranteed to only have one match from the current file since it is
+        // the only one being passed in `getDocumentHighlight`, but we check here just in case.
+        if (file.fileName === fileName) {
+          for (const {textSpan: {start, length}, kind} of file.highlightSpans) {
+            if (kind !== ts.HighlightSpanKind.none) {
+              results.push([start, start + length]);
+            }
+          }
+        }
+      }
+    }
+
+    // Restore full project access to the language service.
+    this._tempOnlyFile = null;
+    return results;
+  }
+
+  /** Used by the language service  */
+  private _readFile(path: string) {
+    if ((this._tempOnlyFile !== null && path !== this._tempOnlyFile) ||
+        this._excludedFiles?.test(path)) {
+      return '';
+    }
+    return this._host.readFile(path);
+  }
+
+  /** Gets a language service that can be used to perform lookups. */
+  private _getLanguageService(): ts.LanguageService {
+    if (!this._languageService) {
+      const rootFileNames = this._rootFileNames.slice();
+
+      this._program.getTsProgram().getSourceFiles().forEach(({fileName}) => {
+        if (!this._excludedFiles?.test(fileName) && !rootFileNames.includes(fileName)) {
+          rootFileNames.push(fileName);
+        }
+      });
+
+      this._languageService = ts.createLanguageService(
+          {
+            getCompilationSettings: () => this._program.getTsProgram().getCompilerOptions(),
+            getScriptFileNames: () => rootFileNames,
+            // The files won't change so we can return the same version.
+            getScriptVersion: () => '0',
+            getScriptSnapshot: (path: string) => {
+              const content = this._readFile(path);
+              return content ? ts.ScriptSnapshot.fromString(content) : undefined;
+            },
+            getCurrentDirectory: () => this._basePath,
+            getDefaultLibFileName: options => ts.getDefaultLibFilePath(options),
+            readFile: path => this._readFile(path),
+            fileExists: (path: string) => this._host.fileExists(path)
+          },
+          ts.createDocumentRegistry(), ts.LanguageServiceMode.PartialSemantic);
+    }
+
+    return this._languageService;
+  }
 }
 
 /** Creates a NodeLookup object from a source file. */
@@ -224,8 +314,7 @@ export function getNodeLookup(sourceFile: ts.SourceFile): NodeLookup {
  * @param results Set in which to store the results.
  */
 export function offsetsToNodes(
-    lookup: NodeLookup, offsets: [start: number, end: number][],
-    results: Set<ts.Node>): Set<ts.Node> {
+    lookup: NodeLookup, offsets: ReferenceSpan[], results: Set<ts.Node>): Set<ts.Node> {
   for (const [start, end] of offsets) {
     const match = lookup.get(start)?.find(node => node.getEnd() === end);