refactor: ensure tsurge migrations have clear ownership of files (#61612)

JeanMeche · atscott · commit 2ae69f77e7ef · 2025-05-22T11:43:48.000-07:00
This is a patch port of #61421 PR Close #61612
diff --git a/packages/core/schematics/ng-generate/cleanup-unused-imports/unused_imports_migration.ts b/packages/core/schematics/ng-generate/cleanup-unused-imports/unused_imports_migration.ts
@@ -12,7 +12,6 @@ import {
   MigrationStats,
   ProgramInfo,
   projectFile,
-  ProjectFileID,
   Replacement,
   Serializable,
   TextUpdate,
@@ -28,8 +27,8 @@ export interface CompilationUnitData {
   /** Text changes that should be performed. */
   replacements: Replacement[];
 
-  /** Identifiers that have been removed from each file. */
-  removedIdentifiers: NodeID[];
+  /** Total number of imports that were removed. */
+  removedImports: number;
 
   /** Total number of files that were changed. */
   changedFiles: number;
@@ -44,7 +43,7 @@ interface RemovalLocations {
   partialRemovals: Map<ts.ArrayLiteralExpression, Set<ts.Expression>>;
 
   /** Text of all identifiers that have been removed. */
-  allRemovedIdentifiers: Set<ts.Identifier>;
+  allRemovedIdentifiers: Set<string>;
 }
 
 /** Tracks how identifiers are used across a single file. */
@@ -60,9 +59,6 @@ interface UsageAnalysis {
   identifierCounts: Map<string, number>;
 }
 
-/** ID of a node based on its location. */
-type NodeID = string & {__nodeID: true};
-
 /** Migration that cleans up unused imports from a project. */
 export class UnusedImportsMigration extends TsurgeFunnelMigration<
   CompilationUnitData,
@@ -84,7 +80,7 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
   override async analyze(info: ProgramInfo): Promise<Serializable<CompilationUnitData>> {
     const nodePositions = new Map<ts.SourceFile, Set<string>>();
     const replacements: Replacement[] = [];
-    const removedIdentifiers: NodeID[] = [];
+    let removedImports = 0;
     let changedFiles = 0;
 
     info.ngCompiler?.getDiagnostics().forEach((diag) => {
@@ -102,7 +98,7 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
         if (!nodePositions.has(diag.file)) {
           nodePositions.set(diag.file, new Set());
         }
-        nodePositions.get(diag.file)!.add(this.getNodeID(diag.start, diag.length));
+        nodePositions.get(diag.file)!.add(this.getNodeKey(diag.start, diag.length));
       }
     });
 
@@ -111,15 +107,14 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
       const usageAnalysis = this.analyzeUsages(sourceFile, resolvedLocations);
 
       if (resolvedLocations.allRemovedIdentifiers.size > 0) {
+        removedImports += resolvedLocations.allRemovedIdentifiers.size;
         changedFiles++;
-        resolvedLocations.allRemovedIdentifiers.forEach((identifier) => {
-          removedIdentifiers.push(this.getNodeID(identifier.getStart(), identifier.getWidth()));
-        });
       }
+
       this.generateReplacements(sourceFile, resolvedLocations, usageAnalysis, info, replacements);
     });
 
-    return confirmAsSerializable({replacements, removedIdentifiers, changedFiles});
+    return confirmAsSerializable({replacements, removedImports, changedFiles});
   }
 
   override async migrate(globalData: CompilationUnitData) {
@@ -130,34 +125,10 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
     unitA: CompilationUnitData,
     unitB: CompilationUnitData,
   ): Promise<Serializable<CompilationUnitData>> {
-    const combinedReplacements: Replacement[] = [];
-    const combinedRemovedIdentifiers: NodeID[] = [];
-    const seenReplacements = new Set<string>();
-    const seenIdentifiers = new Set<NodeID>();
-    const changedFileIds = new Set<ProjectFileID>();
-
-    [unitA, unitB].forEach((unit) => {
-      for (const replacement of unit.replacements) {
-        const key = this.getReplacementID(replacement);
-        changedFileIds.add(replacement.projectFile.id);
-        if (!seenReplacements.has(key)) {
-          seenReplacements.add(key);
-          combinedReplacements.push(replacement);
-        }
-      }
-
-      for (const identifier of unit.removedIdentifiers) {
-        if (!seenIdentifiers.has(identifier)) {
-          seenIdentifiers.add(identifier);
-          combinedRemovedIdentifiers.push(identifier);
-        }
-      }
-    });
-
     return confirmAsSerializable({
-      replacements: combinedReplacements,
-      removedIdentifiers: combinedRemovedIdentifiers,
-      changedFiles: changedFileIds.size,
+      replacements: [...unitA.replacements, ...unitB.replacements],
+      removedImports: unitA.removedImports + unitB.removedImports,
+      changedFiles: unitA.changedFiles + unitB.changedFiles,
     });
   }
 
