fix(migrations): avoid applying the same replacements twice when cleaning up unused imports (#59656)

crisbeto · AndrewKushnir · commit a62c84bc188d · 2025-01-22T07:57:38.000-08:00
If a file ends up in multiple programs, the unused imports migration was counting it twice. This was fine since the string replacements were handling it correctly, but it was printing out incorrect data. These changes rework the migration to de-duplicate the replacements and produce the data from the de-duplicated results. PR Close #59656
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
@@ -13,6 +13,7 @@ import {
   MigrationStats,
   ProgramInfo,
   projectFile,
+  ProjectFileID,
   Replacement,
   Serializable,
   TextUpdate,
@@ -28,8 +29,8 @@ export interface CompilationUnitData {
   /** Text changes that should be performed. */
   replacements: Replacement[];
 
-  /** Total number of imports that were removed. */
-  removedImports: number;
+  /** Identifiers that have been removed from each file. */
+  removedIdentifiers: NodeID[];
 
   /** Total number of files that were changed. */
   changedFiles: number;
@@ -44,7 +45,7 @@ interface RemovalLocations {
   partialRemovals: Map<ts.ArrayLiteralExpression, Set<ts.Expression>>;
 
   /** Text of all identifiers that have been removed. */
-  allRemovedIdentifiers: Set<string>;
+  allRemovedIdentifiers: Set<ts.Identifier>;
 }
 
 /** Tracks how identifiers are used across a single file. */
@@ -60,6 +61,9 @@ 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,
@@ -81,7 +85,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[] = [];
-    let removedImports = 0;
+    const removedIdentifiers: NodeID[] = [];
     let changedFiles = 0;
 
     info.ngCompiler?.getDiagnostics().forEach((diag) => {
@@ -94,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.getNodeKey(diag.start, diag.length));
+        nodePositions.get(diag.file)!.add(this.getNodeID(diag.start, diag.length));
       }
     });
 
@@ -103,14 +107,15 @@ 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, removedImports, changedFiles});
+    return confirmAsSerializable({replacements, removedIdentifiers, changedFiles});
   }
 
   override async migrate(globalData: CompilationUnitData) {
@@ -121,10 +126,34 @@ 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: [...unitA.replacements, ...unitB.replacements],
-      removedImports: unitA.removedImports + unitB.removedImports,
-      changedFiles: unitA.changedFiles + unitB.changedFiles,
+      replacements: combinedReplacements,
+      removedIdentifiers: combinedRemovedIdentifiers,
+      changedFiles: changedFileIds.size,
     });
   }
 
@@ -137,15 +166,21 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
   override async stats(globalMetadata: CompilationUnitData): Promise<MigrationStats> {
     return {
       counters: {
-        removedImports: globalMetadata.removedImports,
+        removedImports: globalMetadata.removedIdentifiers.length,
         changedFiles: globalMetadata.changedFiles,
       },
     };
   }
 
-  /** 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}`;
+  /** 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;
   }
 
   /**
@@ -176,7 +211,7 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
         return;
       }
 
-      if (locations.has(this.getNodeKey(node.getStart(), node.getWidth()))) {
+      if (locations.has(this.getNodeID(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 (
@@ -187,15 +222,15 @@ export class UnusedImportsMigration extends TsurgeFunnelMigration<
           result.fullRemovals.add(parent.initializer);
           parent.initializer.elements.forEach((element) => {
             if (ts.isIdentifier(element)) {
-              result.allRemovedIdentifiers.add(element.text);
+              result.allRemovedIdentifiers.add(element);
             }
           });
         } 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.text);
+          result.allRemovedIdentifiers.add(node);
         }
       }
     };
@@ -326,8 +361,13 @@ 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 (allRemovedIdentifiers.has(localName) && !identifierCounts.get(localName)) {
-          importManager.removeImport(sourceFile, symbolName, moduleName);
+        if (!identifierCounts.get(localName)) {
+          for (const identifier of allRemovedIdentifiers) {
+            if (identifier.text === localName) {
+              importManager.removeImport(sourceFile, symbolName, moduleName);
+              break;
+            }
+          }
         }
       });
     });
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,6 +19,7 @@ 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));
@@ -36,6 +37,7 @@ 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(
@@ -48,6 +50,7 @@ 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.
@@ -92,6 +95,7 @@ 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';
@@ -123,6 +127,7 @@ 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';
@@ -153,6 +158,7 @@ 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';
@@ -190,6 +196,7 @@ 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';
@@ -226,6 +233,7 @@ 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);
   });
 
@@ -242,6 +250,7 @@ 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);
   });
 
@@ -274,6 +283,7 @@ 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';
