Fix insertNodeAtClassStart for empty class with comment

Author: Andy Hanson

src/compiler/utilities.ts

@@ -4020,12 +4020,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/codefixes/helpers.ts

@@ -25,8 +25,7 @@ namespace ts.codefix {
         }
 
         const declaration = declarations[0];
-        // Clone name to remove leading trivia.
-        const name = getSynthesizedDeepClone(getNameOfDeclaration(declaration)) as PropertyName;
+        const name = getSynthesizedDeepCloneWithoutTrivia(getNameOfDeclaration(declaration)) as PropertyName;
         const visibilityModifier = createVisibilityModifier(getModifierFlags(declaration));
         const modifiers = visibilityModifier ? createNodeArray([visibilityModifier]) : undefined;
         const type = checker.getWidenedType(checker.getTypeOfSymbolAtLocation(symbol, enclosingDeclaration));
@@ -69,7 +68,7 @@ namespace ts.codefix {
 
                 for (const signature of signatures) {
                     // Need to ensure nodes are fresh each time so they can have different positions.
-                    outputMethod(signature, getSynthesizedDeepClones(modifiers), getSynthesizedDeepClone(name));
+                    outputMethod(signature, getSynthesizedDeepClones(modifiers), getSynthesizedDeepCloneWithoutTrivia(name));
                 }
 
                 if (declarations.length > signatures.length) {
@@ -104,7 +103,7 @@ namespace ts.codefix {
     }
 
     function getSynthesizedDeepClones<T extends Node>(nodes: NodeArray<T> | undefined): NodeArray<T> | undefined {
-        return nodes && createNodeArray(nodes.map(getSynthesizedDeepClone));
+        return nodes && createNodeArray(nodes.map(getSynthesizedDeepCloneWithoutTrivia));
     }
 
     export function createMethodFromCallExpression(

src/services/textChanges.ts

@@ -208,8 +208,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);
@@ -339,8 +338,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 {
@@ -435,21 +433,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 {
@@ -601,12 +598,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));
+                }
             });
         }
 
@@ -617,11 +616,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 {

src/services/utilities.ts

@@ -1488,6 +1488,12 @@ namespace ts {
         return visited;
     }
 
+    export function getSynthesizedDeepCloneWithoutTrivia<T extends Node>(node: T): T {
+        const clone = getSynthesizedDeepClone(node);
+        suppressLeadingAndTrailingTrivia(clone);
+        return clone;
+    }
+
     export function getSynthesizedDeepClones<T extends Node>(nodes: NodeArray<T> | undefined): NodeArray<T> | undefined {
         return nodes && createNodeArray(nodes.map(getSynthesizedDeepClone), nodes.hasTrailingComma);
     }

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/codeFixClassImplementInterfaceComments.ts

@@ -36,9 +36,9 @@ verify.codeFix({
         /**close-brace prefix*/ }
 /**close-brace prefix*/ }
 class C implements N.I {
-    /** property prefix */ a /** colon prefix */: N.E.a;
-    /** property prefix */ b /** colon prefix */: N.E;
-    /**method signature prefix */ foo /**open angle prefix */<X>(a: X): string {
+    a: N.E.a;
+    b: N.E;
+    foo<X>(a: X): string {
         throw new Error("Method not implemented.");
     }
 }`,

tests/cases/fourslash/codeFixClassImplementInterfaceMemberOrdering.ts

@@ -84,7 +84,6 @@ class C implements I {
     20: any;
     21: any;
     22: any;
-    /** a nice safe prime */
     23: any;
 }`,
 });

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>'",

tests/cases/fourslash/codeFixUndeclaredInStaticMethod.ts

@@ -1,7 +1,7 @@
 /// <reference path='fourslash.ts' />
 
-//// class A {[|
-////     |]static foo0() {
+//// class A {
+////     static foo0() {
 ////         this.m1(1,2,3);
 ////         A.m2(1,2);
 ////         this.prop1 = 10;
@@ -12,51 +12,89 @@
 verify.codeFix({
     description: "Declare static method 'm1'",
     index: 0,
-    newRangeContent: `
+    newFileContent:
+`class A {
     static m1(arg0: any, arg1: any, arg2: any): any {
         throw new Error("Method not implemented.");
     }
-    `,
+
+    static foo0() {
+        this.m1(1,2,3);
+        A.m2(1,2);
+        this.prop1 = 10;
+        A.prop2 = "asdf";
+    }
+}`,
 });
 
 verify.codeFix({
     description: "Declare static method 'm2'",
     index: 0,
-    newRangeContent: `
+    newFileContent:
+`class A {
     static m2(arg0: any, arg1: any): any {
         throw new Error("Method not implemented.");
     }
+
     static m1(arg0: any, arg1: any, arg2: any): any {
         throw new Error("Method not implemented.");
     }
-    `,
+
+    static foo0() {
+        this.m1(1,2,3);
+        A.m2(1,2);
+        this.prop1 = 10;
+        A.prop2 = "asdf";
+    }
+}`,
 });
 
 verify.codeFix({
     description: "Declare static property 'prop1'",
     index: 0,
-    newRangeContent: `
+    newFileContent:
+`class A {
     static prop1: number;
+
     static m2(arg0: any, arg1: any): any {
         throw new Error("Method not implemented.");
     }
+
     static m1(arg0: any, arg1: any, arg2: any): any {
         throw new Error("Method not implemented.");
     }
-    `,
+
+    static foo0() {
+        this.m1(1,2,3);
+        A.m2(1,2);
+        this.prop1 = 10;
+        A.prop2 = "asdf";
+    }
+}`,
 });
 
 verify.codeFix({
     description: "Declare static property 'prop2'",
     index: 1, // fix at index 0 is to change the spelling to 'prop1'
-    newRangeContent: `
+    newFileContent:
+`class A {
     static prop2: string;
+
     static prop1: number;
+
     static m2(arg0: any, arg1: any): any {
         throw new Error("Method not implemented.");
     }
+
     static m1(arg0: any, arg1: any, arg2: any): any {
         throw new Error("Method not implemented.");
     }
-    `,
+
+    static foo0() {
+        this.m1(1,2,3);
+        A.m2(1,2);
+        this.prop1 = 10;
+        A.prop2 = "asdf";
+    }
+}`,
 });