From e0312676020b641e7560d4c4f201e508b55349bd Mon Sep 17 00:00:00 2001 From: Skyler Rain Ross Date: Thu, 2 Dec 2021 09:25:17 -0800 Subject: [PATCH 1/2] fix: add return type works when there's missing whitespace --- .../src/handlers/add_return_type.rs | 77 +++++++++++++++---- 1 file changed, 61 insertions(+), 16 deletions(-) diff --git a/crates/ide_assists/src/handlers/add_return_type.rs b/crates/ide_assists/src/handlers/add_return_type.rs index 84983597eb8..d99443588ae 100644 --- a/crates/ide_assists/src/handlers/add_return_type.rs +++ b/crates/ide_assists/src/handlers/add_return_type.rs @@ -1,5 +1,5 @@ use hir::HirDisplay; -use syntax::{ast, AstNode, TextRange, TextSize}; +use syntax::{ast, AstNode, SyntaxKind, SyntaxToken, TextRange, TextSize}; use crate::{AssistContext, AssistId, AssistKind, Assists}; @@ -16,7 +16,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists}; // fn foo() -> i32 { 42i32 } // ``` pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { - let (fn_type, tail_expr, builder_edit_pos) = extract_tail(ctx)?; + let (fn_type, tail_expr, builder_edit_pos, needs_whitespace) = extract_tail(ctx)?; let module = ctx.sema.scope(tail_expr.syntax()).module()?; let ty = ctx.sema.type_of_expr(&tail_expr)?.adjusted(); if ty.is_unit() { @@ -32,12 +32,14 @@ pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option< }, tail_expr.syntax().text_range(), |builder| { + let preceeding_whitespace = if needs_whitespace { " " } else { "" }; + match builder_edit_pos { InsertOrReplace::Insert(insert_pos) => { - builder.insert(insert_pos, &format!("-> {} ", ty)) + builder.insert(insert_pos, &format!("{}-> {} ", preceeding_whitespace, ty)) } InsertOrReplace::Replace(text_range) => { - builder.replace(text_range, &format!("-> {}", ty)) + builder.replace(text_range, &format!("{}-> {}", preceeding_whitespace, ty)) } } if let FnType::Closure { wrap_expr: true } = fn_type { @@ -56,13 +58,16 @@ enum InsertOrReplace { /// 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, insert_pos: TextSize) -> Option { +fn ret_ty_to_action( + ret_ty: Option, + insert_after: SyntaxToken, +) -> Option<(InsertOrReplace, bool)> { match ret_ty { Some(ret_ty) => match ret_ty.ty() { Some(ast::Type::InferType(_)) | None => { cov_mark::hit!(existing_infer_ret_type); cov_mark::hit!(existing_infer_ret_type_closure); - Some(InsertOrReplace::Replace(ret_ty.syntax().text_range())) + Some((InsertOrReplace::Replace(ret_ty.syntax().text_range()), false)) } _ => { cov_mark::hit!(existing_ret_type); @@ -70,7 +75,17 @@ fn ret_ty_to_action(ret_ty: Option, insert_pos: TextSize) -> Optio 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)) + } } } @@ -79,11 +94,13 @@ enum FnType { Closure { wrap_expr: bool }, } -fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrReplace)> { - let (fn_type, tail_expr, return_type_range, action) = +fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrReplace, bool)> { + let (fn_type, tail_expr, return_type_range, action, needs_whitespace) = if let Some(closure) = ctx.find_node_at_offset::() { - 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, needs_whitespace) = ret_ty_to_action(closure.ret_type(), rpipe)?; let body = closure.body()?; let body_start = body.syntax().first_token()?.text_range().start(); @@ -93,11 +110,13 @@ fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrRepla }; let ret_range = TextRange::new(rpipe_pos, body_start); - (FnType::Closure { wrap_expr }, tail_expr, ret_range, action) + (FnType::Closure { wrap_expr }, tail_expr, ret_range, action, needs_whitespace) } else { let func = ctx.find_node_at_offset::()?; - 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, needs_whitespace) = ret_ty_to_action(func.ret_type(), rparen)?; let body = func.body()?; let stmt_list = body.stmt_list()?; @@ -105,7 +124,7 @@ fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrRepla let ret_range_end = stmt_list.l_curly_token()?.text_range().start(); let ret_range = TextRange::new(rparen_pos, ret_range_end); - (FnType::Function, tail_expr, ret_range, action) + (FnType::Function, tail_expr, ret_range, action, needs_whitespace) }; let range = ctx.selection_trimmed(); if return_type_range.contains_range(range) { @@ -117,7 +136,7 @@ fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrRepla } else { return None; } - Some((fn_type, tail_expr, action)) + Some((fn_type, tail_expr, action, needs_whitespace)) } #[cfg(test)] @@ -196,6 +215,19 @@ mod tests { ); } + #[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( @@ -280,6 +312,19 @@ mod tests { ); } + #[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); From 5b59a5eca8867ba72b6b8838930c75a65d00f13c Mon Sep 17 00:00:00 2001 From: Skyler Rain Ross Date: Thu, 2 Dec 2021 10:46:07 -0800 Subject: [PATCH 2/2] refactor(assist/add_return_type): avoid threading `needs_whitespace` --- .../src/handlers/add_return_type.rs | 31 +++++++++---------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/crates/ide_assists/src/handlers/add_return_type.rs b/crates/ide_assists/src/handlers/add_return_type.rs index d99443588ae..2c5b61eddb7 100644 --- a/crates/ide_assists/src/handlers/add_return_type.rs +++ b/crates/ide_assists/src/handlers/add_return_type.rs @@ -16,7 +16,7 @@ use crate::{AssistContext, AssistId, AssistKind, Assists}; // fn foo() -> i32 { 42i32 } // ``` pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option<()> { - let (fn_type, tail_expr, builder_edit_pos, needs_whitespace) = extract_tail(ctx)?; + 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(); if ty.is_unit() { @@ -32,14 +32,13 @@ pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option< }, tail_expr.syntax().text_range(), |builder| { - let preceeding_whitespace = if needs_whitespace { " " } else { "" }; - match builder_edit_pos { - InsertOrReplace::Insert(insert_pos) => { + 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!("{}-> {}", preceeding_whitespace, ty)) + builder.replace(text_range, &format!("-> {}", ty)) } } if let FnType::Closure { wrap_expr: true } = fn_type { @@ -52,7 +51,7 @@ pub(crate) fn add_return_type(acc: &mut Assists, ctx: &AssistContext) -> Option< } enum InsertOrReplace { - Insert(TextSize), + Insert(TextSize, bool), Replace(TextRange), } @@ -61,13 +60,13 @@ enum InsertOrReplace { fn ret_ty_to_action( ret_ty: Option, insert_after: SyntaxToken, -) -> Option<(InsertOrReplace, bool)> { +) -> Option { match ret_ty { Some(ret_ty) => match ret_ty.ty() { Some(ast::Type::InferType(_)) | None => { cov_mark::hit!(existing_infer_ret_type); cov_mark::hit!(existing_infer_ret_type_closure); - Some((InsertOrReplace::Replace(ret_ty.syntax().text_range()), false)) + Some(InsertOrReplace::Replace(ret_ty.syntax().text_range())) } _ => { cov_mark::hit!(existing_ret_type); @@ -84,7 +83,7 @@ fn ret_ty_to_action( _ => (insert_after_pos, true), }; - Some((InsertOrReplace::Insert(insert_pos), needs_whitespace)) + Some(InsertOrReplace::Insert(insert_pos, needs_whitespace)) } } } @@ -94,13 +93,13 @@ enum FnType { Closure { wrap_expr: bool }, } -fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrReplace, bool)> { - let (fn_type, tail_expr, return_type_range, action, needs_whitespace) = +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::() { let rpipe = closure.param_list()?.syntax().last_token()?; let rpipe_pos = rpipe.text_range().end(); - let (action, needs_whitespace) = ret_ty_to_action(closure.ret_type(), rpipe)?; + let action = ret_ty_to_action(closure.ret_type(), rpipe)?; let body = closure.body()?; let body_start = body.syntax().first_token()?.text_range().start(); @@ -110,13 +109,13 @@ fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrRepla }; let ret_range = TextRange::new(rpipe_pos, body_start); - (FnType::Closure { wrap_expr }, tail_expr, ret_range, action, needs_whitespace) + (FnType::Closure { wrap_expr }, tail_expr, ret_range, action) } else { let func = ctx.find_node_at_offset::()?; let rparen = func.param_list()?.r_paren_token()?; let rparen_pos = rparen.text_range().end(); - let (action, needs_whitespace) = ret_ty_to_action(func.ret_type(), rparen)?; + let action = ret_ty_to_action(func.ret_type(), rparen)?; let body = func.body()?; let stmt_list = body.stmt_list()?; @@ -124,7 +123,7 @@ fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrRepla let ret_range_end = stmt_list.l_curly_token()?.text_range().start(); let ret_range = TextRange::new(rparen_pos, ret_range_end); - (FnType::Function, tail_expr, ret_range, action, needs_whitespace) + (FnType::Function, tail_expr, ret_range, action) }; let range = ctx.selection_trimmed(); if return_type_range.contains_range(range) { @@ -136,7 +135,7 @@ fn extract_tail(ctx: &AssistContext) -> Option<(FnType, ast::Expr, InsertOrRepla } else { return None; } - Some((fn_type, tail_expr, action, needs_whitespace)) + Some((fn_type, tail_expr, action)) } #[cfg(test)]