Fix insertNodeAtClassStart for empty class with comment (#23342)

Author: Andy

src/compiler/utilities.ts

@@ -4025,12 +4025,14 @@ namespace ts {
     }
 
     /** Add a value to a set, and return true if it wasn't already present. */
-    export function addToSeen(seen: Map<true>, key: string | number): boolean {
+    export function addToSeen(seen: Map<true>, key: string | number): boolean;
+    export function addToSeen<T>(seen: Map<T>, key: string | number, value: T): boolean;
+    export function addToSeen<T>(seen: Map<T>, key: string | number, value: T = true as any): boolean {
         key = String(key);
         if (seen.has(key)) {
             return false;
         }
-        seen.set(key, true);
+        seen.set(key, value);
         return true;
     }
 

src/services/textChanges.ts

@@ -212,8 +212,7 @@ namespace ts.textChanges {
     export class ChangeTracker {
         private readonly changes: Change[] = [];
         private readonly deletedNodesInLists: true[] = []; // Stores ids of nodes in lists that we already deleted. Used to avoid deleting `, ` twice in `a, b`.
-        // Map from class id to nodes to insert at the start
-        private readonly nodesInsertedAtClassStarts = createMap<{ sourceFile: SourceFile, cls: ClassLikeDeclaration, members: ClassElement[] }>();
+        private readonly classesWithNodesInsertedAtStart = createMap<ClassDeclaration>(); // Set<ClassDeclaration> implemented as Map<node id, ClassDeclaration>
 
         public static fromContext(context: TextChangesContext): ChangeTracker {
             return new ChangeTracker(getNewLineOrDefaultFromHost(context.host, context.formatContext.options), context.formatContext);
@@ -343,8 +342,7 @@ namespace ts.textChanges {
         }
 
         public insertNodeBefore(sourceFile: SourceFile, before: Node, newNode: Node, blankLineBetween = false) {
-            const pos = getAdjustedStartPosition(sourceFile, before, {}, Position.Start);
-            return this.replaceRange(sourceFile, { pos, end: pos }, newNode, this.getOptionsForInsertNodeBefore(before, blankLineBetween));
+            this.insertNodeAt(sourceFile, getAdjustedStartPosition(sourceFile, before, {}, Position.Start), newNode, this.getOptionsForInsertNodeBefore(before, blankLineBetween));
         }
 
         public insertModifierBefore(sourceFile: SourceFile, modifier: SyntaxKind, before: Node): void {
@@ -443,21 +441,20 @@ namespace ts.textChanges {
         }
 
         public insertNodeAtClassStart(sourceFile: SourceFile, cls: ClassLikeDeclaration, newElement: ClassElement): void {
-            const firstMember = firstOrUndefined(cls.members);
-            if (!firstMember) {
-                const id = getNodeId(cls).toString();
-                const newMembers = this.nodesInsertedAtClassStarts.get(id);
-                if (newMembers) {
-                    Debug.assert(newMembers.sourceFile === sourceFile && newMembers.cls === cls);
-                    newMembers.members.push(newElement);
+            const clsStart = cls.getStart(sourceFile);
+            let prefix = "";
+            let suffix = this.newLineCharacter;
+            if (addToSeen(this.classesWithNodesInsertedAtStart, getNodeId(cls), cls)) {
+                prefix = this.newLineCharacter;
+                // For `class C {\n}`, don't add the trailing "\n"
+                if (cls.members.length === 0 && !(positionsAreOnSameLine as any)(...getClassBraceEnds(cls, sourceFile), sourceFile)) { // TODO: GH#4130 remove 'as any'
+                    suffix = "";
                 }
-                else {
-                    this.nodesInsertedAtClassStarts.set(id, { sourceFile, cls, members: [newElement] });
-                }
-            }
-            else {
-                this.insertNodeBefore(sourceFile, firstMember, newElement);
             }
+
+            const indentation = formatting.SmartIndenter.findFirstNonWhitespaceColumn(getLineStartPositionForPosition(clsStart, sourceFile), clsStart, sourceFile, this.formatContext.options)
+                + this.formatContext.options.indentSize;
+            this.insertNodeAt(sourceFile, cls.members.pos, newElement, { indentation, prefix, suffix });
         }
 
         public insertNodeAfter(sourceFile: SourceFile, after: Node, newNode: Node): this {
@@ -638,12 +635,14 @@ namespace ts.textChanges {
             return this;
         }
 
-        private finishInsertNodeAtClassStart(): void {
-            this.nodesInsertedAtClassStarts.forEach(({ sourceFile, cls, members }) => {
-                const newCls = cls.kind === SyntaxKind.ClassDeclaration
-                    ? updateClassDeclaration(cls, cls.decorators, cls.modifiers, cls.name, cls.typeParameters, cls.heritageClauses, members)
-                    : updateClassExpression(cls, cls.modifiers, cls.name, cls.typeParameters, cls.heritageClauses, members);
-                this.replaceNode(sourceFile, cls, newCls);
+        private finishClassesWithNodesInsertedAtStart(): void {
+            this.classesWithNodesInsertedAtStart.forEach(cls => {
+                const sourceFile = cls.getSourceFile();
+                const [openBraceEnd, closeBraceEnd] = getClassBraceEnds(cls, sourceFile);
+                // For `class C { }` remove the whitespace inside the braces.
+                if (positionsAreOnSameLine(openBraceEnd, closeBraceEnd, sourceFile) && openBraceEnd !== closeBraceEnd - 1) {
+                    this.deleteRange(sourceFile, createTextRange(openBraceEnd, closeBraceEnd - 1));
+                }
             });
         }
 
@@ -654,11 +653,15 @@ namespace ts.textChanges {
          *    so we can only call this once and can't get the non-formatted text separately.
          */
         public getChanges(validate?: ValidateNonFormattedText): FileTextChanges[] {
-            this.finishInsertNodeAtClassStart();
+            this.finishClassesWithNodesInsertedAtStart();
             return changesToText.getTextChangesFromChanges(this.changes, this.newLineCharacter, this.formatContext, validate);
         }
     }
 
+    function getClassBraceEnds(cls: ClassLikeDeclaration, sourceFile: SourceFile): [number, number] {
+        return [findChildOfKind(cls, SyntaxKind.OpenBraceToken, sourceFile).end, findChildOfKind(cls, SyntaxKind.CloseBraceToken, sourceFile).end];
+    }
+
     export type ValidateNonFormattedText = (node: Node, text: string) => void;
 
     namespace changesToText {

tests/cases/fourslash/codeFixAddMissingMember.ts

@@ -11,6 +11,7 @@ verify.codeFix({
     index: 0,
     newFileContent: `class C {
     foo: number;
+
     method() {
         this.foo = 10;
     }

tests/cases/fourslash/codeFixAddMissingMember2.ts

@@ -11,6 +11,7 @@ verify.codeFix({
     index: 1,
     newFileContent: `class C {
     [x: string]: number;
+
     method() {
         this.foo = 10;
     }

tests/cases/fourslash/codeFixAddMissingMember3.ts

@@ -11,6 +11,7 @@ verify.codeFix({
     index: 0,
     newFileContent: `class C {
     static foo: number;
+
     static method() {
         this.foo = 10;
     }

tests/cases/fourslash/codeFixAddMissingMember_all.ts

@@ -17,6 +17,7 @@ verify.codeFixAll({
     y(): any {
         throw new Error("Method not implemented.");
     }
+
     method() {
         this.x = 0;
         this.y();

tests/cases/fourslash/codeFixAddMissingMember_all_js.ts

@@ -21,6 +21,7 @@ verify.codeFixAll({
     y() {
         throw new Error("Method not implemented.");
     }
+
     constructor() {
         this.x = undefined;
     }

tests/cases/fourslash/codeFixClassExtendAbstractMethod_comment.ts

@@ -0,0 +1,24 @@
+/// <reference path='fourslash.ts' />
+
+////abstract class A {
+////    abstract m() : void;
+////}
+////
+////class B extends A {
+////    // comment
+////}
+
+verify.codeFix({
+    description: "Implement inherited abstract class",
+    newFileContent:
+`abstract class A {
+    abstract m() : void;
+}
+
+class B extends A {
+    m(): void {
+        throw new Error("Method not implemented.");
+    }
+    // comment
+}`,
+});

tests/cases/fourslash/codeFixClassImplementClassMultipleSignatures2.ts

@@ -6,7 +6,7 @@
 ////    method(a: string): Function;
 ////    method(a: string | number, b?: string | number): boolean | Function { return a + b as any; }
 ////}
-////class C implements A {[| |]}
+////class C implements A { }
 
 verify.codeFix({
     description: "Implement interface 'A'",

tests/cases/fourslash/codeFixClassImplementInterfaceTypeParamInstantiateNumber.ts

@@ -1,7 +1,7 @@
 /// <reference path='fourslash.ts' />
 
 ////interface I<T> { x: T; }
-////class C implements I<number> {[| |]}
+////class C implements I<number> { }
 
 verify.codeFix({
     description: "Implement interface 'I<number>'",