• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1/* @internal */
2namespace ts.codefix {
3    const fixId = "returnValueCorrect";
4    const fixIdAddReturnStatement = "fixAddReturnStatement";
5    const fixRemoveBracesFromArrowFunctionBody = "fixRemoveBracesFromArrowFunctionBody";
6    const fixIdWrapTheBlockWithParen = "fixWrapTheBlockWithParen";
7    const errorCodes = [
8        Diagnostics.A_function_whose_declared_type_is_neither_void_nor_any_must_return_a_value.code,
9        Diagnostics.Type_0_is_not_assignable_to_type_1.code,
10        Diagnostics.Argument_of_type_0_is_not_assignable_to_parameter_of_type_1.code
11    ];
12
13    enum ProblemKind {
14        MissingReturnStatement,
15        MissingParentheses
16    }
17
18    interface MissingReturnInfo {
19        kind: ProblemKind.MissingReturnStatement;
20        declaration: FunctionLikeDeclaration;
21        expression: Expression;
22        statement: Statement;
23        commentSource: Node;
24    }
25
26    interface MissingParenInfo {
27        kind: ProblemKind.MissingParentheses;
28        declaration: ArrowFunction;
29        expression: Expression;
30        statement: Statement;
31        commentSource: Node;
32    }
33
34    type Info = MissingReturnInfo | MissingParenInfo;
35
36    registerCodeFix({
37        errorCodes,
38        fixIds: [fixIdAddReturnStatement, fixRemoveBracesFromArrowFunctionBody, fixIdWrapTheBlockWithParen],
39        getCodeActions: function getCodeActionsToCorrectReturnValue(context) {
40            const { program, sourceFile, span: { start }, errorCode } = context;
41            const info = getInfo(program.getTypeChecker(), sourceFile, start, errorCode);
42            if (!info) return undefined;
43
44            if (info.kind === ProblemKind.MissingReturnStatement) {
45                return append(
46                    [getActionForfixAddReturnStatement(context, info.expression, info.statement)],
47                    isArrowFunction(info.declaration) ? getActionForFixRemoveBracesFromArrowFunctionBody(context, info.declaration, info.expression, info.commentSource): undefined);
48            }
49            else {
50                return [getActionForfixWrapTheBlockWithParen(context, info.declaration, info.expression)];
51            }
52        },
53        getAllCodeActions: context => codeFixAll(context, errorCodes, (changes, diag) => {
54            const info = getInfo(context.program.getTypeChecker(), diag.file, diag.start, diag.code);
55            if (!info) return undefined;
56
57            switch (context.fixId) {
58                case fixIdAddReturnStatement:
59                    addReturnStatement(changes, diag.file, info.expression, info.statement);
60                    break;
61                case fixRemoveBracesFromArrowFunctionBody:
62                    if (!isArrowFunction(info.declaration)) return undefined;
63                    removeBlockBodyBrace(changes, diag.file, info.declaration, info.expression, info.commentSource, /* withParen */ false);
64                    break;
65                case fixIdWrapTheBlockWithParen:
66                    if (!isArrowFunction(info.declaration)) return undefined;
67                    wrapBlockWithParen(changes, diag.file, info.declaration, info.expression);
68                    break;
69                default:
70                    Debug.fail(JSON.stringify(context.fixId));
71            }
72        }),
73    });
74
75    function createObjectTypeFromLabeledExpression(checker: TypeChecker, label: Identifier, expression: Expression) {
76        const member = checker.createSymbol(SymbolFlags.Property, label.escapedText);
77        member.type = checker.getTypeAtLocation(expression);
78        const members = createSymbolTable([member]);
79        return checker.createAnonymousType(/*symbol*/ undefined, members, [], [], []);
80    }
81
82    function getFixInfo(checker: TypeChecker, declaration: FunctionLikeDeclaration, expectType: Type, isFunctionType: boolean): Info | undefined {
83        if (!declaration.body || !isBlock(declaration.body) || length(declaration.body.statements) !== 1) return undefined;
84
85        const firstStatement = first(declaration.body.statements);
86        if (isExpressionStatement(firstStatement) && checkFixedAssignableTo(checker, declaration, checker.getTypeAtLocation(firstStatement.expression), expectType, isFunctionType)) {
87            return {
88                declaration,
89                kind: ProblemKind.MissingReturnStatement,
90                expression: firstStatement.expression,
91                statement: firstStatement,
92                commentSource: firstStatement.expression
93            };
94        }
95        else if (isLabeledStatement(firstStatement) && isExpressionStatement(firstStatement.statement)) {
96            const node = factory.createObjectLiteralExpression([factory.createPropertyAssignment(firstStatement.label, firstStatement.statement.expression)]);
97            const nodeType = createObjectTypeFromLabeledExpression(checker, firstStatement.label, firstStatement.statement.expression);
98            if (checkFixedAssignableTo(checker, declaration, nodeType, expectType, isFunctionType)) {
99                return isArrowFunction(declaration) ? {
100                    declaration,
101                    kind: ProblemKind.MissingParentheses,
102                    expression: node,
103                    statement: firstStatement,
104                    commentSource: firstStatement.statement.expression
105                } : {
106                        declaration,
107                        kind: ProblemKind.MissingReturnStatement,
108                        expression: node,
109                        statement: firstStatement,
110                        commentSource: firstStatement.statement.expression
111                    };
112            }
113        }
114        else if (isBlock(firstStatement) && length(firstStatement.statements) === 1) {
115            const firstBlockStatement = first(firstStatement.statements);
116            if (isLabeledStatement(firstBlockStatement) && isExpressionStatement(firstBlockStatement.statement)) {
117                const node = factory.createObjectLiteralExpression([factory.createPropertyAssignment(firstBlockStatement.label, firstBlockStatement.statement.expression)]);
118                const nodeType = createObjectTypeFromLabeledExpression(checker, firstBlockStatement.label, firstBlockStatement.statement.expression);
119                if (checkFixedAssignableTo(checker, declaration, nodeType, expectType, isFunctionType)) {
120                    return {
121                        declaration,
122                        kind: ProblemKind.MissingReturnStatement,
123                        expression: node,
124                        statement: firstStatement,
125                        commentSource: firstBlockStatement
126                    };
127                }
128            }
129        }
130
131        return undefined;
132    }
133
134    function checkFixedAssignableTo(checker: TypeChecker, declaration: FunctionLikeDeclaration, exprType: Type, type: Type, isFunctionType: boolean) {
135        if (isFunctionType) {
136            const sig = checker.getSignatureFromDeclaration(declaration);
137            if (sig) {
138                if (hasSyntacticModifier(declaration, ModifierFlags.Async)) {
139                    exprType = checker.createPromiseType(exprType);
140                }
141                const newSig = checker.createSignature(
142                    declaration,
143                    sig.typeParameters,
144                    sig.thisParameter,
145                    sig.parameters,
146                    exprType,
147                    /*typePredicate*/ undefined,
148                    sig.minArgumentCount,
149                    sig.flags);
150                exprType = checker.createAnonymousType(
151                    /*symbol*/ undefined,
152                    createSymbolTable(),
153                    [newSig],
154                    [],
155                    []);
156            }
157            else {
158                exprType = checker.getAnyType();
159            }
160        }
161        return checker.isTypeAssignableTo(exprType, type);
162    }
163
164    function getInfo(checker: TypeChecker, sourceFile: SourceFile, position: number, errorCode: number): Info | undefined {
165        const node = getTokenAtPosition(sourceFile, position);
166        if (!node.parent) return undefined;
167
168        const declaration = findAncestor(node.parent, isFunctionLikeDeclaration);
169        switch (errorCode) {
170            case Diagnostics.A_function_whose_declared_type_is_neither_void_nor_any_must_return_a_value.code:
171                if (!declaration || !declaration.body || !declaration.type || !rangeContainsRange(declaration.type, node)) return undefined;
172                return getFixInfo(checker, declaration, checker.getTypeFromTypeNode(declaration.type), /* isFunctionType */ false);
173            case Diagnostics.Argument_of_type_0_is_not_assignable_to_parameter_of_type_1.code:
174                if (!declaration || !isCallExpression(declaration.parent) || !declaration.body) return undefined;
175                const pos = declaration.parent.arguments.indexOf(declaration as Expression);
176                const type = checker.getContextualTypeForArgumentAtIndex(declaration.parent, pos);
177                if (!type) return undefined;
178                return getFixInfo(checker, declaration, type, /* isFunctionType */ true);
179            case Diagnostics.Type_0_is_not_assignable_to_type_1.code:
180                if (!isDeclarationName(node) || !isVariableLike(node.parent) && !isJsxAttribute(node.parent)) return undefined;
181                const initializer = getVariableLikeInitializer(node.parent);
182                if (!initializer || !isFunctionLikeDeclaration(initializer) || !initializer.body) return undefined;
183                return getFixInfo(checker, initializer, checker.getTypeAtLocation(node.parent), /* isFunctionType */ true);
184        }
185        return undefined;
186    }
187
188    function getVariableLikeInitializer(declaration: VariableLikeDeclaration): Expression | undefined {
189        switch (declaration.kind) {
190            case SyntaxKind.VariableDeclaration:
191            case SyntaxKind.Parameter:
192            case SyntaxKind.BindingElement:
193            case SyntaxKind.PropertyDeclaration:
194            case SyntaxKind.PropertyAssignment:
195                return declaration.initializer;
196            case SyntaxKind.JsxAttribute:
197                return declaration.initializer && (isJsxExpression(declaration.initializer) ? declaration.initializer.expression : undefined);
198            case SyntaxKind.ShorthandPropertyAssignment:
199            case SyntaxKind.PropertySignature:
200            case SyntaxKind.EnumMember:
201            case SyntaxKind.JSDocPropertyTag:
202            case SyntaxKind.JSDocParameterTag:
203                return undefined;
204        }
205    }
206
207    function addReturnStatement(changes: textChanges.ChangeTracker, sourceFile: SourceFile, expression: Expression, statement: Statement) {
208        suppressLeadingAndTrailingTrivia(expression);
209        const probablyNeedSemi = probablyUsesSemicolons(sourceFile);
210        changes.replaceNode(sourceFile, statement, factory.createReturnStatement(expression), {
211            leadingTriviaOption: textChanges.LeadingTriviaOption.Exclude,
212            trailingTriviaOption: textChanges.TrailingTriviaOption.Exclude,
213            suffix: probablyNeedSemi ? ";" : undefined
214        });
215    }
216
217    function removeBlockBodyBrace(changes: textChanges.ChangeTracker, sourceFile: SourceFile, declaration: ArrowFunction, expression: Expression, commentSource: Node, withParen: boolean) {
218        const newBody = (withParen || needsParentheses(expression)) ? factory.createParenthesizedExpression(expression) : expression;
219        suppressLeadingAndTrailingTrivia(commentSource);
220        copyComments(commentSource, newBody);
221
222        changes.replaceNode(sourceFile, declaration.body, newBody);
223    }
224
225    function wrapBlockWithParen(changes: textChanges.ChangeTracker, sourceFile: SourceFile, declaration: ArrowFunction, expression: Expression) {
226        changes.replaceNode(sourceFile, declaration.body, factory.createParenthesizedExpression(expression));
227    }
228
229    function getActionForfixAddReturnStatement(context: CodeFixContext, expression: Expression, statement: Statement) {
230        const changes = textChanges.ChangeTracker.with(context, t => addReturnStatement(t, context.sourceFile, expression, statement));
231        return createCodeFixAction(fixId, changes, Diagnostics.Add_a_return_statement, fixIdAddReturnStatement, Diagnostics.Add_all_missing_return_statement);
232    }
233
234    function getActionForFixRemoveBracesFromArrowFunctionBody(context: CodeFixContext, declaration: ArrowFunction, expression: Expression, commentSource: Node) {
235        const changes = textChanges.ChangeTracker.with(context, t => removeBlockBodyBrace(t, context.sourceFile, declaration, expression, commentSource, /* withParen */ false));
236        return createCodeFixAction(fixId, changes, Diagnostics.Remove_braces_from_arrow_function_body, fixRemoveBracesFromArrowFunctionBody, Diagnostics.Remove_braces_from_all_arrow_function_bodies_with_relevant_issues);
237    }
238
239    function getActionForfixWrapTheBlockWithParen(context: CodeFixContext, declaration: ArrowFunction, expression: Expression) {
240        const changes = textChanges.ChangeTracker.with(context, t => wrapBlockWithParen(t, context.sourceFile, declaration, expression));
241        return createCodeFixAction(fixId, changes, Diagnostics.Wrap_the_following_body_with_parentheses_which_should_be_an_object_literal, fixIdWrapTheBlockWithParen, Diagnostics.Wrap_all_object_literal_with_parentheses);
242    }
243}
244