use hir::HirDisplay;
-use syntax::{ast, AstNode, TextRange, TextSize};
+use syntax::{ast, match_ast, AstNode, SyntaxKind, SyntaxToken, TextRange, TextSize};
use crate::{AssistContext, AssistId, AssistKind, Assists};
pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
let (fn_type, tail_expr, builder_edit_pos) = extract_tail(ctx)?;
let module = ctx.sema.scope(tail_expr.syntax()).module()?;
- let ty = ctx.sema.type_of_expr(&tail_expr)?.adjusted();
+ let ty = ctx.sema.type_of_expr(&peel_blocks(tail_expr.clone()))?.original();
if ty.is_unit() {
return None;
}
tail_expr.syntax().text_range(),
|builder| {
match builder_edit_pos {
- InsertOrReplace::Insert(insert_pos) => {
- builder.insert(insert_pos, &format!("-> {} ", ty))
+ InsertOrReplace::Insert(insert_pos, needs_whitespace) => {
+ let preceeding_whitespace = if needs_whitespace { " " } else { "" };
+ builder.insert(insert_pos, &format!("{}-> {} ", preceeding_whitespace, ty))
}
InsertOrReplace::Replace(text_range) => {
builder.replace(text_range, &format!("-> {}", ty))
}
enum InsertOrReplace {
- Insert(TextSize),
+ Insert(TextSize, bool),
Replace(TextRange),
}
/// Check the potentially already specified return type and reject it or turn it into a builder command
/// if allowed.
-fn ret_ty_to_action(ret_ty: Option<ast::RetType>, insert_pos: TextSize) -> Option<InsertOrReplace> {
+fn ret_ty_to_action(
+ ret_ty: Option<ast::RetType>,
+ insert_after: SyntaxToken,
+) -> Option<InsertOrReplace> {
match ret_ty {
Some(ret_ty) => match ret_ty.ty() {
Some(ast::Type::InferType(_)) | None => {
None
}
},
- None => Some(InsertOrReplace::Insert(insert_pos + TextSize::from(1))),
+ None => {
+ let insert_after_pos = insert_after.text_range().end();
+ let (insert_pos, needs_whitespace) = match insert_after.next_token() {
+ Some(it) if it.kind() == SyntaxKind::WHITESPACE => {
+ (insert_after_pos + TextSize::from(1), false)
+ }
+ _ => (insert_after_pos, true),
+ };
+
+ Some(InsertOrReplace::Insert(insert_pos, needs_whitespace))
+ }
}
}
Closure { wrap_expr: bool },
}
+/// If we're looking at a block that is supposed to return `()`, type inference
+/// will just tell us it has type `()`. We have to look at the tail expression
+/// to see the mismatched actual type. This 'unpeels' the various blocks to
+/// hopefully let us see the type the user intends. (This still doesn't handle
+/// all situations fully correctly; the 'ideal' way to handle this would be to
+/// run type inference on the function again, but with a variable as the return
+/// type.)
+fn peel_blocks(mut expr: ast::Expr) -> ast::Expr {
+ loop {
+ match_ast! {
+ match (expr.syntax()) {
+ ast::BlockExpr(it) => {
+ if let Some(tail) = it.tail_expr() {
+ expr = tail.clone();
+ } else {
+ break;
+ }
+ },
+ ast::IfExpr(it) => {
+ if let Some(then_branch) = it.then_branch() {
+ expr = ast::Expr::BlockExpr(then_branch.clone());
+ } else {
+ break;
+ }
+ },
+ ast::MatchExpr(it) => {
+ if let Some(arm_expr) = it.match_arm_list().and_then(|l| l.arms().next()).and_then(|a| a.expr()) {
+ expr = arm_expr;
+ } else {
+ break;
+ }
+ },
+ _ => break,
+ }
+ }
+ }
+ expr
+}
+
fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrReplace)> {
let (fn_type, tail_expr, return_type_range, action) =
if let Some(closure) = ctx.find_node_at_offset::<ast::ClosureExpr>() {
- let rpipe_pos = closure.param_list()?.syntax().last_token()?.text_range().end();
- let action = ret_ty_to_action(closure.ret_type(), rpipe_pos)?;
+ let rpipe = closure.param_list()?.syntax().last_token()?;
+ let rpipe_pos = rpipe.text_range().end();
+
+ let action = ret_ty_to_action(closure.ret_type(), rpipe)?;
let body = closure.body()?;
let body_start = body.syntax().first_token()?.text_range().start();
(FnType::Closure { wrap_expr }, tail_expr, ret_range, action)
} else {
let func = ctx.find_node_at_offset::<ast::Fn>()?;
- let rparen_pos = func.param_list()?.r_paren_token()?.text_range().end();
- let action = ret_ty_to_action(func.ret_type(), rparen_pos)?;
+
+ let rparen = func.param_list()?.r_paren_token()?;
+ let rparen_pos = rparen.text_range().end();
+ let action = ret_ty_to_action(func.ret_type(), rparen)?;
let body = func.body()?;
let stmt_list = body.stmt_list()?;
);
}
+ #[test]
+ fn infer_return_type_no_whitespace() {
+ check_assist(
+ add_return_type,
+ r#"fn foo(){
+ 45$0
+}"#,
+ r#"fn foo() -> i32 {
+ 45
+}"#,
+ );
+ }
+
#[test]
fn infer_return_type_nested() {
check_assist(
);
}
+ #[test]
+ fn infer_return_type_nested_match() {
+ check_assist(
+ add_return_type,
+ r#"fn foo() {
+ match true {
+ true => { 3$0 },
+ false => { 5 },
+ }
+}"#,
+ r#"fn foo() -> i32 {
+ match true {
+ true => { 3 },
+ false => { 5 },
+ }
+}"#,
+ );
+ }
+
#[test]
fn not_applicable_ret_type_specified() {
cov_mark::check!(existing_ret_type);
);
}
+ #[test]
+ fn infer_return_type_closure_no_whitespace() {
+ check_assist(
+ add_return_type,
+ r#"fn foo() {
+ |x: i32|{ x$0 };
+}"#,
+ r#"fn foo() {
+ |x: i32| -> i32 { x };
+}"#,
+ );
+ }
+
#[test]
fn infer_return_type_closure_wrap() {
cov_mark::check!(wrap_closure_non_block_expr);
projects / rust.git / blobdiff