diff --git a/crates/ide_assists/src/handlers/unwrap_result_return_type.rs b/crates/ide_assists/src/handlers/unwrap_result_return_type.rs index 6d813aa672..82499d77c3 100644 --- a/crates/ide_assists/src/handlers/unwrap_result_return_type.rs +++ b/crates/ide_assists/src/handlers/unwrap_result_return_type.rs @@ -1,7 +1,8 @@ use ide_db::helpers::{for_each_tail_expr, node_ext::walk_expr, FamousDefs}; +use itertools::Itertools; use syntax::{ ast::{self, Expr}, - match_ast, AstNode, + match_ast, AstNode, TextRange, TextSize, }; use crate::{AssistContext, AssistId, AssistKind, Assists}; @@ -60,25 +61,55 @@ pub(crate) fn unwrap_result_return_type(acc: &mut Assists, ctx: &AssistContext) }); for_each_tail_expr(&body, tail_cb); - for ret_expr_arg in exprs_to_unwrap { - let new_ret_expr = ret_expr_arg.to_string(); - let new_ret_expr = - new_ret_expr.trim_start_matches("Ok(").trim_start_matches("Err("); - builder.replace( - ret_expr_arg.syntax().text_range(), - new_ret_expr.strip_suffix(')').unwrap_or(new_ret_expr), - ) - } - + let mut is_unit_type = false; if let Some((_, inner_type)) = type_ref.to_string().split_once('<') { let inner_type = match inner_type.split_once(',') { Some((success_inner_type, _)) => success_inner_type, None => inner_type, }; - builder.replace( - type_ref.syntax().text_range(), - inner_type.strip_suffix('>').unwrap_or(inner_type), - ) + let new_ret_type = inner_type.strip_suffix('>').unwrap_or(inner_type); + if new_ret_type == "()" { + is_unit_type = true; + let text_range = TextRange::new( + ret_type.syntax().text_range().start(), + ret_type.syntax().text_range().end() + TextSize::from(1u32), + ); + builder.delete(text_range) + } else { + builder.replace( + type_ref.syntax().text_range(), + inner_type.strip_suffix('>').unwrap_or(inner_type), + ) + } + } + + for ret_expr_arg in exprs_to_unwrap { + let ret_expr_str = ret_expr_arg.to_string(); + if ret_expr_str.starts_with("Ok(") || ret_expr_str.starts_with("Err(") { + let arg_list = ret_expr_arg.syntax().children().find_map(ast::ArgList::cast); + if let Some(arg_list) = arg_list { + if is_unit_type { + match ret_expr_arg.syntax().prev_sibling_or_token() { + // Useful to delete the entire line without leaving trailing whitespaces + Some(whitespace) => { + let new_range = TextRange::new( + whitespace.text_range().start(), + ret_expr_arg.syntax().text_range().end(), + ); + builder.delete(new_range); + } + None => { + builder.delete(ret_expr_arg.syntax().text_range()); + } + } + } else { + builder.replace( + ret_expr_arg.syntax().text_range(), + arg_list.args().join(", "), + ); + } + } + } } }, ) @@ -126,6 +157,49 @@ fn foo() -> i32 { ); } + #[test] + fn unwrap_result_return_type_unit_type() { + check_assist( + unwrap_result_return_type, + r#" +//- minicore: result +fn foo() -> Result<(), Box> { + Ok(()) +} +"#, + r#" +fn foo() { +} +"#, + ); + } + + #[test] + fn unwrap_result_return_type_ending_with_parent() { + check_assist( + unwrap_result_return_type, + r#" +//- minicore: result +fn foo() -> Result> { + if true { + Ok(42) + } else { + foo() + } +} +"#, + r#" +fn foo() -> i32 { + if true { + 42 + } else { + foo() + } +} +"#, + ); + } + #[test] fn unwrap_return_type_break_split_tail() { check_assist(