From 728bfbbb3192ac72a299fdecd41df19400b697ed Mon Sep 17 00:00:00 2001
From: Paul Gschwendtner <paulgschwendtner@gmail.com>
Date: Wed, 6 Nov 2024 12:58:31 +0000
Subject: [PATCH] refactor(migrations): speed up code refactoring action for
 queries (#58525)

We were not properly passing around the TCB full program optimization,
so TCB generation was done per individual file. This significantly
slowed down reference resolution.

PR Close #58525
---
 .../compiler-cli/src/ngtsc/typecheck/src/checker.ts   |  2 +-
 .../migrations/signal-queries-migration/migration.ts  | 11 +++++------
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts b/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts
index a8539cd23fe..c4454c61035 100644
--- a/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts
+++ b/packages/compiler-cli/src/ngtsc/typecheck/src/checker.ts
@@ -150,7 +150,7 @@ export class TemplateTypeCheckerImpl implements TemplateTypeChecker {
   ) {}
 
   getTemplate(component: ts.ClassDeclaration, optimizeFor?: OptimizeFor): TmplAstNode[] | null {
-    const {data} = this.getLatestComponentState(component);
+    const {data} = this.getLatestComponentState(component, optimizeFor);
     if (data === null) {
       return null;
     }
diff --git a/packages/core/schematics/migrations/signal-queries-migration/migration.ts b/packages/core/schematics/migrations/signal-queries-migration/migration.ts
index e73bdcb852f..feaa8e1c636 100644
--- a/packages/core/schematics/migrations/signal-queries-migration/migration.ts
+++ b/packages/core/schematics/migrations/signal-queries-migration/migration.ts
@@ -104,6 +104,8 @@ export class SignalQueriesMigration extends TsurgeComplexMigration<
 
     // Pre-Analyze the program and get access to the template type checker.
     const {templateTypeChecker} = info.ngCompiler['ensureAnalyzed']();
+    // Generate all type check blocks.
+    templateTypeChecker.generateAllTypeCheckBlocks();
 
     const {sourceFiles, program} = info;
     const checker = program.getTypeChecker();
@@ -205,8 +207,7 @@ export class SignalQueriesMigration extends TsurgeComplexMigration<
         // Ignore references to non-query class fields.
         if (
           this.config.assumeNonBatch &&
-          descriptor !== null &&
-          !filteredQueriesForCompilationUnit.has(descriptor.key)
+          (descriptor === null || !filteredQueriesForCompilationUnit.has(descriptor.key))
         ) {
           return null;
         }
@@ -448,8 +449,6 @@ export class SignalQueriesMigration extends TsurgeComplexMigration<
       ts.forEachChild(node, queryWholeProgramVisitor);
     };
 
-    this.config.reportProgressFn?.(40, 'Tracking query declarations..');
-
     for (const sf of info.fullProgramSourceFiles) {
       ts.forEachChild(sf, queryWholeProgramVisitor);
     }
@@ -493,14 +492,14 @@ export class SignalQueriesMigration extends TsurgeComplexMigration<
       isClassWithKnownFields: (clazz) => knownQueries.getQueryFieldsOfClass(clazz) !== undefined,
     });
 
-    this.config.reportProgressFn?.(70, 'Checking inheritance..');
+    this.config.reportProgressFn?.(80, 'Checking inheritance..');
     groupedAstVisitor.execute();
 
     if (this.config.bestEffortMode) {
       filterBestEffortIncompatibilities(knownQueries);
     }
 
-    this.config.reportProgressFn?.(80, 'Migrating queries..');
+    this.config.reportProgressFn?.(90, 'Migrating queries..');
 
     // Migrate declarations.
     for (const extractedQuery of sourceQueries) {