annotateWithTypeFromJSDoc: Do less special-casing for arrow functions

src/services/codefixes/annotateWithTypeFromJSDoc.ts

@@ -42,28 +42,22 @@ namespace ts.codefix {
 
     function doChange(changes: textChanges.ChangeTracker, sourceFile: SourceFile, decl: DeclarationWithType): void {
         if (isFunctionLikeDeclaration(decl) && (getJSDocReturnType(decl) || decl.parameters.some(p => !!getJSDocType(p)))) {
-            const typeParameters = getJSDocTypeParameterDeclarations(decl);
-            const returnType = getJSDocReturnType(decl);
-            const returnTypeNode = returnType && transformJSDocType(returnType);
-
-            if (isArrowFunction(decl) && !findChildOfKind(decl, SyntaxKind.OpenParenToken, sourceFile)) {
-                const params = decl.parameters.map(p => {
-                    const paramType = getJSDocType(p);
-                    return paramType && !p.type ? updateParameter(p, p.decorators, p.modifiers, p.dotDotDotToken, p.name, p.questionToken, transformJSDocType(paramType), p.initializer) : p;
-                });
-                changes.replaceNode(sourceFile, decl, updateArrowFunction(decl, decl.modifiers, decl.typeParameters || typeParameters, params, decl.type || returnTypeNode, decl.equalsGreaterThanToken, decl.body));
+            if (!decl.typeParameters) {
+                const typeParameters = getJSDocTypeParameterDeclarations(decl);
+                if (typeParameters) changes.insertTypeParameters(sourceFile, decl, typeParameters);
             }
-            else {
-                if (typeParameters && !decl.typeParameters) {
-                    changes.insertTypeParameters(sourceFile, decl, typeParameters);
+            const needParens = isArrowFunction(decl) && !findChildOfKind(decl, SyntaxKind.OpenParenToken, sourceFile);
+            if (needParens) changes.insertNodeBefore(sourceFile, first(decl.parameters), createToken(SyntaxKind.OpenParenToken));
+            for (const param of decl.parameters) {
+                if (!param.type) {
+                    const paramType = getJSDocType(param);
+                    if (paramType) changes.insertTypeAnnotation(sourceFile, param, transformJSDocType(paramType));
                 }
-                for (const param of decl.parameters) {
-                    if (!param.type) {
-                        const paramType = getJSDocType(param);
-                        if (paramType) changes.insertTypeAnnotation(sourceFile, param, transformJSDocType(paramType));
-                    }
-                }
-                if (returnTypeNode && !decl.type) changes.insertTypeAnnotation(sourceFile, decl, returnTypeNode);
+            }
+            if (needParens) changes.insertNodeAfter(sourceFile, last(decl.parameters), createToken(SyntaxKind.CloseParenToken));
+            if (!decl.type) {
+                const returnType = getJSDocReturnType(decl);
+                if (returnType) changes.insertTypeAnnotation(sourceFile, decl, transformJSDocType(returnType));
             }
         }
         else {

src/services/textChanges.ts

@@ -352,14 +352,16 @@ namespace ts.textChanges {
         /** Prefer this over replacing a node with another that has a type annotation, as it avoids reformatting the other parts of the node. */
         public insertTypeAnnotation(sourceFile: SourceFile, node: TypeAnnotatable, type: TypeNode): void {
             const end = (isFunctionLike(node)
-                ? findChildOfKind(node, SyntaxKind.CloseParenToken, sourceFile)!
+                // If no `)`, is an arrow function `x => x`, so use the end of the first parameter
+                ? findChildOfKind(node, SyntaxKind.CloseParenToken, sourceFile) || first(node.parameters)
                 : node.kind !== SyntaxKind.VariableDeclaration && node.questionToken ? node.questionToken : node.name).end;
             this.insertNodeAt(sourceFile, end, type, { prefix: ": " });
         }
 
         public insertTypeParameters(sourceFile: SourceFile, node: SignatureDeclaration, typeParameters: ReadonlyArray<TypeParameterDeclaration>): void {
-            const lparen = findChildOfKind(node, SyntaxKind.OpenParenToken, sourceFile)!.pos;
-            this.insertNodesAt(sourceFile, lparen, typeParameters, { prefix: "<", suffix: ">" });
+            // If no `(`, is an arrow function `x => x`, so use the pos of the first parameter
+            const start = (findChildOfKind(node, SyntaxKind.OpenParenToken, sourceFile) || first(node.parameters)).getStart(sourceFile);
+            this.insertNodesAt(sourceFile, start, typeParameters, { prefix: "<", suffix: ">" });
         }
 
         private getOptionsForInsertNodeBefore(before: Node, doubleNewlines: boolean): ChangeNodeOptions {
@@ -369,6 +371,9 @@ namespace ts.textChanges {
             else if (isVariableDeclaration(before)) { // insert `x = 1, ` into `const x = 1, y = 2;
                 return { suffix: ", " };
             }
+            else if (isParameter(before)) {
+                return {};
+            }
             return Debug.failBadSyntaxKind(before); // We haven't handled this kind of node yet -- add it
         }
 
@@ -453,6 +458,9 @@ namespace ts.textChanges {
             else if (isVariableDeclaration(node)) {
                 return { prefix: ", " };
             }
+            else if (isParameter(node)) {
+                return {};
+            }
             return Debug.failBadSyntaxKind(node); // We haven't handled this kind of node yet -- add it
         }
 

tests/cases/fourslash/annotateWithTypeFromJSDoc9.5.ts

@@ -0,0 +1,19 @@
+/// <reference path='fourslash.ts' />
+
+/////**
+//// * @template {T}
+//// * @param {T} x
+//// * @returns {T}
+//// */
+////var f = /*a*/x/*b*/ => x
+
+verify.codeFix({
+    description: "Annotate with type from JSDoc",
+    newFileContent:
+`/**
+ * @template {T}
+ * @param {T} x
+ * @returns {T}
+ */
+var f = <T>(x: T): T => x`,
+});