2 use syntax::{ast, AstNode, SyntaxKind, SyntaxToken, TextRange, TextSize};
4 use crate::{AssistContext, AssistId, AssistKind, Assists};
6 // Assist: add_return_type
8 // Adds the return type to a function or closure inferred from its tail expression if it doesn't have a return
9 // type specified. This assists is useable in a functions or closures tail expression or return type position.
12 // fn foo() { 4$02i32 }
16 // fn foo() -> i32 { 42i32 }
18 pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
19 let (fn_type, tail_expr, builder_edit_pos) = extract_tail(ctx)?;
20 let module = ctx.sema.scope(tail_expr.syntax()).module()?;
21 let ty = ctx.sema.type_of_expr(&tail_expr)?.adjusted();
25 let ty = ty.display_source_code(ctx.db(), module.into()).ok()?;
28 AssistId("add_return_type", AssistKind::RefactorRewrite),
30 FnType::Function => "Add this function's return type",
31 FnType::Closure { .. } => "Add this closure's return type",
33 tail_expr.syntax().text_range(),
35 match builder_edit_pos {
36 InsertOrReplace::Insert(insert_pos, needs_whitespace) => {
37 let preceeding_whitespace = if needs_whitespace { " " } else { "" };
38 builder.insert(insert_pos, &format!("{}-> {} ", preceeding_whitespace, ty))
40 InsertOrReplace::Replace(text_range) => {
41 builder.replace(text_range, &format!("-> {}", ty))
44 if let FnType::Closure { wrap_expr: true } = fn_type {
45 cov_mark::hit!(wrap_closure_non_block_expr);
46 // `|x| x` becomes `|x| -> T x` which is invalid, so wrap it in a block
47 builder.replace(tail_expr.syntax().text_range(), &format!("{{{}}}", tail_expr));
53 enum InsertOrReplace {
54 Insert(TextSize, bool),
58 /// Check the potentially already specified return type and reject it or turn it into a builder command
61 ret_ty: Option<ast::RetType>,
62 insert_after: SyntaxToken,
63 ) -> Option<InsertOrReplace> {
65 Some(ret_ty) => match ret_ty.ty() {
66 Some(ast::Type::InferType(_)) | None => {
67 cov_mark::hit!(existing_infer_ret_type);
68 cov_mark::hit!(existing_infer_ret_type_closure);
69 Some(InsertOrReplace::Replace(ret_ty.syntax().text_range()))
72 cov_mark::hit!(existing_ret_type);
73 cov_mark::hit!(existing_ret_type_closure);
78 let insert_after_pos = insert_after.text_range().end();
79 let (insert_pos, needs_whitespace) = match insert_after.next_token() {
80 Some(it) if it.kind() == SyntaxKind::WHITESPACE => {
81 (insert_after_pos + TextSize::from(1), false)
83 _ => (insert_after_pos, true),
86 Some(InsertOrReplace::Insert(insert_pos, needs_whitespace))
93 Closure { wrap_expr: bool },
96 fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrReplace)> {
97 let (fn_type, tail_expr, return_type_range, action) =
98 if let Some(closure) = ctx.find_node_at_offset::<ast::ClosureExpr>() {
99 let rpipe = closure.param_list()?.syntax().last_token()?;
100 let rpipe_pos = rpipe.text_range().end();
102 let action = ret_ty_to_action(closure.ret_type(), rpipe)?;
104 let body = closure.body()?;
105 let body_start = body.syntax().first_token()?.text_range().start();
106 let (tail_expr, wrap_expr) = match body {
107 ast::Expr::BlockExpr(block) => (block.tail_expr()?, false),
108 body => (body, true),
111 let ret_range = TextRange::new(rpipe_pos, body_start);
112 (FnType::Closure { wrap_expr }, tail_expr, ret_range, action)
114 let func = ctx.find_node_at_offset::<ast::Fn>()?;
116 let rparen = func.param_list()?.r_paren_token()?;
117 let rparen_pos = rparen.text_range().end();
118 let action = ret_ty_to_action(func.ret_type(), rparen)?;
120 let body = func.body()?;
121 let stmt_list = body.stmt_list()?;
122 let tail_expr = stmt_list.tail_expr()?;
124 let ret_range_end = stmt_list.l_curly_token()?.text_range().start();
125 let ret_range = TextRange::new(rparen_pos, ret_range_end);
126 (FnType::Function, tail_expr, ret_range, action)
128 let range = ctx.selection_trimmed();
129 if return_type_range.contains_range(range) {
130 cov_mark::hit!(cursor_in_ret_position);
131 cov_mark::hit!(cursor_in_ret_position_closure);
132 } else if tail_expr.syntax().text_range().contains_range(range) {
133 cov_mark::hit!(cursor_on_tail);
134 cov_mark::hit!(cursor_on_tail_closure);
138 Some((fn_type, tail_expr, action))
143 use crate::tests::{check_assist, check_assist_not_applicable};
148 fn infer_return_type_specified_inferred() {
149 cov_mark::check!(existing_infer_ret_type);
162 fn infer_return_type_specified_inferred_closure() {
163 cov_mark::check!(existing_infer_ret_type_closure);
176 fn infer_return_type_cursor_at_return_type_pos() {
177 cov_mark::check!(cursor_in_ret_position);
190 fn infer_return_type_cursor_at_return_type_pos_closure() {
191 cov_mark::check!(cursor_in_ret_position_closure);
204 fn infer_return_type() {
205 cov_mark::check!(cursor_on_tail);
218 fn infer_return_type_no_whitespace() {
231 fn infer_return_type_nested() {
252 fn not_applicable_ret_type_specified() {
253 cov_mark::check!(existing_ret_type);
254 check_assist_not_applicable(
263 fn not_applicable_non_tail_expr() {
264 check_assist_not_applicable(
274 fn not_applicable_unit_return_type() {
275 check_assist_not_applicable(
284 fn infer_return_type_closure_block() {
285 cov_mark::check!(cursor_on_tail_closure);
302 fn infer_return_type_closure() {
309 |x: i32| -> i32 { x };
315 fn infer_return_type_closure_no_whitespace() {
322 |x: i32| -> i32 { x };
328 fn infer_return_type_closure_wrap() {
329 cov_mark::check!(wrap_closure_non_block_expr);
342 fn infer_return_type_nested_closure() {
367 fn not_applicable_ret_type_specified_closure() {
368 cov_mark::check!(existing_ret_type_closure);
369 check_assist_not_applicable(
378 fn not_applicable_non_tail_expr_closure() {
379 check_assist_not_applicable(