From 90b181c8e137d10991733064fd5cefef98658ffa Mon Sep 17 00:00:00 2001 From: Giga Bowser <45986823+Giga-Bowser@users.noreply.github.com> Date: Sun, 13 Oct 2024 10:53:38 -0400 Subject: [PATCH] Add wrap/unwrap return type in Option --- .../src/handlers/unwrap_option_return_type.rs | 1143 ++++++++++++++++ .../handlers/wrap_return_type_in_option.rs | 1173 +++++++++++++++++ .../crates/ide-assists/src/lib.rs | 4 + .../crates/ide-assists/src/tests/generated.rs | 28 + 4 files changed, 2348 insertions(+) create mode 100644 src/tools/rust-analyzer/crates/ide-assists/src/handlers/unwrap_option_return_type.rs create mode 100644 src/tools/rust-analyzer/crates/ide-assists/src/handlers/wrap_return_type_in_option.rs diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unwrap_option_return_type.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unwrap_option_return_type.rs new file mode 100644 index 0000000000000..8dc40842406ba --- /dev/null +++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/unwrap_option_return_type.rs @@ -0,0 +1,1143 @@ +use ide_db::{ + famous_defs::FamousDefs, + syntax_helpers::node_ext::{for_each_tail_expr, walk_expr}, +}; +use itertools::Itertools; +use syntax::{ + ast::{self, Expr, HasGenericArgs}, + match_ast, AstNode, NodeOrToken, SyntaxKind, TextRange, +}; + +use crate::{AssistContext, AssistId, AssistKind, Assists}; + +// Assist: unwrap_option_return_type +// +// Unwrap the function's return type. +// +// ``` +// # //- minicore: option +// fn foo() -> Option$0 { Some(42i32) } +// ``` +// -> +// ``` +// fn foo() -> i32 { 42i32 } +// ``` +pub(crate) fn unwrap_option_return_type(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { + let ret_type = ctx.find_node_at_offset::()?; + let parent = ret_type.syntax().parent()?; + let body = match_ast! { + match parent { + ast::Fn(func) => func.body()?, + ast::ClosureExpr(closure) => match closure.body()? { + Expr::BlockExpr(block) => block, + // closures require a block when a return type is specified + _ => return None, + }, + _ => return None, + } + }; + + let type_ref = &ret_type.ty()?; + let Some(hir::Adt::Enum(ret_enum)) = ctx.sema.resolve_type(type_ref)?.as_adt() else { + return None; + }; + let option_enum = + FamousDefs(&ctx.sema, ctx.sema.scope(type_ref.syntax())?.krate()).core_option_Option()?; + if ret_enum != option_enum { + return None; + } + + let ok_type = unwrap_option_type(type_ref)?; + + acc.add( + AssistId("unwrap_option_return_type", AssistKind::RefactorRewrite), + "Unwrap Option return type", + type_ref.syntax().text_range(), + |builder| { + let body = ast::Expr::BlockExpr(body); + + let mut exprs_to_unwrap = Vec::new(); + let tail_cb = &mut |e: &_| tail_cb_impl(&mut exprs_to_unwrap, e); + walk_expr(&body, &mut |expr| { + if let Expr::ReturnExpr(ret_expr) = expr { + if let Some(ret_expr_arg) = &ret_expr.expr() { + for_each_tail_expr(ret_expr_arg, tail_cb); + } + } + }); + for_each_tail_expr(&body, tail_cb); + + let is_unit_type = is_unit_type(&ok_type); + if is_unit_type { + let mut text_range = ret_type.syntax().text_range(); + + if let Some(NodeOrToken::Token(token)) = ret_type.syntax().next_sibling_or_token() { + if token.kind() == SyntaxKind::WHITESPACE { + text_range = TextRange::new(text_range.start(), token.text_range().end()); + } + } + + builder.delete(text_range); + } else { + builder.replace(type_ref.syntax().text_range(), ok_type.syntax().text()); + } + + for ret_expr_arg in exprs_to_unwrap { + let ret_expr_str = ret_expr_arg.to_string(); + if ret_expr_str.starts_with("Some(") { + 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(", "), + ); + } + } + } else if ret_expr_str == "None" { + builder.replace(ret_expr_arg.syntax().text_range(), "()"); + } + } + }, + ) +} + +fn tail_cb_impl(acc: &mut Vec, e: &ast::Expr) { + match e { + Expr::BreakExpr(break_expr) => { + if let Some(break_expr_arg) = break_expr.expr() { + for_each_tail_expr(&break_expr_arg, &mut |e| tail_cb_impl(acc, e)) + } + } + Expr::ReturnExpr(_) => { + // all return expressions have already been handled by the walk loop + } + e => acc.push(e.clone()), + } +} + +// Tries to extract `T` from `Option`. +fn unwrap_option_type(ty: &ast::Type) -> Option { + let ast::Type::PathType(path_ty) = ty else { + return None; + }; + let path = path_ty.path()?; + let segment = path.first_segment()?; + let generic_arg_list = segment.generic_arg_list()?; + let generic_args: Vec<_> = generic_arg_list.generic_args().collect(); + let ast::GenericArg::TypeArg(some_type) = generic_args.first()? else { + return None; + }; + some_type.ty() +} + +fn is_unit_type(ty: &ast::Type) -> bool { + let ast::Type::TupleType(tuple) = ty else { return false }; + tuple.fields().next().is_none() +} + +#[cfg(test)] +mod tests { + use crate::tests::{check_assist, check_assist_not_applicable}; + + use super::*; + + #[test] + fn unwrap_option_return_type_simple() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + let test = "test"; + return Some(42i32); +} +"#, + r#" +fn foo() -> i32 { + let test = "test"; + return 42i32; +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_unit_type() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option<()$0> { + Some(()) +} +"#, + r#" +fn foo() { +} +"#, + ); + + // Unformatted return type + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option<()$0>{ + Some(()) +} +"#, + r#" +fn foo() { +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_none() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + if true { + Some(42) + } else { + None + } +} +"#, + r#" +fn foo() -> i32 { + if true { + 42 + } else { + () + } +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_ending_with_parent() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + if true { + Some(42) + } else { + foo() + } +} +"#, + r#" +fn foo() -> i32 { + if true { + 42 + } else { + foo() + } +} +"#, + ); + } + + #[test] + fn unwrap_return_type_break_split_tail() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + loop { + break if true { + Some(1) + } else { + Some(0) + }; + } +} +"#, + r#" +fn foo() -> i32 { + loop { + break if true { + 1 + } else { + 0 + }; + } +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_closure() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() { + || -> Option { + let test = "test"; + return Some(42i32); + }; +} +"#, + r#" +fn foo() { + || -> i32 { + let test = "test"; + return 42i32; + }; +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_return_type_bad_cursor() { + check_assist_not_applicable( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> i32 { + let test = "test";$0 + return 42i32; +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_return_type_bad_cursor_closure() { + check_assist_not_applicable( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() { + || -> i32 { + let test = "test";$0 + return 42i32; + }; +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_closure_non_block() { + check_assist_not_applicable( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() { || -> i$032 3; } +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_return_type_already_not_option_std() { + check_assist_not_applicable( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> i32$0 { + let test = "test"; + return 42i32; +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_return_type_already_not_option_closure() { + check_assist_not_applicable( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() { + || -> i32$0 { + let test = "test"; + return 42i32; + }; +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_tail() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() ->$0 Option { + let test = "test"; + Some(42i32) +} +"#, + r#" +fn foo() -> i32 { + let test = "test"; + 42i32 +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_tail_closure() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() { + || ->$0 Option { + let test = "test"; + Some(42i32) + }; +} +"#, + r#" +fn foo() { + || -> i32 { + let test = "test"; + 42i32 + }; +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_tail_only() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { Some(42i32) } +"#, + r#" +fn foo() -> i32 { 42i32 } +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_tail_block_like() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option$0 { + if true { + Some(42i32) + } else { + Some(24i32) + } +} +"#, + r#" +fn foo() -> i32 { + if true { + 42i32 + } else { + 24i32 + } +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_without_block_closure() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() { + || -> Option$0 { + if true { + Some(42i32) + } else { + Some(24i32) + } + }; +} +"#, + r#" +fn foo() { + || -> i32 { + if true { + 42i32 + } else { + 24i32 + } + }; +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_nested_if() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option$0 { + if true { + if false { + Some(1) + } else { + Some(2) + } + } else { + Some(24i32) + } +} +"#, + r#" +fn foo() -> i32 { + if true { + if false { + 1 + } else { + 2 + } + } else { + 24i32 + } +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_await() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +async fn foo() -> Option { + if true { + if false { + Some(1.await) + } else { + Some(2.await) + } + } else { + Some(24i32.await) + } +} +"#, + r#" +async fn foo() -> i32 { + if true { + if false { + 1.await + } else { + 2.await + } + } else { + 24i32.await + } +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_array() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option<[i32; 3]$0> { Some([1, 2, 3]) } +"#, + r#" +fn foo() -> [i32; 3] { [1, 2, 3] } +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_cast() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -$0> Option { + if true { + if false { + Some(1 as i32) + } else { + Some(2 as i32) + } + } else { + Some(24 as i32) + } +} +"#, + r#" +fn foo() -> i32 { + if true { + if false { + 1 as i32 + } else { + 2 as i32 + } + } else { + 24 as i32 + } +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_tail_block_like_match() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + let my_var = 5; + match my_var { + 5 => Some(42i32), + _ => Some(24i32), + } +} +"#, + r#" +fn foo() -> i32 { + let my_var = 5; + match my_var { + 5 => 42i32, + _ => 24i32, + } +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_loop_with_tail() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + let my_var = 5; + loop { + println!("test"); + 5 + } + Some(my_var) +} +"#, + r#" +fn foo() -> i32 { + let my_var = 5; + loop { + println!("test"); + 5 + } + my_var +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_loop_in_let_stmt() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + let my_var = let x = loop { + break 1; + }; + Some(my_var) +} +"#, + r#" +fn foo() -> i32 { + let my_var = let x = loop { + break 1; + }; + my_var +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_tail_block_like_match_return_expr() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option$0 { + let my_var = 5; + let res = match my_var { + 5 => 42i32, + _ => return Some(24i32), + }; + Some(res) +} +"#, + r#" +fn foo() -> i32 { + let my_var = 5; + let res = match my_var { + 5 => 42i32, + _ => return 24i32, + }; + res +} +"#, + ); + + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + let my_var = 5; + let res = if my_var == 5 { + 42i32 + } else { + return Some(24i32); + }; + Some(res) +} +"#, + r#" +fn foo() -> i32 { + let my_var = 5; + let res = if my_var == 5 { + 42i32 + } else { + return 24i32; + }; + res +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_tail_block_like_match_deeper() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + let my_var = 5; + match my_var { + 5 => { + if true { + Some(42i32) + } else { + Some(25i32) + } + }, + _ => { + let test = "test"; + if test == "test" { + return Some(bar()); + } + Some(53i32) + }, + } +} +"#, + r#" +fn foo() -> i32 { + let my_var = 5; + match my_var { + 5 => { + if true { + 42i32 + } else { + 25i32 + } + }, + _ => { + let test = "test"; + if test == "test" { + return bar(); + } + 53i32 + }, + } +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_tail_block_like_early_return() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + let test = "test"; + if test == "test" { + return Some(24i32); + } + Some(53i32) +} +"#, + r#" +fn foo() -> i32 { + let test = "test"; + if test == "test" { + return 24i32; + } + 53i32 +} +"#, + ); + } + + #[test] + fn wrap_return_in_tail_position() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo(num: i32) -> $0Option { + return Some(num) +} +"#, + r#" +fn foo(num: i32) -> i32 { + return num +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_closure() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo(the_field: u32) -> Option { + let true_closure = || { return true; }; + if the_field < 5 { + let mut i = 0; + if true_closure() { + return Some(99); + } else { + return Some(0); + } + } + Some(the_field) +} +"#, + r#" +fn foo(the_field: u32) -> u32 { + let true_closure = || { return true; }; + if the_field < 5 { + let mut i = 0; + if true_closure() { + return 99; + } else { + return 0; + } + } + the_field +} +"#, + ); + + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo(the_field: u32) -> Option { + let true_closure = || { + return true; + }; + if the_field < 5 { + let mut i = 0; + + + if true_closure() { + return Some(99); + } else { + return Some(0); + } + } + let t = None; + + Some(t.unwrap_or_else(|| the_field)) +} +"#, + r#" +fn foo(the_field: u32) -> u32 { + let true_closure = || { + return true; + }; + if the_field < 5 { + let mut i = 0; + + + if true_closure() { + return 99; + } else { + return 0; + } + } + let t = None; + + t.unwrap_or_else(|| the_field) +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_simple_with_weird_forms() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo() -> Option { + let test = "test"; + if test == "test" { + return Some(24i32); + } + let mut i = 0; + loop { + if i == 1 { + break Some(55); + } + i += 1; + } +} +"#, + r#" +fn foo() -> i32 { + let test = "test"; + if test == "test" { + return 24i32; + } + let mut i = 0; + loop { + if i == 1 { + break 55; + } + i += 1; + } +} +"#, + ); + + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo(the_field: u32) -> Option { + if the_field < 5 { + let mut i = 0; + loop { + if i > 5 { + return Some(55u32); + } + i += 3; + } + match i { + 5 => return Some(99), + _ => return Some(0), + }; + } + Some(the_field) +} +"#, + r#" +fn foo(the_field: u32) -> u32 { + if the_field < 5 { + let mut i = 0; + loop { + if i > 5 { + return 55u32; + } + i += 3; + } + match i { + 5 => return 99, + _ => return 0, + }; + } + the_field +} +"#, + ); + + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo(the_field: u32) -> Option { + if the_field < 5 { + let mut i = 0; + match i { + 5 => return Some(99), + _ => return Some(0), + } + } + Some(the_field) +} +"#, + r#" +fn foo(the_field: u32) -> u32 { + if the_field < 5 { + let mut i = 0; + match i { + 5 => return 99, + _ => return 0, + } + } + the_field +} +"#, + ); + + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo(the_field: u32) -> Option { + if the_field < 5 { + let mut i = 0; + if i == 5 { + return Some(99) + } else { + return Some(0) + } + } + Some(the_field) +} +"#, + r#" +fn foo(the_field: u32) -> u32 { + if the_field < 5 { + let mut i = 0; + if i == 5 { + return 99 + } else { + return 0 + } + } + the_field +} +"#, + ); + + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option +fn foo(the_field: u32) -> Option { + if the_field < 5 { + let mut i = 0; + if i == 5 { + return Some(99); + } else { + return Some(0); + } + } + Some(the_field) +} +"#, + r#" +fn foo(the_field: u32) -> u32 { + if the_field < 5 { + let mut i = 0; + if i == 5 { + return 99; + } else { + return 0; + } + } + the_field +} +"#, + ); + } + + #[test] + fn unwrap_option_return_type_nested_type() { + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option, result +fn foo() -> Option> { + Some(Ok(42)) +} +"#, + r#" +fn foo() -> Result { + Ok(42) +} +"#, + ); + + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option, result +fn foo() -> Option, ()>> { + Some(Err()) +} +"#, + r#" +fn foo() -> Result, ()> { + Err() +} +"#, + ); + + check_assist( + unwrap_option_return_type, + r#" +//- minicore: option, result, iterators +fn foo() -> Option$0> { + Some(Some(42).into_iter()) +} +"#, + r#" +fn foo() -> impl Iterator { + Some(42).into_iter() +} +"#, + ); + } +} diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/handlers/wrap_return_type_in_option.rs b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/wrap_return_type_in_option.rs new file mode 100644 index 0000000000000..90cfc2ef17c74 --- /dev/null +++ b/src/tools/rust-analyzer/crates/ide-assists/src/handlers/wrap_return_type_in_option.rs @@ -0,0 +1,1173 @@ +use std::iter; + +use hir::HasSource; +use ide_db::{ + famous_defs::FamousDefs, + syntax_helpers::node_ext::{for_each_tail_expr, walk_expr}, +}; +use itertools::Itertools; +use syntax::{ + ast::{self, make, Expr, HasGenericParams}, + match_ast, ted, AstNode, ToSmolStr, +}; + +use crate::{AssistContext, AssistId, AssistKind, Assists}; + +// Assist: wrap_return_type_in_option +// +// Wrap the function's return type into Option. +// +// ``` +// # //- minicore: option +// fn foo() -> i32$0 { 42i32 } +// ``` +// -> +// ``` +// fn foo() -> Option { Some(42i32) } +// ``` +pub(crate) fn wrap_return_type_in_option(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { + let ret_type = ctx.find_node_at_offset::()?; + let parent = ret_type.syntax().parent()?; + let body = match_ast! { + match parent { + ast::Fn(func) => func.body()?, + ast::ClosureExpr(closure) => match closure.body()? { + Expr::BlockExpr(block) => block, + // closures require a block when a return type is specified + _ => return None, + }, + _ => return None, + } + }; + + let type_ref = &ret_type.ty()?; + let core_option = + FamousDefs(&ctx.sema, ctx.sema.scope(type_ref.syntax())?.krate()).core_option_Option()?; + + let ty = ctx.sema.resolve_type(type_ref)?.as_adt(); + if matches!(ty, Some(hir::Adt::Enum(ret_type)) if ret_type == core_option) { + // The return type is already wrapped in an Option + cov_mark::hit!(wrap_return_type_in_option_simple_return_type_already_option); + return None; + } + + acc.add( + AssistId("wrap_return_type_in_option", AssistKind::RefactorRewrite), + "Wrap return type in Option", + type_ref.syntax().text_range(), + |edit| { + let new_option_ty = option_type(ctx, &core_option, type_ref).clone_for_update(); + let body = edit.make_mut(ast::Expr::BlockExpr(body)); + + let mut exprs_to_wrap = Vec::new(); + let tail_cb = &mut |e: &_| tail_cb_impl(&mut exprs_to_wrap, e); + walk_expr(&body, &mut |expr| { + if let Expr::ReturnExpr(ret_expr) = expr { + if let Some(ret_expr_arg) = &ret_expr.expr() { + for_each_tail_expr(ret_expr_arg, tail_cb); + } + } + }); + for_each_tail_expr(&body, tail_cb); + + for ret_expr_arg in exprs_to_wrap { + let some_wrapped = make::expr_call( + make::expr_path(make::ext::ident_path("Some")), + make::arg_list(iter::once(ret_expr_arg.clone())), + ) + .clone_for_update(); + ted::replace(ret_expr_arg.syntax(), some_wrapped.syntax()); + } + + let old_option_ty = edit.make_mut(type_ref.clone()); + ted::replace(old_option_ty.syntax(), new_option_ty.syntax()); + }, + ) +} + +fn option_type( + ctx: &AssistContext<'_>, + core_option: &hir::Enum, + ret_type: &ast::Type, +) -> ast::Type { + // Try to find an Option type alias in the current scope (shadowing the default). + let option_path = hir::ModPath::from_segments( + hir::PathKind::Plain, + iter::once(hir::Name::new_symbol_root(hir::sym::Option.clone())), + ); + let alias = ctx.sema.resolve_mod_path(ret_type.syntax(), &option_path).and_then(|def| { + def.filter_map(|def| match def.as_module_def()? { + hir::ModuleDef::TypeAlias(alias) => { + let enum_ty = alias.ty(ctx.db()).as_adt()?.as_enum()?; + (&enum_ty == core_option).then_some(alias) + } + _ => None, + }) + .find_map(|alias| { + let mut inserted_ret_type = false; + let generic_params = alias + .source(ctx.db())? + .value + .generic_param_list()? + .generic_params() + .map(|param| match param { + // Replace the very first type parameter with the functions return type. + ast::GenericParam::TypeParam(_) if !inserted_ret_type => { + inserted_ret_type = true; + ret_type.to_smolstr() + } + ast::GenericParam::LifetimeParam(_) => make::lifetime("'_").to_smolstr(), + _ => make::ty_placeholder().to_smolstr(), + }) + .join(", "); + + let name = alias.name(ctx.db()); + let name = name.as_str(); + Some(make::ty(&format!("{name}<{generic_params}>"))) + }) + }); + // If there is no applicable alias in scope use the default Option type. + alias.unwrap_or_else(|| make::ext::ty_option(ret_type.clone())) +} + +fn tail_cb_impl(acc: &mut Vec, e: &ast::Expr) { + match e { + Expr::BreakExpr(break_expr) => { + if let Some(break_expr_arg) = break_expr.expr() { + for_each_tail_expr(&break_expr_arg, &mut |e| tail_cb_impl(acc, e)) + } + } + Expr::ReturnExpr(_) => { + // all return expressions have already been handled by the walk loop + } + e => acc.push(e.clone()), + } +} + +#[cfg(test)] +mod tests { + use crate::tests::{check_assist, check_assist_not_applicable}; + + use super::*; + + #[test] + fn wrap_return_type_in_option_simple() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i3$02 { + let test = "test"; + return 42i32; +} +"#, + r#" +fn foo() -> Option { + let test = "test"; + return Some(42i32); +} +"#, + ); + } + + #[test] + fn wrap_return_type_break_split_tail() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i3$02 { + loop { + break if true { + 1 + } else { + 0 + }; + } +} +"#, + r#" +fn foo() -> Option { + loop { + break if true { + Some(1) + } else { + Some(0) + }; + } +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_closure() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() { + || -> i32$0 { + let test = "test"; + return 42i32; + }; +} +"#, + r#" +fn foo() { + || -> Option { + let test = "test"; + return Some(42i32); + }; +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_return_type_bad_cursor() { + check_assist_not_applicable( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32 { + let test = "test";$0 + return 42i32; +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_return_type_bad_cursor_closure() { + check_assist_not_applicable( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() { + || -> i32 { + let test = "test";$0 + return 42i32; + }; +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_closure_non_block() { + check_assist_not_applicable( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() { || -> i$032 3; } +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_return_type_already_option_std() { + check_assist_not_applicable( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> core::option::Option { + let test = "test"; + return 42i32; +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_return_type_already_option() { + cov_mark::check!(wrap_return_type_in_option_simple_return_type_already_option); + check_assist_not_applicable( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> Option { + let test = "test"; + return 42i32; +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_return_type_already_option_closure() { + check_assist_not_applicable( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() { + || -> Option { + let test = "test"; + return 42i32; + }; +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_cursor() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> $0i32 { + let test = "test"; + return 42i32; +} +"#, + r#" +fn foo() -> Option { + let test = "test"; + return Some(42i32); +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_tail() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() ->$0 i32 { + let test = "test"; + 42i32 +} +"#, + r#" +fn foo() -> Option { + let test = "test"; + Some(42i32) +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_tail_closure() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() { + || ->$0 i32 { + let test = "test"; + 42i32 + }; +} +"#, + r#" +fn foo() { + || -> Option { + let test = "test"; + Some(42i32) + }; +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_tail_only() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { 42i32 } +"#, + r#" +fn foo() -> Option { Some(42i32) } +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_tail_block_like() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + if true { + 42i32 + } else { + 24i32 + } +} +"#, + r#" +fn foo() -> Option { + if true { + Some(42i32) + } else { + Some(24i32) + } +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_without_block_closure() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() { + || -> i32$0 { + if true { + 42i32 + } else { + 24i32 + } + }; +} +"#, + r#" +fn foo() { + || -> Option { + if true { + Some(42i32) + } else { + Some(24i32) + } + }; +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_nested_if() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + if true { + if false { + 1 + } else { + 2 + } + } else { + 24i32 + } +} +"#, + r#" +fn foo() -> Option { + if true { + if false { + Some(1) + } else { + Some(2) + } + } else { + Some(24i32) + } +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_await() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +async fn foo() -> i$032 { + if true { + if false { + 1.await + } else { + 2.await + } + } else { + 24i32.await + } +} +"#, + r#" +async fn foo() -> Option { + if true { + if false { + Some(1.await) + } else { + Some(2.await) + } + } else { + Some(24i32.await) + } +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_array() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> [i32;$0 3] { [1, 2, 3] } +"#, + r#" +fn foo() -> Option<[i32; 3]> { Some([1, 2, 3]) } +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_cast() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -$0> i32 { + if true { + if false { + 1 as i32 + } else { + 2 as i32 + } + } else { + 24 as i32 + } +} +"#, + r#" +fn foo() -> Option { + if true { + if false { + Some(1 as i32) + } else { + Some(2 as i32) + } + } else { + Some(24 as i32) + } +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_tail_block_like_match() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + let my_var = 5; + match my_var { + 5 => 42i32, + _ => 24i32, + } +} +"#, + r#" +fn foo() -> Option { + let my_var = 5; + match my_var { + 5 => Some(42i32), + _ => Some(24i32), + } +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_loop_with_tail() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + let my_var = 5; + loop { + println!("test"); + 5 + } + my_var +} +"#, + r#" +fn foo() -> Option { + let my_var = 5; + loop { + println!("test"); + 5 + } + Some(my_var) +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_loop_in_let_stmt() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + let my_var = let x = loop { + break 1; + }; + my_var +} +"#, + r#" +fn foo() -> Option { + let my_var = let x = loop { + break 1; + }; + Some(my_var) +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_tail_block_like_match_return_expr() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + let my_var = 5; + let res = match my_var { + 5 => 42i32, + _ => return 24i32, + }; + res +} +"#, + r#" +fn foo() -> Option { + let my_var = 5; + let res = match my_var { + 5 => 42i32, + _ => return Some(24i32), + }; + Some(res) +} +"#, + ); + + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + let my_var = 5; + let res = if my_var == 5 { + 42i32 + } else { + return 24i32; + }; + res +} +"#, + r#" +fn foo() -> Option { + let my_var = 5; + let res = if my_var == 5 { + 42i32 + } else { + return Some(24i32); + }; + Some(res) +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_tail_block_like_match_deeper() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + let my_var = 5; + match my_var { + 5 => { + if true { + 42i32 + } else { + 25i32 + } + }, + _ => { + let test = "test"; + if test == "test" { + return bar(); + } + 53i32 + }, + } +} +"#, + r#" +fn foo() -> Option { + let my_var = 5; + match my_var { + 5 => { + if true { + Some(42i32) + } else { + Some(25i32) + } + }, + _ => { + let test = "test"; + if test == "test" { + return Some(bar()); + } + Some(53i32) + }, + } +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_tail_block_like_early_return() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i$032 { + let test = "test"; + if test == "test" { + return 24i32; + } + 53i32 +} +"#, + r#" +fn foo() -> Option { + let test = "test"; + if test == "test" { + return Some(24i32); + } + Some(53i32) +} +"#, + ); + } + + #[test] + fn wrap_return_in_tail_position() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo(num: i32) -> $0i32 { + return num +} +"#, + r#" +fn foo(num: i32) -> Option { + return Some(num) +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_closure() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo(the_field: u32) ->$0 u32 { + let true_closure = || { return true; }; + if the_field < 5 { + let mut i = 0; + if true_closure() { + return 99; + } else { + return 0; + } + } + the_field +} +"#, + r#" +fn foo(the_field: u32) -> Option { + let true_closure = || { return true; }; + if the_field < 5 { + let mut i = 0; + if true_closure() { + return Some(99); + } else { + return Some(0); + } + } + Some(the_field) +} +"#, + ); + + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo(the_field: u32) -> u32$0 { + let true_closure = || { + return true; + }; + if the_field < 5 { + let mut i = 0; + + + if true_closure() { + return 99; + } else { + return 0; + } + } + let t = None; + + t.unwrap_or_else(|| the_field) +} +"#, + r#" +fn foo(the_field: u32) -> Option { + let true_closure = || { + return true; + }; + if the_field < 5 { + let mut i = 0; + + + if true_closure() { + return Some(99); + } else { + return Some(0); + } + } + let t = None; + + Some(t.unwrap_or_else(|| the_field)) +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_option_simple_with_weird_forms() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i32$0 { + let test = "test"; + if test == "test" { + return 24i32; + } + let mut i = 0; + loop { + if i == 1 { + break 55; + } + i += 1; + } +} +"#, + r#" +fn foo() -> Option { + let test = "test"; + if test == "test" { + return Some(24i32); + } + let mut i = 0; + loop { + if i == 1 { + break Some(55); + } + i += 1; + } +} +"#, + ); + + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo(the_field: u32) -> u32$0 { + if the_field < 5 { + let mut i = 0; + loop { + if i > 5 { + return 55u32; + } + i += 3; + } + match i { + 5 => return 99, + _ => return 0, + }; + } + the_field +} +"#, + r#" +fn foo(the_field: u32) -> Option { + if the_field < 5 { + let mut i = 0; + loop { + if i > 5 { + return Some(55u32); + } + i += 3; + } + match i { + 5 => return Some(99), + _ => return Some(0), + }; + } + Some(the_field) +} +"#, + ); + + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo(the_field: u32) -> u3$02 { + if the_field < 5 { + let mut i = 0; + match i { + 5 => return 99, + _ => return 0, + } + } + the_field +} +"#, + r#" +fn foo(the_field: u32) -> Option { + if the_field < 5 { + let mut i = 0; + match i { + 5 => return Some(99), + _ => return Some(0), + } + } + Some(the_field) +} +"#, + ); + + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo(the_field: u32) -> u32$0 { + if the_field < 5 { + let mut i = 0; + if i == 5 { + return 99 + } else { + return 0 + } + } + the_field +} +"#, + r#" +fn foo(the_field: u32) -> Option { + if the_field < 5 { + let mut i = 0; + if i == 5 { + return Some(99) + } else { + return Some(0) + } + } + Some(the_field) +} +"#, + ); + + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo(the_field: u32) -> $0u32 { + if the_field < 5 { + let mut i = 0; + if i == 5 { + return 99; + } else { + return 0; + } + } + the_field +} +"#, + r#" +fn foo(the_field: u32) -> Option { + if the_field < 5 { + let mut i = 0; + if i == 5 { + return Some(99); + } else { + return Some(0); + } + } + Some(the_field) +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_local_option_type() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +type Option = core::option::Option; + +fn foo() -> i3$02 { + return 42i32; +} +"#, + r#" +type Option = core::option::Option; + +fn foo() -> Option { + return Some(42i32); +} +"#, + ); + + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +type Option2 = core::option::Option; + +fn foo() -> i3$02 { + return 42i32; +} +"#, + r#" +type Option2 = core::option::Option; + +fn foo() -> Option { + return Some(42i32); +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_imported_local_option_type() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +mod some_module { + pub type Option = core::option::Option; +} + +use some_module::Option; + +fn foo() -> i3$02 { + return 42i32; +} +"#, + r#" +mod some_module { + pub type Option = core::option::Option; +} + +use some_module::Option; + +fn foo() -> Option { + return Some(42i32); +} +"#, + ); + + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +mod some_module { + pub type Option = core::option::Option; +} + +use some_module::*; + +fn foo() -> i3$02 { + return 42i32; +} +"#, + r#" +mod some_module { + pub type Option = core::option::Option; +} + +use some_module::*; + +fn foo() -> Option { + return Some(42i32); +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_local_option_type_from_function_body() { + check_assist( + wrap_return_type_in_option, + r#" +//- minicore: option +fn foo() -> i3$02 { + type Option = core::option::Option; + 0 +} +"#, + r#" +fn foo() -> Option { + type Option = core::option::Option; + Some(0) +} +"#, + ); + } + + #[test] + fn wrap_return_type_in_local_option_type_already_using_alias() { + check_assist_not_applicable( + wrap_return_type_in_option, + r#" +//- minicore: option +pub type Option = core::option::Option; + +fn foo() -> Option { + return Some(42i32); +} +"#, + ); + } +} diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs index c98655b423029..63a3fec3f41d5 100644 --- a/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs +++ b/src/tools/rust-analyzer/crates/ide-assists/src/lib.rs @@ -223,8 +223,10 @@ mod handlers { mod unnecessary_async; mod unqualify_method_call; mod unwrap_block; + mod unwrap_option_return_type; mod unwrap_result_return_type; mod unwrap_tuple; + mod wrap_return_type_in_option; mod wrap_return_type_in_result; mod wrap_unwrap_cfg_attr; @@ -355,9 +357,11 @@ mod handlers { unmerge_use::unmerge_use, unnecessary_async::unnecessary_async, unwrap_block::unwrap_block, + unwrap_option_return_type::unwrap_option_return_type, unwrap_result_return_type::unwrap_result_return_type, unwrap_tuple::unwrap_tuple, unqualify_method_call::unqualify_method_call, + wrap_return_type_in_option::wrap_return_type_in_option, wrap_return_type_in_result::wrap_return_type_in_result, wrap_unwrap_cfg_attr::wrap_unwrap_cfg_attr, diff --git a/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs b/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs index 48e12a8107353..933d45d750832 100644 --- a/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs +++ b/src/tools/rust-analyzer/crates/ide-assists/src/tests/generated.rs @@ -3264,6 +3264,20 @@ fn foo() { ) } +#[test] +fn doctest_unwrap_option_return_type() { + check_doc_test( + "unwrap_option_return_type", + r#####" +//- minicore: option +fn foo() -> Option$0 { Some(42i32) } +"#####, + r#####" +fn foo() -> i32 { 42i32 } +"#####, + ) +} + #[test] fn doctest_unwrap_result_return_type() { check_doc_test( @@ -3297,6 +3311,20 @@ fn main() { ) } +#[test] +fn doctest_wrap_return_type_in_option() { + check_doc_test( + "wrap_return_type_in_option", + r#####" +//- minicore: option +fn foo() -> i32$0 { 42i32 } +"#####, + r#####" +fn foo() -> Option { Some(42i32) } +"#####, + ) +} + #[test] fn doctest_wrap_return_type_in_result() { check_doc_test(