Improve ChangeTracker#deleteNodeInList (#24221)

Author: Andy

src/compiler/core.ts

@@ -320,6 +320,15 @@ namespace ts {
         return -1;
     }
 
+    export function findLastIndex<T>(array: ReadonlyArray<T>, predicate: (element: T, index: number) => boolean, startIndex?: number): number {
+        for (let i = startIndex === undefined ? array.length - 1 : startIndex; i >= 0; i--) {
+            if (predicate(array[i], i)) {
+                return i;
+            }
+        }
+        return -1;
+    }
+
     /**
      * Returns the first truthy result of `callback`, or else fails.
      * This is like `forEach`, but never returns undefined.

src/services/textChanges.ts

@@ -212,7 +212,7 @@ namespace ts.textChanges {
     export class ChangeTracker {
         private readonly changes: Change[] = [];
         private readonly newFiles: { readonly oldFile: SourceFile, readonly fileName: string, readonly statements: ReadonlyArray<Statement> }[] = [];
-        private readonly deletedNodesInLists: true[] = []; // Stores ids of nodes in lists that we already deleted. Used to avoid deleting `, ` twice in `a, b`.
+        private readonly deletedNodesInLists = new NodeSet(); // Stores ids of nodes in lists that we already deleted. Used to avoid deleting `, ` twice in `a, b`.
         private readonly classesWithNodesInsertedAtStart = createMap<ClassDeclaration>(); // Set<ClassDeclaration> implemented as Map<node id, ClassDeclaration>
 
         public static fromContext(context: TextChangesContext): ChangeTracker {
@@ -262,35 +262,15 @@ namespace ts.textChanges {
                 this.deleteNode(sourceFile, node);
                 return this;
             }
-            const id = getNodeId(node);
-            Debug.assert(!this.deletedNodesInLists[id], "Deleting a node twice");
-            this.deletedNodesInLists[id] = true;
-            if (index !== containingList.length - 1) {
-                const nextToken = getTokenAtPosition(sourceFile, node.end, /*includeJsDocComment*/ false);
-                if (nextToken && isSeparator(node, nextToken)) {
-                    // find first non-whitespace position in the leading trivia of the node
-                    const startPosition = skipTrivia(sourceFile.text, getAdjustedStartPosition(sourceFile, node, {}, Position.FullStart), /*stopAfterLineBreak*/ false, /*stopAtComments*/ true);
-                    const nextElement = containingList[index + 1];
-                    /// find first non-whitespace position in the leading trivia of the next node
-                    const endPosition = skipTrivia(sourceFile.text, getAdjustedStartPosition(sourceFile, nextElement, {}, Position.FullStart), /*stopAfterLineBreak*/ false, /*stopAtComments*/ true);
-                    // shift next node so its first non-whitespace position will be moved to the first non-whitespace position of the deleted node
-                    this.deleteRange(sourceFile, { pos: startPosition, end: endPosition });
-                }
-            }
-            else {
-                const prev = containingList[index - 1];
-                if (this.deletedNodesInLists[getNodeId(prev)]) {
-                    const pos = skipTrivia(sourceFile.text, getAdjustedStartPosition(sourceFile, node, {}, Position.FullStart), /*stopAfterLineBreak*/ false, /*stopAtComments*/ true);
-                    const end = getAdjustedEndPosition(sourceFile, node, {});
-                    this.deleteRange(sourceFile, { pos, end });
-                }
-                else {
-                    const previousToken = getTokenAtPosition(sourceFile, containingList[index - 1].end, /*includeJsDocComment*/ false);
-                    if (previousToken && isSeparator(node, previousToken)) {
-                        this.deleteNodeRange(sourceFile, previousToken, node);
-                    }
-                }
-            }
+
+            // Note: We will only delete a comma *after* a node. This will leave a trailing comma if we delete the last node.
+            // That's handled in the end by `finishTrailingCommaAfterDeletingNodesInList`.
+            Debug.assert(!this.deletedNodesInLists.has(node), "Deleting a node twice");
+            this.deletedNodesInLists.add(node);
+            this.deleteRange(sourceFile, {
+                pos: startPositionToDeleteNodeInList(sourceFile, node),
+                end: index === containingList.length - 1 ? getAdjustedEndPosition(sourceFile, node, {}) : startPositionToDeleteNodeInList(sourceFile, containingList[index + 1]),
+            });
             return this;
         }
 
