perf(migrations): speed up signal input migration by combining two analyze phases (#57318)

devversion · AndrewKushnir · commit 604270619a21 · 2024-08-12T12:12:23.000-07:00
Instead of revisiting each source file, and each of its child nodes twice, we now visit them together using a grouped AST visitor that only traverses each source file once. This seemed to speed up migration by 6-8% locally, but is likely noticable better with large compilation scopes. PR Close #57318
diff --git a/packages/core/schematics/migrations/signal-migration/src/passes/2_find_source_file_references.ts b/packages/core/schematics/migrations/signal-migration/src/passes/2_find_source_file_references.ts
@@ -19,6 +19,7 @@ import {identifyTemplateReferences} from './references/identify_template_referen
 import {identifyPotentialTypeScriptReference} from './references/identify_ts_references';
 import {PartialDirectiveTypeInCatalystTests} from '../pattern_advisors/partial_directive_type';
 import {InputReferenceKind} from '../utils/input_reference';
+import {GroupedTsAstVisitor} from '../utils/grouped_ts_ast_visitor';
 
 /**
  * Phase where we iterate through all source file references and
@@ -34,11 +35,11 @@ import {InputReferenceKind} from '../utils/input_reference';
  *    - Host binding expressions.
  */
 export function pass2_IdentifySourceFileReferences(
-  sf: ts.SourceFile,
   host: MigrationHost,
   checker: ts.TypeChecker,
   reflector: ReflectionHost,
   templateTypeChecker: TemplateTypeChecker,
+  groupedTsAstVisitor: GroupedTsAstVisitor,
   knownInputs: KnownInputs,
   result: MigrationResult,
 ) {
@@ -86,8 +87,7 @@ export function pass2_IdentifySourceFileReferences(
         target: partialDirectiveInCatalyst.targetClass,
       });
     }
-
-    ts.forEachChild(node, visitor);
   };
-  ts.forEachChild(sf, visitor);
+
+  groupedTsAstVisitor.register(visitor);
 }
diff --git a/packages/core/schematics/migrations/signal-migration/src/passes/3_check_incompatible_patterns.ts b/packages/core/schematics/migrations/signal-migration/src/passes/3_check_incompatible_patterns.ts
@@ -15,6 +15,7 @@ import {getMemberName} from '../utils/class_member_names';
 import {InheritanceGraph} from '../utils/inheritance_graph';
 import {SpyOnInputPattern} from '../pattern_advisors/spy_on_pattern';
 import {MigrationHost} from '../migration_host';
