InsertNode instead of replace + handling call expressions

Author: iliashkolyar

src/services/codefixes/fixAddMissingNewOperator.ts

@@ -6,23 +6,29 @@ namespace ts.codefix {
         errorCodes,
         getCodeActions(context) {
             const { sourceFile, span } = context;
-            const missingNewExpression = getMissingNewExpression(sourceFile, span.start);
-            const changes = textChanges.ChangeTracker.with(context, t => addMissingNewOperator(t, sourceFile, missingNewExpression));
+            const changes = textChanges.ChangeTracker.with(context, t => addMissingNewOperator(t, sourceFile, span));
             return [createCodeFixAction(fixId, changes, Diagnostics.Add_missing_new_operator_to_call, fixId, Diagnostics.Add_missing_new_operator_to_all_calls)];
         },
         fixIds: [fixId],
         getAllCodeActions: context => codeFixAll(context, errorCodes, (changes, diag) =>
-            addMissingNewOperator(changes, context.sourceFile, getMissingNewExpression(diag.file, diag.start))),
+            addMissingNewOperator(changes, context.sourceFile, diag)),
     });
 
-    function getMissingNewExpression(sourceFile: SourceFile, pos: number): Expression {
-        const token = getTokenAtPosition(sourceFile, pos);
-        Debug.assert(isCallExpression(token.parent));
-        return <Expression>token;
+    function addMissingNewOperator(changes: textChanges.ChangeTracker, sourceFile: SourceFile, span: TextSpan): void {
+        const call = cast(findAncestorMatchingSpan(sourceFile, span), isCallExpression);
+        changes.insertNodeAt(sourceFile, span.start, createToken(SyntaxKind.NewKeyword), { suffix: " " });
+        if (isCallExpression(call.expression)) {
+            changes.insertNodeAt(sourceFile, call.expression.getStart(sourceFile), createToken(SyntaxKind.OpenParenToken));
+            changes.insertNodeAt(sourceFile, call.expression.end, createToken(SyntaxKind.CloseParenToken));
+        }
     }
 
-    function addMissingNewOperator(changes: textChanges.ChangeTracker, sourceFile: SourceFile, missingNewExpression: Expression): void {
-        const newTypeNode = createNew(missingNewExpression, /*typeArguments*/ undefined, /*argumentsArray*/ undefined);
-        changes.replaceNode(sourceFile, missingNewExpression, newTypeNode);
+    function findAncestorMatchingSpan(sourceFile: SourceFile, span: TextSpan): Node {
+        let token = getTokenAtPosition(sourceFile, span.start);
+        const end = textSpanEnd(span);
+        while (token.end < end) {
+            token = token.parent;
+        }
+        return token;
     }
 }

tests/cases/fourslash/codeFixAddMissingNew.ts

@@ -7,7 +7,8 @@
 verify.codeFix({
     description: "Add missing 'new' operator to call",
     index: 0,
-    newFileContent: `class C {
+    newFileContent:
+`class C {
 }
 var c = new C();`
 });

tests/cases/fourslash/codeFixAddMissingNew2.ts

@@ -0,0 +1,14 @@
+/// <reference path='fourslash.ts' />
+
+////class C {
+////}
+////let x = (() => C)()();
+
+verify.codeFix({
+    description: "Add missing 'new' operator to call",
+    index: 0,
+    newFileContent:
+`class C {
+}
+let x = new ((() => C)())();`
+});

tests/cases/fourslash/codeFixAddMissingNew3.ts

@@ -0,0 +1,16 @@
+/// <reference path='fourslash.ts' />
+
+////class C {
+////}
+////let x = [C];
+////let a = x[0]();
+
+verify.codeFix({
+    description: "Add missing 'new' operator to call",
+    index: 0,
+    newFileContent:
+`class C {
+}
+let x = [C];
+let a = new x[0]();`
+});

tests/cases/fourslash/codeFixAddMissingNew4.ts

@@ -0,0 +1,18 @@
+/// <reference path='fourslash.ts' />
+
+////class C {
+////}
+////class D {
+////}
+////let x = (true ? C : D)();
+
+verify.codeFix({
+    description: "Add missing 'new' operator to call",
+    index: 0,
+    newFileContent:
+`class C {
+}
+class D {
+}
+let x = new (true ? C : D)();`
+});