feat(51870): Add a Quick Fix to add an additional parameter to a method or function (#56411)

Author: a-tarasyuk

src/compiler/diagnosticMessages.json

@@ -7792,6 +7792,30 @@
         "category": "Message",
         "code": 95187
     },
+    "Add missing parameter to '{0}'": {
+        "category": "Message",
+        "code": 95188
+    },
+    "Add missing parameters to '{0}'": {
+        "category": "Message",
+        "code": 95189
+    },
+    "Add all missing parameters": {
+        "category": "Message",
+        "code": 95190
+    },
+    "Add optional parameter to '{0}'": {
+        "category": "Message",
+        "code": 95191
+    },
+    "Add optional parameters to '{0}'": {
+        "category": "Message",
+        "code": 95192
+    },
+    "Add all optional parameters": {
+        "category": "Message",
+        "code": 95193
+    },
 
     "No value exists in scope for the shorthand property '{0}'. Either declare one or provide an initializer.": {
         "category": "Error",

src/services/_namespaces/ts.codefix.ts

@@ -31,6 +31,7 @@ export * from "../codefixes/fixSpelling";
 export * from "../codefixes/returnValueCorrect";
 export * from "../codefixes/fixAddMissingMember";
 export * from "../codefixes/fixAddMissingNewOperator";
+export * from "../codefixes/fixAddMissingParam";
 export * from "../codefixes/fixCannotFindModule";
 export * from "../codefixes/fixClassDoesntImplementInheritedAbstractMember";
 export * from "../codefixes/fixClassSuperMustPrecedeThisAccess";

src/services/codefixes/fixAddMissingParam.ts

@@ -0,0 +1,327 @@
+import {
+    append,
+    ArrowFunction,
+    CodeFixAction,
+    declarationNameToString,
+    Diagnostics,
+    factory,
+    filter,
+    findAncestor,
+    first,
+    forEach,
+    FunctionDeclaration,
+    FunctionExpression,
+    FunctionLikeDeclaration,
+    getNameOfAccessExpression,
+    getNameOfDeclaration,
+    getTokenAtPosition,
+    isAccessExpression,
+    isCallExpression,
+    isIdentifier,
+    isParameter,
+    isPropertyDeclaration,
+    isSourceFileFromLibrary,
+    isVariableDeclaration,
+    last,
+    lastOrUndefined,
+    length,
+    map,
+    MethodDeclaration,
+    Node,
+    NodeBuilderFlags,
+    ParameterDeclaration,
+    Program,
+    QuestionToken,
+    some,
+    SourceFile,
+    SyntaxKind,
+    textChanges,
+    Type,
+    TypeChecker,
+    TypeNode,
+} from "../_namespaces/ts";
+import {
+    codeFixAll,
+    createCodeFixAction,
+    registerCodeFix,
+} from "../_namespaces/ts.codefix";
+
+const addMissingParamFixId = "addMissingParam";
+const addOptionalParamFixId = "addOptionalParam";
+const errorCodes = [Diagnostics.Expected_0_arguments_but_got_1.code];
+
+registerCodeFix({
+    errorCodes,
+    fixIds: [addMissingParamFixId, addOptionalParamFixId],
+    getCodeActions(context) {
+        const info = getInfo(context.sourceFile, context.program, context.span.start);
+        if (info === undefined) return undefined;
+
+        const { name, declarations, newParameters, newOptionalParameters } = info;
+        const actions: CodeFixAction[] = [];
+
+        if (length(newParameters)) {
+            append(
+                actions,
+                createCodeFixAction(
+                    addMissingParamFixId,
+                    textChanges.ChangeTracker.with(context, t => doChange(t, context.sourceFile, declarations, newParameters)),
+                    [length(newParameters) > 1 ? Diagnostics.Add_missing_parameters_to_0 : Diagnostics.Add_missing_parameter_to_0, name],
+                    addMissingParamFixId,
+                    Diagnostics.Add_all_missing_parameters,
+                ),
+            );
+        }
+
+        if (length(newOptionalParameters)) {
+            append(
+                actions,
+                createCodeFixAction(
+                    addOptionalParamFixId,
+                    textChanges.ChangeTracker.with(context, t => doChange(t, context.sourceFile, declarations, newOptionalParameters)),
+                    [length(newOptionalParameters) > 1 ? Diagnostics.Add_optional_parameters_to_0 : Diagnostics.Add_optional_parameter_to_0, name],
+                    addOptionalParamFixId,
+                    Diagnostics.Add_all_optional_parameters,
+                ),
+            );
+        }
+
+        return actions;
+    },
+    getAllCodeActions: context =>
+        codeFixAll(context, errorCodes, (changes, diag) => {
+            const info = getInfo(context.sourceFile, context.program, diag.start);
+            if (info) {
+                const { declarations, newParameters, newOptionalParameters } = info;
+                if (context.fixId === addMissingParamFixId) {
+                    doChange(changes, context.sourceFile, declarations, newParameters);
+                }
+                if (context.fixId === addOptionalParamFixId) {
+                    doChange(changes, context.sourceFile, declarations, newOptionalParameters);
+                }
+            }
+        }),
+});
+
+type ConvertibleSignatureDeclaration =
+    | FunctionDeclaration
+    | FunctionExpression
+    | ArrowFunction
+    | MethodDeclaration;
+
+interface SignatureInfo {
+    readonly newParameters: ParameterInfo[];
+    readonly newOptionalParameters: ParameterInfo[];
+    readonly name: string;
+    readonly declarations: ConvertibleSignatureDeclaration[];
+}
+
+interface ParameterInfo {
+    readonly pos: number;
+    readonly declaration: ParameterDeclaration;
+}
+
+function getInfo(sourceFile: SourceFile, program: Program, pos: number): SignatureInfo | undefined {
+    const token = getTokenAtPosition(sourceFile, pos);
+    const callExpression = findAncestor(token, isCallExpression);
+    if (callExpression === undefined || length(callExpression.arguments) === 0) {
+        return undefined;
+    }
+
+    const checker = program.getTypeChecker();
+    const type = checker.getTypeAtLocation(callExpression.expression);
+    const convertibleSignatureDeclarations = filter(type.symbol.declarations, isConvertibleSignatureDeclaration);
+    if (convertibleSignatureDeclarations === undefined) {
+        return undefined;
+    }
+
+    const nonOverloadDeclaration = lastOrUndefined(convertibleSignatureDeclarations);
+    if (
+        nonOverloadDeclaration === undefined ||
+        nonOverloadDeclaration.body === undefined ||
+        isSourceFileFromLibrary(program, nonOverloadDeclaration.getSourceFile())
+    ) {
+        return undefined;
+    }
+
+    const name = tryGetName(nonOverloadDeclaration);
+    if (name === undefined) {
+        return undefined;
+    }
+
+    const newParameters: ParameterInfo[] = [];
+    const newOptionalParameters: ParameterInfo[] = [];
+    const parametersLength = length(nonOverloadDeclaration.parameters);
+    const argumentsLength = length(callExpression.arguments);
+    if (parametersLength > argumentsLength) {
+        return undefined;
+    }
+
+    const declarations = [nonOverloadDeclaration, ...getOverloads(nonOverloadDeclaration, convertibleSignatureDeclarations)];
+    for (let i = 0, pos = 0, paramIndex = 0; i < argumentsLength; i++) {
+        const arg = callExpression.arguments[i];
+        const expr = isAccessExpression(arg) ? getNameOfAccessExpression(arg) : arg;
+        const type = checker.getWidenedType(checker.getBaseTypeOfLiteralType(checker.getTypeAtLocation(arg)));
+        const parameter = pos < parametersLength ? nonOverloadDeclaration.parameters[pos] : undefined;
+        if (
+            parameter &&
+            checker.isTypeAssignableTo(type, checker.getTypeAtLocation(parameter))
+        ) {
+            pos++;
+            continue;
+        }
+
+        const name = expr && isIdentifier(expr) ? expr.text : `p${paramIndex++}`;
+        const typeNode = typeToTypeNode(checker, type, nonOverloadDeclaration);
+        append(newParameters, {
+            pos: i,
+            declaration: createParameter(name, typeNode, /*questionToken*/ undefined),
+        });
+
+        if (isOptionalPos(declarations, pos)) {
+            continue;
+        }
+
+        append(newOptionalParameters, {
+            pos: i,
+            declaration: createParameter(name, typeNode, factory.createToken(SyntaxKind.QuestionToken)),
+        });
+    }
+
+    return {
+        newParameters,
+        newOptionalParameters,
+        name: declarationNameToString(name),
+        declarations,
+    };
+}
+
+function tryGetName(node: FunctionLikeDeclaration) {
+    const name = getNameOfDeclaration(node);
+    if (name) {
+        return name;
+    }
+
+    if (
+        isVariableDeclaration(node.parent) && isIdentifier(node.parent.name) ||
+        isPropertyDeclaration(node.parent) ||
+        isParameter(node.parent)
+    ) {
+        return node.parent.name;
+    }
+}
+
+function typeToTypeNode(checker: TypeChecker, type: Type, enclosingDeclaration: Node) {
+    return checker.typeToTypeNode(checker.getWidenedType(type), enclosingDeclaration, NodeBuilderFlags.NoTruncation)
+        ?? factory.createKeywordTypeNode(SyntaxKind.UnknownKeyword);
+}
+
+function doChange(
+    changes: textChanges.ChangeTracker,
+    sourceFile: SourceFile,
+    declarations: ConvertibleSignatureDeclaration[],
+    newParameters: ParameterInfo[],
+) {
+    forEach(declarations, declaration => {
+        if (length(declaration.parameters)) {
+            changes.replaceNodeRangeWithNodes(
+                sourceFile,
+                first(declaration.parameters),
+                last(declaration.parameters),
+                updateParameters(declaration, newParameters),
+                {
+                    joiner: ", ",
+                    indentation: 0,
+                    leadingTriviaOption: textChanges.LeadingTriviaOption.IncludeAll,
+                    trailingTriviaOption: textChanges.TrailingTriviaOption.Include,
+                },
+            );
+        }
+        else {
+            forEach(updateParameters(declaration, newParameters), (parameter, index) => {
+                if (length(declaration.parameters) === 0 && index === 0) {
+                    changes.insertNodeAt(sourceFile, declaration.parameters.end, parameter);
+                }
+                else {
+                    changes.insertNodeAtEndOfList(sourceFile, declaration.parameters, parameter);
+                }
+            });
+        }
+    });
+}
+
+function isConvertibleSignatureDeclaration(node: Node): node is ConvertibleSignatureDeclaration {
+    switch (node.kind) {
+        case SyntaxKind.FunctionDeclaration:
+        case SyntaxKind.FunctionExpression:
+        case SyntaxKind.MethodDeclaration:
+        case SyntaxKind.ArrowFunction:
+            return true;
+        default:
+            return false;
+    }
+}
+
+function updateParameters(node: ConvertibleSignatureDeclaration, newParameters: readonly ParameterInfo[]) {
+    const parameters = map(node.parameters, p =>
+        factory.createParameterDeclaration(
+            p.modifiers,
+            p.dotDotDotToken,
+            p.name,
+            p.questionToken,
+            p.type,
+            p.initializer,
+        ));
+    for (const { pos, declaration } of newParameters) {
+        const prev = pos > 0 ? parameters[pos - 1] : undefined;
+        parameters.splice(
+            pos,
+            0,
+            factory.updateParameterDeclaration(
+                declaration,
+                declaration.modifiers,
+                declaration.dotDotDotToken,
+                declaration.name,
+                prev && prev.questionToken ? factory.createToken(SyntaxKind.QuestionToken) : declaration.questionToken,
+                declaration.type,
+                declaration.initializer,
+            ),
+        );
+    }
+    return parameters;
+}
+
+function getOverloads(implementation: ConvertibleSignatureDeclaration, declarations: readonly ConvertibleSignatureDeclaration[]): ConvertibleSignatureDeclaration[] {
+    const overloads: ConvertibleSignatureDeclaration[] = [];
+    for (const declaration of declarations) {
+        if (isOverload(declaration)) {
+            if (length(declaration.parameters) === length(implementation.parameters)) {
+                overloads.push(declaration);
+                continue;
+            }
+            if (length(declaration.parameters) > length(implementation.parameters)) {
+                return [];
+            }
+        }
+    }
+    return overloads;
+}
+
+function isOverload(declaration: ConvertibleSignatureDeclaration) {
+    return isConvertibleSignatureDeclaration(declaration) && declaration.body === undefined;
+}
+
+function createParameter(name: string, type: TypeNode, questionToken: QuestionToken | undefined) {
+    return factory.createParameterDeclaration(
+        /*modifiers*/ undefined,
+        /*dotDotDotToken*/ undefined,
+        name,
+        questionToken,
+        type,
+        /*initializer*/ undefined,
+    );
+}
+
+function isOptionalPos(declarations: ConvertibleSignatureDeclaration[], pos: number) {
+    return length(declarations) && some(declarations, d => pos < length(d.parameters) && !!d.parameters[pos] && d.parameters[pos].questionToken === undefined);
+}

tests/cases/fourslash/arityErrorAfterSignatureHelp.ts

@@ -18,5 +18,5 @@ verify.signatureHelp({
         kind: "retrigger"
     }
 })
-verify.not.codeFixAvailable() // trigger typecheck
+verify.not.codeFixAvailable(); // trigger typecheck
 verify.errorExistsBetweenMarkers("1", "2");

tests/cases/fourslash/codeFixAddMissingParam1.ts

@@ -0,0 +1,12 @@
+/// <reference path="fourslash.ts" />
+
+////[|function f() {}|]
+////
+////const a = 1;
+////f(a);
+
+verify.codeFix({
+    description: [ts.Diagnostics.Add_missing_parameter_to_0.message, "f"],
+    index: 0,
+    newRangeContent: "function f(a: number) {}"
+});