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

crisbeto · pkozlowski-opensource · commit 26cb7ab2e6ac · 2023-02-09T15:37:18.000+01:00
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
diff --git a/packages/core/schematics/ng-generate/standalone-migration/index.ts b/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);
   }
 
diff --git a/packages/core/schematics/ng-generate/standalone-migration/prune-modules.ts b/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.
diff --git a/packages/core/schematics/ng-generate/standalone-migration/standalone-bootstrap.ts b/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]);
-
-  if (highlights) {
-    const offsets: [start: number, end: number][] = [];
-
-    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]);
-          }
-        }
-      }
-    }
+    referenceResolver: ReferenceResolver): Set<ts.Node>|null {
+  const offsets =
+      referenceResolver.findSameFileReferences(node, node.getSourceFile().fileName)
+          .filter(([start, end]) => isOutsideRange(excludeStart, excludeEnd, start, end));
 
-    if (offsets.length > 0) {
-      const nodes = offsetsToNodes(nodeLookup, offsets, new Set());
+  if (offsets.length > 0) {
+    const nodes = offsetsToNodes(nodeLookup, offsets, new Set());
 
-      if (nodes.size > 0) {
-        return nodes;
-      }
+    if (nodes.size > 0) {
+      return nodes;
     }
   }
 
diff --git a/packages/core/schematics/ng-generate/standalone-migration/util.ts b/packages/core/schematics/ng-generate/standalone-migration/util.ts