@@ -683,6 +663,19 @@ namespace ts.textChanges {
             });
         }
 
+        private finishTrailingCommaAfterDeletingNodesInList() {
+            this.deletedNodesInLists.forEach(node => {
+                const sourceFile = node.getSourceFile();
+                const list = formatting.SmartIndenter.getContainingList(node, sourceFile);
+                if (node !== last(list)) return;
+
+                const lastNonDeletedIndex = findLastIndex(list, n => !this.deletedNodesInLists.has(n), list.length - 2);
+                if (lastNonDeletedIndex !== -1) {
+                    this.deleteRange(sourceFile, { pos: list[lastNonDeletedIndex].end, end: startPositionToDeleteNodeInList(sourceFile, list[lastNonDeletedIndex + 1]) });
+                }
+            });
+        }
+
         /**
          * Note: after calling this, the TextChanges object must be discarded!
          * @param validate only for tests
@@ -691,6 +684,7 @@ namespace ts.textChanges {
          */
         public getChanges(validate?: ValidateNonFormattedText): FileTextChanges[] {
             this.finishClassesWithNodesInsertedAtStart();
+            this.finishTrailingCommaAfterDeletingNodesInList();
             const changes = changesToText.getTextChangesFromChanges(this.changes, this.newLineCharacter, this.formatContext, validate);
             for (const { oldFile, fileName, statements } of this.newFiles) {
                 changes.push(changesToText.newFileChanges(oldFile, fileName, statements, this.newLineCharacter));
@@ -703,6 +697,11 @@ namespace ts.textChanges {
         }
     }
 
+    // find first non-whitespace position in the leading trivia of the node
+    function startPositionToDeleteNodeInList(sourceFile: SourceFile, node: Node): number {
+        return skipTrivia(sourceFile.text, getAdjustedStartPosition(sourceFile, node, {}, Position.FullStart), /*stopAfterLineBreak*/ false, /*stopAtComments*/ true);
+    }
+
     function getClassBraceEnds(cls: ClassLikeDeclaration, sourceFile: SourceFile): [number, number] {
         return [findChildOfKind(cls, SyntaxKind.OpenBraceToken, sourceFile).end, findChildOfKind(cls, SyntaxKind.CloseBraceToken, sourceFile).end];
     }

src/services/utilities.ts

@@ -1280,6 +1280,23 @@ namespace ts {
         }
         return propSymbol;
     }
+
+    export class NodeSet {
+        private map = createMap<Node>();
+
+        add(node: Node): void {
+            this.map.set(String(getNodeId(node)), node);
+        }
+        has(node: Node): boolean {
+            return this.map.has(String(getNodeId(node)));
+        }
+        forEach(cb: (node: Node) => void): void {
+            this.map.forEach(cb);
+        }
+        some(pred: (node: Node) => boolean): boolean {
+            return forEachEntry(this.map, pred) || false;
+        }
+    }
 }
 
 // Display-part writer helpers

tests/cases/fourslash/codeFixUnusedIdentifier_all_delete.ts

@@ -6,11 +6,13 @@
 ////function f(a, b) {
 ////    const x = 0;
 ////}
+////function g(a, b, c) { return a; }
 
 verify.codeFixAll({
     fixId: "unusedIdentifier_delete",
     fixAllDescription: "Delete all unused declarations",
     newFileContent:
 `function f() {
-}`,
+}
+function g(a) { return a; }`,
 });

tests/cases/fourslash/incompleteFunctionCallCodefix2.ts

@@ -5,5 +5,6 @@
 
 verify.codeFix({
     description: "Prefix 'C' with an underscore",
+    index: 1,
     newFileContent: "function f(new _C(100, 3, undefined)",
 });