+import {GroupedTsAstVisitor} from '../utils/grouped_ts_ast_visitor';
 
 /**
  * Phase where problematic patterns are detected and advise
@@ -28,9 +29,9 @@ import {MigrationHost} from '../migration_host';
  */
 export function pass3__checkIncompatiblePatterns(
   host: MigrationHost,
-  files: readonly ts.SourceFile[],
   inheritanceGraph: InheritanceGraph,
   checker: ts.TypeChecker,
+  groupedTsAstVisitor: GroupedTsAstVisitor,
   knownInputs: KnownInputs,
 ) {
   const inputClassSymbolsToClass = new Map<ts.Symbol, ts.ClassDeclaration>();
@@ -51,7 +52,6 @@ export function pass3__checkIncompatiblePatterns(
 
   const incompatibilityPatterns = [new SpyOnInputPattern(host, checker, knownInputs)];
 
-  let insidePropertyDeclaration: ts.PropertyDeclaration | null = null;
   const visitor = (node: ts.Node) => {
     // Check for manual class instantiations.
     if (ts.isNewExpression(node) && ts.isIdentifier(unwrapExpression(node.expression))) {
@@ -71,6 +71,7 @@ export function pass3__checkIncompatiblePatterns(
     // Detect possible problematic patterns.
     incompatibilityPatterns.forEach((p) => p.detect(node));
 
+    const insidePropertyDeclaration = groupedTsAstVisitor.state.insidePropertyDeclaration;
     // Check for problematic class references inside property declarations.
     // These are likely problematic, causing type conflicts, if the containing
     // class inherits a non-input member with the same name.
@@ -109,16 +110,7 @@ export function pass3__checkIncompatiblePatterns(
         );
       }
     }
-
-    if (ts.isPropertyDeclaration(node)) {
-      const oldPropertyDeclaration = insidePropertyDeclaration;
-      insidePropertyDeclaration = node;
-      ts.forEachChild(node, visitor);
-      insidePropertyDeclaration = oldPropertyDeclaration;
-    } else {
-      ts.forEachChild(node, visitor);
-    }
   };
 
-  files.forEach((f) => ts.forEachChild(f, visitor));
+  groupedTsAstVisitor.register(visitor);
 }
diff --git a/packages/core/schematics/migrations/signal-migration/src/phase_analysis.ts b/packages/core/schematics/migrations/signal-migration/src/phase_analysis.ts
@@ -15,6 +15,7 @@ import {pass3__checkIncompatiblePatterns} from './passes/3_check_incompatible_pa
 import {pass2_IdentifySourceFileReferences} from './passes/2_find_source_file_references';
 import {MigrationResult} from './result';
 import {InheritanceGraph} from './utils/inheritance_graph';
+import {GroupedTsAstVisitor} from './utils/grouped_ts_ast_visitor';
 
 /**
  * Executes the analysis phase of the migration.
@@ -59,26 +60,32 @@ export function executeAnalysisPhase(
       ),
   );
 
-  // Pass 2
-  sourceFiles.forEach((sf) =>
-    pass2_IdentifySourceFileReferences(
-      sf,
-      host,
-      typeChecker,
-      reflector,
-      templateTypeChecker,
-      knownInputs,
-      result,
-    ),
-  );
-
-  // Source files is fine. We will resolve into declaration files if a source
-  // file depends on such.
+  // A graph starting with source files is sufficient. We will resolve into
+  // declaration files if a source file depends on such.
   const inheritanceGraph = new InheritanceGraph(typeChecker).expensivePopulate(sourceFiles);
+  const pass2And3SourceFileVisitor = new GroupedTsAstVisitor(sourceFiles);
+
+  // Register pass 2. Find all source file references.
+  pass2_IdentifySourceFileReferences(
+    host,
+    typeChecker,
+    reflector,
+    templateTypeChecker,
+    pass2And3SourceFileVisitor,
+    knownInputs,
+    result,
+  );
+  // Register pass 3. Check incompatible patterns pass.
+  pass3__checkIncompatiblePatterns(
+    host,
+    inheritanceGraph,
+    typeChecker,
+    pass2And3SourceFileVisitor,
+    knownInputs,
+  );
 
-  // Check pass
-  // TODO: Combine with source file iterations above for speed up??
-  pass3__checkIncompatiblePatterns(host, sourceFiles, inheritanceGraph, typeChecker, knownInputs);
+  // Perform Pass 2 and Pass 3, efficiently in one pass.
+  pass2And3SourceFileVisitor.execute();
 
   return {inheritanceGraph};
 }
diff --git a/packages/core/schematics/migrations/signal-migration/src/utils/grouped_ts_ast_visitor.ts b/packages/core/schematics/migrations/signal-migration/src/utils/grouped_ts_ast_visitor.ts
@@ -0,0 +1,49 @@
+/**
+ * @license
+ * Copyright Google LLC 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 ts from 'typescript';
+
+/**
+ * Class that allows for efficient grouping of TypeScript node AST
+ * traversal.
+ *
+ * Allows visitors to execute in a single pass when visiting all
+ * children of source files.
+ */
+export class GroupedTsAstVisitor {
+  private visitors: Array<(node: ts.Node) => void> = [];
+
+  constructor(private files: readonly ts.SourceFile[]) {}
+
+  state = {
+    insidePropertyDeclaration: null as ts.PropertyDeclaration | null,
+  };
+
+  register(visitor: (node: ts.Node) => void) {
+    this.visitors.push(visitor);
+  }
+
+  execute() {
+    const visitor = (node: ts.Node) => {
+      for (const v of this.visitors) {
+        v(node);
+      }
+      if (ts.isPropertyDeclaration(node)) {
+        this.state.insidePropertyDeclaration = node;
+        ts.forEachChild(node, visitor);
+        this.state.insidePropertyDeclaration = null;
+      } else {
+        ts.forEachChild(node, visitor);
+      }
+    };
+
+    for (const file of this.files) {
+      ts.forEachChild(file, visitor);
+    }
+  }
+}
