Unnamed repository; edit this file 'description' to name the repository.
| -rw-r--r-- | crates/ide-assists/src/handlers/replace_or_with_or_else.rs | 117 |
1 files changed, 86 insertions, 31 deletions
diff --git a/crates/ide-assists/src/handlers/replace_or_with_or_else.rs b/crates/ide-assists/src/handlers/replace_or_with_or_else.rs index bc1122a9d2..bee52a0d7f 100644 --- a/crates/ide-assists/src/handlers/replace_or_with_or_else.rs +++ b/crates/ide-assists/src/handlers/replace_or_with_or_else.rs @@ -30,33 +30,33 @@ use crate::{AssistContext, Assists}; pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { let call: ast::MethodCallExpr = ctx.find_node_at_offset()?; - is_option_or_result(call.receiver()?, ctx)?; + let kind = is_option_or_result(call.receiver()?, ctx)?; let (name, arg_list) = (call.name_ref()?, call.arg_list()?); + let mut map_or = false; + let replace = match &*name.text() { "unwrap_or" => "unwrap_or_else".to_string(), - "ok_or" => "ok_or_else".to_string(), + "or" => "or_else".to_string(), + "ok_or" if kind == Kind::Option => "ok_or_else".to_string(), + "map_or" => { + map_or = true; + "map_or_else".to_string() + } _ => return None, }; let arg = match arg_list.args().collect::<Vec<_>>().as_slice() { [] => make::arg_list(Vec::new()), [first] => { - let param = (|| { - if let ast::Expr::CallExpr(call) = first { - if call.arg_list()?.args().count() == 0 { - Some(call.expr()?.clone()) - } else { - None - } - } else { - None - } - })() - .unwrap_or_else(|| make::expr_closure(None, first.clone())); + let param = into_closure(first); make::arg_list(vec![param]) } + [first, second] if map_or => { + let param = into_closure(first); + make::arg_list(vec![param, second.clone()]) + } _ => return None, }; @@ -71,6 +71,21 @@ pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_> ) } +fn into_closure(param: &Expr) -> Expr { + (|| { + if let ast::Expr::CallExpr(call) = param { + if call.arg_list()?.args().count() == 0 { + Some(call.expr()?.clone()) + } else { + None + } + } else { + None + } + })() + .unwrap_or_else(|| make::expr_closure(None, param.clone())) +} + // Assist: replace_or_else_with_or // // Replace `unwrap_or_else` with `unwrap_or` and `ok_or_else` with `ok_or`. @@ -92,33 +107,32 @@ pub(crate) fn replace_or_with_or_else(acc: &mut Assists, ctx: &AssistContext<'_> pub(crate) fn replace_or_else_with_or(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> { let call: ast::MethodCallExpr = ctx.find_node_at_offset()?; - is_option_or_result(call.receiver()?, ctx)?; + let kind = is_option_or_result(call.receiver()?, ctx)?; let (name, arg_list) = (call.name_ref()?, call.arg_list()?); + let mut map_or = false; let replace = match &*name.text() { "unwrap_or_else" => "unwrap_or".to_string(), - "ok_or_else" => "ok_or".to_string(), + "or_else" => "or".to_string(), + "ok_or_else" if kind == Kind::Option => "ok_or".to_string(), + "map_or_else" => { + map_or = true; + "map_or".to_string() + } _ => return None, }; let arg = match arg_list.args().collect::<Vec<_>>().as_slice() { [] => make::arg_list(Vec::new()), [first] => { - let param = (|| { - if let ast::Expr::ClosureExpr(closure) = first { - if closure.param_list()?.params().count() == 0 { - Some(closure.body()?.clone()) - } else { - None - } - } else { - None - } - })() - .unwrap_or_else(|| make::expr_call(first.clone(), make::arg_list(Vec::new()))); + let param = into_call(first); make::arg_list(vec![param]) } + [first, second] if map_or => { + let param = into_call(first); + make::arg_list(vec![param, second.clone()]) + } _ => return None, }; @@ -133,14 +147,35 @@ pub(crate) fn replace_or_else_with_or(acc: &mut Assists, ctx: &AssistContext<'_> ) } -fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<()> { +fn into_call(param: &Expr) -> Expr { + (|| { + if let ast::Expr::ClosureExpr(closure) = param { + if closure.param_list()?.params().count() == 0 { + Some(closure.body()?.clone()) + } else { + None + } + } else { + None + } + })() + .unwrap_or_else(|| make::expr_call(param.clone(), make::arg_list(Vec::new()))) +} + +#[derive(PartialEq, Eq)] +enum Kind { + Option, + Result, +} + +fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<Kind> { let ty = ctx.sema.type_of_expr(&receiver)?.adjusted().as_adt()?.as_enum()?; let option_enum = FamousDefs(&ctx.sema, ctx.sema.scope(receiver.syntax())?.krate()).core_option_Option(); if let Some(option_enum) = option_enum { if ty == option_enum { - return Some(()); + return Some(Kind::Option); } } @@ -149,7 +184,7 @@ fn is_option_or_result(receiver: Expr, ctx: &AssistContext<'_>) -> Option<()> { if let Some(result_enum) = result_enum { if ty == result_enum { - return Some(()); + return Some(Kind::Result); } } @@ -295,6 +330,26 @@ fn foo() { } #[test] + fn replace_or_else_with_or_map() { + check_assist( + replace_or_else_with_or, + r#" +//- minicore: result +fn foo() { + let foo = Ok("foo"); + return foo.map$0_or_else(|| 42, |v| v.len()); +} +"#, + r#" +fn foo() { + let foo = Ok("foo"); + return foo.map_or(42, |v| v.len()); +} +"#, + ) + } + + #[test] fn replace_or_else_with_or_not_applicable() { check_assist_not_applicable( replace_or_else_with_or, |