@@ -170,21 +141,15 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
   override async stats(globalMetadata: CompilationUnitData): Promise<MigrationStats> {
     return {
       counters: {
-        removedImports: globalMetadata.removedIdentifiers.length,
+        removedImports: globalMetadata.removedImports,
         changedFiles: globalMetadata.changedFiles,
       },
     };
   }
 
-  /** Gets an ID that can be used to look up a node based on its location. */
-  private getNodeID(start: number, length: number): NodeID {
-    return `${start}/${length}` as NodeID;
-  }
-
-  /** Gets a unique ID for a replacement. */
-  private getReplacementID(replacement: Replacement): string {
-    const {position, end, toInsert} = replacement.update.data;
-    return replacement.projectFile.id + '/' + position + '/' + end + '/' + toInsert;
+  /** Gets a key that can be used to look up a node based on its location. */
+  private getNodeKey(start: number, length: number): string {
+    return `${start}/${length}`;
   }
 
   /**
@@ -215,7 +180,7 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
         return;
       }
 
-      if (locations.has(this.getNodeID(node.getStart(), node.getWidth()))) {
+      if (locations.has(this.getNodeKey(node.getStart(), node.getWidth()))) {
         // When the entire array needs to be cleared, the diagnostic is
         // reported on the property assignment, rather than an array element.
         if (
@@ -226,15 +191,15 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
           result.fullRemovals.add(parent.initializer);
           parent.initializer.elements.forEach((element) => {
             if (ts.isIdentifier(element)) {
-              result.allRemovedIdentifiers.add(element);
+              result.allRemovedIdentifiers.add(element.text);
             }
           });
         } else if (ts.isArrayLiteralExpression(parent)) {
           if (!result.partialRemovals.has(parent)) {
             result.partialRemovals.set(parent, new Set());
           }
           result.partialRemovals.get(parent)!.add(node);
-          result.allRemovedIdentifiers.add(node);
+          result.allRemovedIdentifiers.add(node.text);
         }
       }
     };
@@ -365,13 +330,8 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
       names.forEach((symbolName, localName) => {
         // Note that in the `identifierCounts` lookup both zero and undefined
         // are valid and mean that the identifiers isn't being used anymore.
-        if (!identifierCounts.get(localName)) {
-          for (const identifier of allRemovedIdentifiers) {
-            if (identifier.text === localName) {
-              importManager.removeImport(sourceFile, symbolName, moduleName);
-              break;
-            }
-          }
+        if (allRemovedIdentifiers.has(localName) && !identifierCounts.get(localName)) {
+          importManager.removeImport(sourceFile, symbolName, moduleName);
         }
       });
     });
diff --git a/packages/core/schematics/test/cleanup_unused_imports_migration_spec.ts b/packages/core/schematics/test/cleanup_unused_imports_migration_spec.ts
@@ -19,7 +19,6 @@ describe('cleanup unused imports schematic', () => {
   let tree: UnitTestTree;
   let tmpDirPath: string;
   let previousWorkingDir: string;
-  let logs: string[];
 
   function writeFile(filePath: string, contents: string) {
     host.sync.write(normalize(filePath), virtualFs.stringToFileBuffer(contents));
@@ -37,7 +36,6 @@ describe('cleanup unused imports schematic', () => {
     runner = new SchematicTestRunner('test', runfiles.resolvePackageRelative('../collection.json'));
     host = new TempScopedNodeJsSyncHost();
     tree = new UnitTestTree(new HostTree(host));
-    logs = [];
 
     writeFile('/tsconfig.json', '{}');
     writeFile(
@@ -50,7 +48,6 @@ describe('cleanup unused imports schematic', () => {
 
     previousWorkingDir = shx.pwd();
     tmpDirPath = getSystemPath(host.root);
-    runner.logger.subscribe((log) => logs.push(log.message));
 
     // Switch into the temporary directory path. This allows us to run
     // the schematic against our custom unit test tree.
@@ -95,7 +92,6 @@ describe('cleanup unused imports schematic', () => {
 
     await runMigration();
 
-    expect(logs.pop()).toBe('Removed 2 imports in 1 file');
     expect(stripWhitespace(tree.readContent('comp.ts'))).toBe(
       stripWhitespace(`
         import {Component} from '@angular/core';
@@ -127,7 +123,6 @@ describe('cleanup unused imports schematic', () => {
 
     await runMigration();
 
-    expect(logs.pop()).toBe('Removed 3 imports in 1 file');
     expect(stripWhitespace(tree.readContent('comp.ts'))).toBe(
       stripWhitespace(`
         import {Component} from '@angular/core';
@@ -158,7 +153,6 @@ describe('cleanup unused imports schematic', () => {
 
     await runMigration();
 
-    expect(logs.pop()).toBe('Removed 2 imports in 1 file');
     expect(stripWhitespace(tree.readContent('comp.ts'))).toBe(
       stripWhitespace(`
         import {Component} from '@angular/core';
@@ -196,7 +190,6 @@ describe('cleanup unused imports schematic', () => {
 
     await runMigration();
 
-    expect(logs.pop()).toBe('Removed 1 import in 1 file');
     expect(stripWhitespace(tree.readContent('comp.ts'))).toBe(
       stripWhitespace(`
         import {Component} from '@angular/core';
@@ -233,7 +226,6 @@ describe('cleanup unused imports schematic', () => {
 
     await runMigration();
 
-    expect(logs.pop()).toBe('Schematic could not find unused imports in the project');
     expect(tree.readContent('comp.ts')).toBe(initialContent);
   });
 
@@ -250,7 +242,6 @@ describe('cleanup unused imports schematic', () => {
 
     await runMigration();
 
-    expect(logs.pop()).toBe('Schematic could not find unused imports in the project');
     expect(tree.readContent('comp.ts')).toBe(initialContent);
   });
 
@@ -283,7 +274,6 @@ describe('cleanup unused imports schematic', () => {
 
     await runMigration();
 
-    expect(logs.pop()).toBe('Removed 2 imports in 1 file');
     expect(stripWhitespace(tree.readContent('comp.ts'))).toBe(
       stripWhitespace(`
         import {Component} from '@angular/core';
diff --git a/packages/core/schematics/utils/tsurge/helpers/angular_devkit/run_in_devkit.ts b/packages/core/schematics/utils/tsurge/helpers/angular_devkit/run_in_devkit.ts
@@ -17,6 +17,7 @@ import {ProjectRootRelativePath} from '../../project_paths';
 import {ProgramInfo} from '../../program_info';
 import {getProjectTsConfigPaths} from '../../../../utils/project_tsconfig_paths';
 import ts from 'typescript';
+import {MigrationStats} from '../../base_migration';
 
 export enum MigrationStage {
   /** The migration is analyzing an entrypoint */
